From 2899605b6825f4750588fc7049826983669741ef Mon Sep 17 00:00:00 2001 From: GuyAv46 Date: Thu, 2 Jul 2026 13:48:08 +0300 Subject: [PATCH 1/3] Multi-accumulator x86 SIMD kernels + lower small-dim SIMD thresholds Break the single-accumulator dependency chain in the x86 distance kernels by accumulating into 2-4 independent SIMD registers (following the pattern already used by the SQ8_FP16 kernels and the ARM implementations), and convert the main loops from do-while to while so the kernels are correct for dimensions smaller than one full chunk. Kernels updated (IP + L2, plus SQ8 Cosine/L2 wrappers that inherit): - FP32/FP64: SSE, AVX, AVX512F - FP16: F16C, AVX512F - SQ8_FP32: SSE4, AVX2, AVX2_FMA, AVX512F_BW_VL_VNNI Dispatcher thresholds lowered where the kernels' residual handling is safe below one full chunk (residual loads never touch memory past the minimum dimension): - FP32: dim >= 8 (was 16), FP64: dim >= 4 (was 8) - FP16: per-tier - F16C >= 8, AVX512F >= 16, AVX512FP16_VL stays 32 - SQ8_FP32: dim >= 8 (was 16) Measured on Intel Core Ultra 7 255U (AVX2/F16C tiers, google-benchmark, 9 repetitions, interleaved, median): - FP16 F16C L2/IP: up to 1.9x at dim 1024 - SQ8_FP32 AVX2_FMA IP: up to 1.8x - FP32 SSE L2: up to 1.45x, FP64 SSE up to 1.7x - FP32/FP64 AVX: 1.1-1.3x - Newly-SIMD small dims: FP32 dim 8-15 1.4-2.8x vs scalar, FP16 dim 8-31 5.7-9.4x vs scalar AVX-512 tiers follow the same structure but were not benchmarked locally (no AVX-512 hardware); loop/tail coverage verified by simulation for all dims up to 4096. Co-Authored-By: Claude Fable 5 --- src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h | 39 ++++++++----- src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h | 39 ++++++++----- .../IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h | 36 +++++++++--- src/VecSim/spaces/IP/IP_AVX512F_FP16.h | 32 +++++++---- src/VecSim/spaces/IP/IP_AVX512F_FP32.h | 31 +++++++--- src/VecSim/spaces/IP/IP_AVX512F_FP64.h | 31 +++++++--- src/VecSim/spaces/IP/IP_AVX_FP32.h | 35 ++++++++---- src/VecSim/spaces/IP/IP_AVX_FP64.h | 33 +++++++---- src/VecSim/spaces/IP/IP_F16C_FP16.h | 40 +++++++------ src/VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h | 35 +++++++----- src/VecSim/spaces/IP/IP_SSE_FP32.h | 35 +++++++----- src/VecSim/spaces/IP/IP_SSE_FP64.h | 31 ++++++---- src/VecSim/spaces/IP_space.cpp | 33 ++++++----- src/VecSim/spaces/L2/L2_AVX512F_FP16.h | 32 +++++++---- src/VecSim/spaces/L2/L2_AVX512F_FP32.h | 31 +++++++--- src/VecSim/spaces/L2/L2_AVX512F_FP64.h | 31 +++++++--- src/VecSim/spaces/L2/L2_AVX_FP32.h | 32 ++++++++--- src/VecSim/spaces/L2/L2_AVX_FP64.h | 32 ++++++++--- src/VecSim/spaces/L2/L2_F16C_FP16.h | 40 +++++++------ src/VecSim/spaces/L2/L2_SSE_FP32.h | 36 +++++++----- src/VecSim/spaces/L2/L2_SSE_FP64.h | 31 ++++++---- src/VecSim/spaces/L2_space.cpp | 30 +++++----- tests/unit/test_spaces.cpp | 57 +++++++------------ 23 files changed, 516 insertions(+), 286 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h index 5767a4828..a948e6166 100644 --- a/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h @@ -52,10 +52,15 @@ float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, siz const float *pVect2 = static_cast(pVect2v); // FP32 query const uint8_t *pEnd1 = pVect1 + dimension; - // Initialize sum accumulator for Σ(q_i * y_i) - __m256 sum256 = _mm256_setzero_ps(); - - // Handle residual elements first (0-7 elements) + // Initialize sum accumulators for Σ(q_i * y_i). 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; @@ -71,23 +76,31 @@ float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, siz pVect2 += residual % 8; // Compute q_i * y_i (no dequantization) - sum256 = _mm256_mul_ps(v1_f, v2); + sum0 = _mm256_mul_ps(v1_f, v2); } // If the residual is >=8, have another step of 8 floats if constexpr (residual >= 8) { - InnerProductStepSQ8_FMA(pVect1, pVect2, sum256); + InnerProductStepSQ8_FMA(pVect1, pVect2, sum1); } - // Process remaining full chunks of 16 elements (2x8) - // Using do-while since dim > 16 guarantees at least one iteration - do { - InnerProductStepSQ8_FMA(pVect1, pVect2, sum256); - InnerProductStepSQ8_FMA(pVect1, pVect2, sum256); - } while (pVect1 < pEnd1); + // We are left with some multiple of 16 elements. The main loop handles 32 per iteration; + // a possible leftover block of 16 is handled after it. The loops may run zero times + // (dim can be as small as 8). + while (pVect1 + 32 <= pEnd1) { + InnerProductStepSQ8_FMA(pVect1, pVect2, sum0); + InnerProductStepSQ8_FMA(pVect1, pVect2, sum1); + InnerProductStepSQ8_FMA(pVect1, pVect2, sum2); + InnerProductStepSQ8_FMA(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + InnerProductStepSQ8_FMA(pVect1, pVect2, sum2); + InnerProductStepSQ8_FMA(pVect1, pVect2, sum3); + } // Reduce to get Σ(q_i * y_i) - float quantized_dot = my_mm256_reduce_add_ps(sum256); + float quantized_dot = + my_mm256_reduce_add_ps(_mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3))); // Get quantization parameters from stored vector (after quantized data) const uint8_t *pVect1Base = static_cast(pVect1v); diff --git a/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h index dea167eb3..3342e9e20 100644 --- a/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h @@ -52,10 +52,15 @@ float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, si const float *pVect2 = static_cast(pVect2v); // FP32 query const uint8_t *pEnd1 = pVect1 + dimension; - // Initialize sum accumulator for Σ(q_i * y_i) - __m256 sum256 = _mm256_setzero_ps(); - - // Handle residual elements first (0-7 elements) + // Initialize sum accumulators for Σ(q_i * y_i). Four accumulators break the mul->add + // dependency chain, letting more loads/adds 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; @@ -71,23 +76,31 @@ float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, si pVect2 += residual % 8; // Compute q_i * y_i (no dequantization) - sum256 = _mm256_mul_ps(v1_f, v2); + sum0 = _mm256_mul_ps(v1_f, v2); } // If the residual is >=8, have another step of 8 floats if constexpr (residual >= 8) { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum256); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); } - // Process remaining full chunks of 16 elements (2x8) - // Using do-while since dim > 16 guarantees at least one iteration - do { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum256); - InnerProductStepSQ8_FP32(pVect1, pVect2, sum256); - } while (pVect1 < pEnd1); + // We are left with some multiple of 16 elements. The main loop handles 32 per iteration; + // a possible leftover block of 16 is handled after it. The loops may run zero times + // (dim can be as small as 8). + while (pVect1 + 32 <= pEnd1) { + InnerProductStepSQ8_FP32(pVect1, pVect2, sum0); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); + } // Reduce to get Σ(q_i * y_i) - float quantized_dot = my_mm256_reduce_add_ps(sum256); + float quantized_dot = + my_mm256_reduce_add_ps(_mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3))); // Get quantization parameters from stored vector (after quantized data) const uint8_t *pVect1Base = static_cast(pVect1v); diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h index 76a590519..75fe35b44 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h @@ -51,14 +51,20 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si const float *pVec2 = static_cast(pVec2v); // FP32 query const uint8_t *pEnd1 = pVec1 + dimension; - // Initialize sum accumulator for Σ(q_i * y_i) - __m512 sum = _mm512_setzero_ps(); + // Initialize sum accumulators for Σ(q_i * y_i). Four accumulators break the FMA dependency + // chain, letting more FMAs be in flight at once. + __m512 sum0 = _mm512_setzero_ps(); + __m512 sum1 = _mm512_setzero_ps(); + __m512 sum2 = _mm512_setzero_ps(); + __m512 sum3 = _mm512_setzero_ps(); // Handle residual elements first (0 to 15) if constexpr (residual > 0) { __mmask16 mask = (1U << residual) - 1; - // Load uint8 elements (safe to load 16 bytes due to padding) + // 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); @@ -67,19 +73,31 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si __m512 v2 = _mm512_maskz_loadu_ps(mask, pVec2); // Compute q_i * y_i with mask (no dequantization) - sum = _mm512_maskz_mul_ps(mask, v1_f, v2); + sum0 = _mm512_maskz_mul_ps(mask, v1_f, v2); pVec1 += residual; pVec2 += residual; } - // Process full chunks of 16 elements - // Using do-while since dim > 16 guarantees at least one iteration - do { - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum); - } while (pVec1 < pEnd1); + // Process full chunks of 16 elements. The main loop handles 64 per iteration; leftover + // blocks of 16/32/48 are handled after it. The loops may run zero times (dim can be as + // small as 8). + while (pVec1 + 64 <= pEnd1) { + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum0); + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); + } + if (pVec1 + 32 <= pEnd1) { + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); + } + if (pVec1 < pEnd1) { + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); + } // Reduce to get Σ(q_i * y_i) + __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); float quantized_dot = _mm512_reduce_add_ps(sum); // Get quantization parameters from stored vector (after quantized data) diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h index b09352533..5f298c303 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h @@ -31,11 +31,15 @@ float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, s const float16 *pEnd1 = pVect1 + dimension; - auto sum = _mm512_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + auto sum0 = _mm512_setzero_ps(); + auto sum1 = _mm512_setzero_ps(); + auto sum2 = _mm512_setzero_ps(); + auto sum3 = _mm512_setzero_ps(); if constexpr (residual % 16) { - // Deal with remainder first. `dim` is more than 32, so we have at least one block of 32 - // 16-bit float so mask loading is guaranteed to be safe. + // Deal with remainder first. The full-width load of 16 16-bit floats is safe because + // `dim` is at least 16, so the vector spans at least 16 elements. __mmask16 constexpr residuals_mask = (1 << (residual % 16)) - 1; // Convert the first half-floats in the residual positions into floats and store them // 512 bits register, where the floats in the positions corresponding to the non-residuals @@ -44,20 +48,28 @@ float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, s _mm512_cvtph_ps(_mm256_lddqu_si256((__m256i *)pVect1))); auto v2 = _mm512_maskz_mov_ps(residuals_mask, _mm512_cvtph_ps(_mm256_lddqu_si256((__m256i *)pVect2))); - sum = _mm512_mul_ps(v1, v2); + sum0 = _mm512_mul_ps(v1, v2); pVect1 += residual % 16; pVect2 += residual % 16; } if constexpr (residual >= 16) { - InnerProductStep(pVect1, pVect2, sum); + InnerProductStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // In every iteration we process 2 chunks of 256bit (32 FP16) - do { - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 64 elements per iteration; a possible leftover block of 32 is + // handled after it. The loops may run zero times (dim can be as small as 16). + while (pVect1 + 64 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + auto sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); return 1.0f - _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h index 88421fd39..20995ee79 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h @@ -23,23 +23,40 @@ float FP32_InnerProductSIMD16_AVX512(const void *pVect1v, const void *pVect2v, s const float *pEnd1 = pVect1 + dimension; - __m512 sum512 = _mm512_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + __m512 sum0 = _mm512_setzero_ps(); + __m512 sum1 = _mm512_setzero_ps(); + __m512 sum2 = _mm512_setzero_ps(); + __m512 sum3 = _mm512_setzero_ps(); - // Deal with remainder first. `dim` is more than 16, so we have at least one 16-float block, - // so mask loading is guaranteed to be safe + // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so + // this is safe for any dimension. if constexpr (residual) { __mmask16 constexpr mask = (1 << residual) - 1; __m512 v1 = _mm512_maskz_loadu_ps(mask, pVect1); pVect1 += residual; __m512 v2 = _mm512_maskz_loadu_ps(mask, pVect2); pVect2 += residual; - sum512 = _mm512_mul_ps(v1, v2); + sum0 = _mm512_mul_ps(v1, v2); } // We dealt with the residual part. We are left with some multiple of 16 floats. - do { - InnerProductStep(pVect1, pVect2, sum512); - } while (pVect1 < pEnd1); + // The main loop handles 64 floats per iteration; leftover blocks of 16/32/48 floats are + // handled after it. The loops may run zero times (dim can be as small as 8). + while (pVect1 + 64 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + if (pVect1 + 32 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + } + if (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum3); + } + __m512 sum512 = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); return 1.0f - _mm512_reduce_add_ps(sum512); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h index e6eebcc44..8c783b753 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h @@ -23,23 +23,40 @@ double FP64_InnerProductSIMD8_AVX512(const void *pVect1v, const void *pVect2v, s const double *pEnd1 = pVect1 + dimension; - __m512d sum512 = _mm512_setzero_pd(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + __m512d sum0 = _mm512_setzero_pd(); + __m512d sum1 = _mm512_setzero_pd(); + __m512d sum2 = _mm512_setzero_pd(); + __m512d sum3 = _mm512_setzero_pd(); - // Deal with remainder first. `dim` is more than 8, so we have at least one 8-double block, - // so mask loading is guaranteed to be safe + // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so + // this is safe for any dimension. if constexpr (residual) { __mmask8 constexpr mask = (1 << residual) - 1; __m512d v1 = _mm512_maskz_loadu_pd(mask, pVect1); pVect1 += residual; __m512d v2 = _mm512_maskz_loadu_pd(mask, pVect2); pVect2 += residual; - sum512 = _mm512_mul_pd(v1, v2); + sum0 = _mm512_mul_pd(v1, v2); } // We dealt with the residual part. We are left with some multiple of 8 doubles. - do { - InnerProductStep(pVect1, pVect2, sum512); - } while (pVect1 < pEnd1); + // The main loop handles 32 doubles per iteration; leftover blocks of 8/16/24 doubles are + // handled after it. The loops may run zero times (dim can be as small as 4). + while (pVect1 + 32 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + if (pVect1 + 16 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + } + if (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum3); + } + __m512d sum512 = _mm512_add_pd(_mm512_add_pd(sum0, sum1), _mm512_add_pd(sum2, sum3)); return 1.0 - _mm512_reduce_add_pd(sum512); } diff --git a/src/VecSim/spaces/IP/IP_AVX_FP32.h b/src/VecSim/spaces/IP/IP_AVX_FP32.h index e495fb9a1..c0dec6057 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP32.h @@ -24,30 +24,43 @@ float FP32_InnerProductSIMD16_AVX(const void *pVect1v, const void *pVect2v, size const float *pEnd1 = pVect1 + dimension; - __m256 sum256 = _mm256_setzero_ps(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds 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(); - // Deal with 1-7 floats with mask loading, if needed. `dim` is >16, so we have at least one - // 16-float block, so mask loading is guaranteed to be safe. + // Deal with 1-7 floats with mask loading, if needed. The full-width load is safe because + // `dim` is at least 8, so the vector spans at least 8 floats. if constexpr (residual % 8) { __mmask8 constexpr mask = (1 << (residual % 8)) - 1; __m256 v1 = my_mm256_maskz_loadu_ps(pVect1); pVect1 += residual % 8; __m256 v2 = my_mm256_maskz_loadu_ps(pVect2); pVect2 += residual % 8; - sum256 = _mm256_mul_ps(v1, v2); + sum0 = _mm256_mul_ps(v1, v2); } // If the reminder is >=8, have another step of 8 floats if constexpr (residual >= 8) { - InnerProductStep(pVect1, pVect2, sum256); + InnerProductStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - InnerProductStep(pVect1, pVect2, sum256); - InnerProductStep(pVect1, pVect2, sum256); - } while (pVect1 < pEnd1); + // The main loop handles 32 floats per iteration; a possible leftover block of 16 floats is + // handled after it. The loops may run zero times (dim can be as small as 8). + while (pVect1 + 32 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } - return 1.0f - my_mm256_reduce_add_ps(sum256); + __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); + return 1.0f - my_mm256_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/IP/IP_AVX_FP64.h b/src/VecSim/spaces/IP/IP_AVX_FP64.h index b570e6b61..6eb74ed3e 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP64.h @@ -24,10 +24,15 @@ double FP64_InnerProductSIMD8_AVX(const void *pVect1v, const void *pVect2v, size const double *pEnd1 = pVect1 + dimension; - __m256d sum256 = _mm256_setzero_pd(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds be in + // flight at once. + __m256d sum0 = _mm256_setzero_pd(); + __m256d sum1 = _mm256_setzero_pd(); + __m256d sum2 = _mm256_setzero_pd(); + __m256d sum3 = _mm256_setzero_pd(); - // Deal with 1-3 doubles with mask loading, if needed. `dim` is >8, so we have at least one - // 8-double block, so mask loading is guaranteed to be safe. + // Deal with 1-3 doubles with mask loading, if needed. The full-width load is safe because + // `dim` is at least 4, so the vector spans at least 4 doubles. if constexpr (residual % 4) { // _mm256_maskz_loadu_pd is not available in AVX __mmask8 constexpr mask = (1 << (residual % 4)) - 1; @@ -35,21 +40,29 @@ double FP64_InnerProductSIMD8_AVX(const void *pVect1v, const void *pVect2v, size pVect1 += residual % 4; __m256d v2 = my_mm256_maskz_loadu_pd(pVect2); pVect2 += residual % 4; - sum256 = _mm256_mul_pd(v1, v2); + sum0 = _mm256_mul_pd(v1, v2); } // If the reminder is >=4, have another step of 4 doubles if constexpr (residual >= 4) { - InnerProductStep(pVect1, pVect2, sum256); + InnerProductStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 8 doubles. - // In each iteration we calculate 8 doubles = 512 bits. - do { - InnerProductStep(pVect1, pVect2, sum256); - InnerProductStep(pVect1, pVect2, sum256); - } while (pVect1 < pEnd1); + // The main loop handles 16 doubles per iteration; a possible leftover block of 8 doubles is + // handled after it. The loops may run zero times (dim can be as small as 4). + while (pVect1 + 16 <= pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + __m256d sum256 = _mm256_add_pd(_mm256_add_pd(sum0, sum1), _mm256_add_pd(sum2, sum3)); double PORTABLE_ALIGN32 TmpRes[4]; _mm256_store_pd(TmpRes, sum256); double sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3]; diff --git a/src/VecSim/spaces/IP/IP_F16C_FP16.h b/src/VecSim/spaces/IP/IP_F16C_FP16.h index a6f2ec0f4..f1fc9b100 100644 --- a/src/VecSim/spaces/IP/IP_F16C_FP16.h +++ b/src/VecSim/spaces/IP/IP_F16C_FP16.h @@ -31,11 +31,15 @@ float FP16_InnerProductSIMD32_F16C(const void *pVect1v, const void *pVect2v, siz const float16 *pEnd1 = pVect1 + dimension; - auto sum = _mm256_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + auto sum0 = _mm256_setzero_ps(); + auto sum1 = _mm256_setzero_ps(); + auto sum2 = _mm256_setzero_ps(); + auto sum3 = _mm256_setzero_ps(); if constexpr (residual % 8) { - // Deal with remainder first. `dim` is more than 32, so we have at least one block of 32 - // 16-bit float so mask loading is guaranteed to be safe. + // Deal with remainder first. The full-width load of 8 16-bit floats is safe because + // `dim` is at least 8, so the vector spans at least 8 elements. __mmask16 constexpr residuals_mask = (1 << (residual % 8)) - 1; // Convert the first 8 half-floats into floats and store them 256 bits register, // where the floats in the positions corresponding to residuals are zeros. @@ -45,29 +49,31 @@ float FP16_InnerProductSIMD32_F16C(const void *pVect1v, const void *pVect2v, siz auto v2 = _mm256_blend_ps(_mm256_setzero_ps(), _mm256_cvtph_ps(_mm_loadu_si128((__m128i_u const *)pVect2)), residuals_mask); - sum = _mm256_mul_ps(v1, v2); + sum0 = _mm256_mul_ps(v1, v2); pVect1 += residual % 8; pVect2 += residual % 8; } if constexpr (residual >= 8 && residual < 16) { - InnerProductStep(pVect1, pVect2, sum); + InnerProductStep(pVect1, pVect2, sum1); } else if constexpr (residual >= 16 && residual < 24) { - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); } else if constexpr (residual >= 24) { - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); } // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // In every iteration we process 4 chunk of 128bit (32 FP16) - do { - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); - InnerProductStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // In every iteration we process 4 chunk of 128bit (32 FP16). The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); return 1.0f - my_mm256_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h index 45f9f31f4..c71be8bdb 100644 --- a/src/VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h @@ -50,10 +50,15 @@ float SQ8_FP32_InnerProductSIMD16_SSE4_IMP(const void *pVect1v, const void *pVec const float *pVect2 = static_cast(pVect2v); // FP32 query const uint8_t *pEnd1 = pVect1 + dimension; - // Initialize sum accumulator for Σ(q_i * y_i) - __m128 sum = _mm_setzero_ps(); - - // Process residual elements first (1-3 elements) + // Initialize sum accumulators for Σ(q_i * y_i). Four accumulators break the mul->add + // dependency chain, letting more loads/adds 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(); + + // Process residual elements first (1-3 elements). Loads touch only the residual elements, + // so they are safe for any dimension. if constexpr (residual % 4) { __m128 v1_f; __m128 v2; @@ -75,27 +80,31 @@ float SQ8_FP32_InnerProductSIMD16_SSE4_IMP(const void *pVect1v, const void *pVec pVect2 += residual % 4; // Compute q_i * y_i (no dequantization) - sum = _mm_mul_ps(v1_f, v2); + sum0 = _mm_mul_ps(v1_f, v2); } // Handle remaining residual in chunks of 4 (for residual 4-15) if constexpr (residual >= 4) { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); } if constexpr (residual >= 8) { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); } if constexpr (residual >= 12) { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); } - // Process remaining full chunks of 4 elements - // Using do-while since dim > 16 guarantees at least one iteration - do { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // Process remaining full chunks of 16 elements (4x4). The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + InnerProductStepSQ8_FP32(pVect1, pVect2, sum0); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); + InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); + } // Horizontal sum to get Σ(q_i * y_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); float quantized_dot = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3]; diff --git a/src/VecSim/spaces/IP/IP_SSE_FP32.h b/src/VecSim/spaces/IP/IP_SSE_FP32.h index a82cc85ff..6e87d137e 100644 --- a/src/VecSim/spaces/IP/IP_SSE_FP32.h +++ b/src/VecSim/spaces/IP/IP_SSE_FP32.h @@ -23,10 +23,15 @@ float FP32_InnerProductSIMD16_SSE(const void *pVect1v, const void *pVect2v, size const float *pEnd1 = pVect1 + dimension; - __m128 sum_prod = _mm_setzero_ps(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds 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(); - // Deal with %4 remainder first. `dim` is >16, so we have at least one 16-float block, - // so loading 4 floats and then masking them is safe. + // Deal with %4 remainder first. All the loads below touch at most `residual % 4` floats, + // so they are safe for any dimension. if constexpr (residual % 4) { __m128 v1, v2; if constexpr (residual % 4 == 3) { @@ -46,26 +51,28 @@ float FP32_InnerProductSIMD16_SSE(const void *pVect1v, const void *pVect2v, size } pVect1 += residual % 4; pVect2 += residual % 4; - sum_prod = _mm_mul_ps(v1, v2); + sum0 = _mm_mul_ps(v1, v2); } // have another 1, 2 or 3 4-float steps according to residual if constexpr (residual >= 12) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum1); if constexpr (residual >= 8) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum2); if constexpr (residual >= 4) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum3); // We dealt with the residual part. We are left with some multiple of 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - } while (pVect1 < pEnd1); + // In each iteration we calculate 16 floats = 512 bits. The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + __m128 sum_prod = _mm_add_ps(_mm_add_ps(sum0, sum1), _mm_add_ps(sum2, sum3)); // TmpRes must be 16 bytes aligned. float PORTABLE_ALIGN16 TmpRes[4]; _mm_store_ps(TmpRes, sum_prod); diff --git a/src/VecSim/spaces/IP/IP_SSE_FP64.h b/src/VecSim/spaces/IP/IP_SSE_FP64.h index eb0ebab7f..f90fd6aea 100644 --- a/src/VecSim/spaces/IP/IP_SSE_FP64.h +++ b/src/VecSim/spaces/IP/IP_SSE_FP64.h @@ -24,7 +24,12 @@ double FP64_InnerProductSIMD8_SSE(const void *pVect1v, const void *pVect2v, size const double *pEnd1 = pVect1 + dimension; - __m128d sum_prod = _mm_setzero_pd(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds be in + // flight at once. + __m128d sum0 = _mm_setzero_pd(); + __m128d sum1 = _mm_setzero_pd(); + __m128d sum2 = _mm_setzero_pd(); + __m128d sum3 = _mm_setzero_pd(); // If residual is odd, we load 1 double and set the last one to 0 if constexpr (residual % 2 == 1) { @@ -32,26 +37,28 @@ double FP64_InnerProductSIMD8_SSE(const void *pVect1v, const void *pVect2v, size pVect1++; __m128d v2 = _mm_load_sd(pVect2); pVect2++; - sum_prod = _mm_mul_pd(v1, v2); + sum0 = _mm_mul_pd(v1, v2); } // have another 1, 2 or 3 2-double steps according to residual if constexpr (residual >= 6) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum1); if constexpr (residual >= 4) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum2); if constexpr (residual >= 2) - InnerProductStep(pVect1, pVect2, sum_prod); + InnerProductStep(pVect1, pVect2, sum3); // We dealt with the residual part. We are left with some multiple of 8 doubles. - // In each iteration we calculate 8 doubles = 512 bits in total. - do { - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - InnerProductStep(pVect1, pVect2, sum_prod); - } while (pVect1 < pEnd1); + // In each iteration we calculate 8 doubles = 512 bits in total. The loop may run zero times + // (dim can be as small as 4). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); + InnerProductStep(pVect1, pVect2, sum2); + InnerProductStep(pVect1, pVect2, sum3); + } + __m128d sum_prod = _mm_add_pd(_mm_add_pd(sum0, sum1), _mm_add_pd(sum2, sum3)); double PORTABLE_ALIGN16 TmpRes[2]; _mm_store_pd(TmpRes, sum_prod); double sum = TmpRes[0] + TmpRes[1]; diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index e13d57326..6706d8f31 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -69,8 +69,9 @@ dist_func_t IP_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment, #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 16 floats. If we have less, we use the naive implementation. - if (dim < 16) { + // Optimizations assume at least 8 elements (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 8) { return ret_dist_func; } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. @@ -137,8 +138,9 @@ dist_func_t Cosine_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignm #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 16 floats. If we have less, we use the naive implementation. - if (dim < 16) { + // Optimizations assume at least 8 elements (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 8) { return ret_dist_func; } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. @@ -447,8 +449,9 @@ dist_func_t IP_FP32_GetDistFunc(size_t dim, unsigned char *alignment, con #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 16 floats. If we have less, we use the naive implementation. - if (dim < 16) { + // Optimizations assume at least 8 floats (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 8) { return ret_dist_func; } #ifdef OPT_AVX512F @@ -506,8 +509,9 @@ dist_func_t IP_FP64_GetDistFunc(size_t dim, unsigned char *alignment, #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 doubles. If we have less, we use the naive implementation. - if (dim < 8) { + // Optimizations assume at least 4 doubles (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 4) { return ret_dist_func; } #ifdef OPT_AVX512F @@ -626,28 +630,27 @@ dist_func_t IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con #endif #if defined(CPU_FEATURES_ARCH_X86_64) - // Optimizations assume at least 32 16FPs. If we have less, we use the naive implementation. - if (dim < 32) { - return ret_dist_func; - } + // Each tier has a minimal dimension implied by its residual handling: the AVX512FP16_VL + // kernel loads full 512-bit blocks (32 elements), the AVX512F kernel loads full 256-bit + // blocks (16 elements), and the F16C kernel loads full 128-bit blocks (8 elements). #ifdef OPT_AVX512_FP16_VL // More details about the dimension limitation can be found in this PR's description: // https://github.com/RedisAI/VectorSimilarity/pull/477 - if (features.avx512_fp16 && features.avx512vl) { + if (dim >= 32 && features.avx512_fp16 && features.avx512vl) { if (dim % 32 == 0) // no point in aligning if we have an offsetting residual *alignment = 32 * sizeof(float16); // handles 32 floats return Choose_FP16_IP_implementation_AVX512FP16_VL(dim); } #endif #ifdef OPT_AVX512F - if (features.avx512f) { + if (dim >= 16 && features.avx512f) { if (dim % 32 == 0) // no point in aligning if we have an offsetting residual *alignment = 32 * sizeof(float16); // handles 32 floats return Choose_FP16_IP_implementation_AVX512F(dim); } #endif #ifdef OPT_F16C - if (features.f16c && features.fma3 && features.avx) { + if (dim >= 8 && features.f16c && features.fma3 && features.avx) { if (dim % 16 == 0) // no point in aligning if we have an offsetting residual *alignment = 16 * sizeof(float16); // handles 16 floats return Choose_FP16_IP_implementation_F16C(dim); diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h index e2a21414f..91e0fa527 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h @@ -32,11 +32,15 @@ float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t d const float16 *pEnd1 = pVect1 + dimension; - auto sum = _mm512_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + auto sum0 = _mm512_setzero_ps(); + auto sum1 = _mm512_setzero_ps(); + auto sum2 = _mm512_setzero_ps(); + auto sum3 = _mm512_setzero_ps(); if constexpr (residual % 16) { - // Deal with remainder first. `dim` is more than 32, so we have at least one block of 32 - // 16-bit float so mask loading is guaranteed to be safe. + // Deal with remainder first. The full-width load of 16 16-bit floats is safe because + // `dim` is at least 16, so the vector spans at least 16 elements. __mmask16 constexpr residuals_mask = (1 << (residual % 16)) - 1; // Convert the first half-floats in the residual positions into floats and store them // 512 bits register, where the floats in the positions corresponding to the non-residuals @@ -46,20 +50,28 @@ float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t d auto v2 = _mm512_maskz_mov_ps(residuals_mask, _mm512_cvtph_ps(_mm256_lddqu_si256((__m256i *)pVect2))); auto c = _mm512_sub_ps(v1, v2); - sum = _mm512_mul_ps(c, c); + sum0 = _mm512_mul_ps(c, c); pVect1 += residual % 16; pVect2 += residual % 16; } if constexpr (residual >= 16) { - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // In every iteration we process 2 chunk of 256bit (32 FP16) - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 64 elements per iteration; a possible leftover block of 32 is + // handled after it. The loops may run zero times (dim can be as small as 16). + while (pVect1 + 64 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + auto sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); return _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h index 0100e4264..9fd276359 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h @@ -25,10 +25,14 @@ float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t d const float *pEnd1 = pVect1 + dimension; - __m512 sum = _mm512_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + __m512 sum0 = _mm512_setzero_ps(); + __m512 sum1 = _mm512_setzero_ps(); + __m512 sum2 = _mm512_setzero_ps(); + __m512 sum3 = _mm512_setzero_ps(); - // Deal with remainder first. `dim` is more than 16, so we have at least one 16-float block, - // so mask loading is guaranteed to be safe + // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so + // this is safe for any dimension. if constexpr (residual) { __mmask16 constexpr mask = (1 << residual) - 1; __m512 v1 = _mm512_maskz_loadu_ps(mask, pVect1); @@ -36,13 +40,26 @@ float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t d __m512 v2 = _mm512_maskz_loadu_ps(mask, pVect2); pVect2 += residual; __m512 diff = _mm512_sub_ps(v1, v2); - sum = _mm512_mul_ps(diff, diff); + sum0 = _mm512_mul_ps(diff, diff); } // We dealt with the residual part. We are left with some multiple of 16 floats. - do { - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 64 floats per iteration; leftover blocks of 16/32/48 floats are + // handled after it. The loops may run zero times (dim can be as small as 8). + while (pVect1 + 64 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + if (pVect1 + 32 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + } + if (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum3); + } + __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); return _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h index 1a54c7048..5d3d41633 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h @@ -25,10 +25,14 @@ double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t d const double *pEnd1 = pVect1 + dimension; - __m512d sum = _mm512_setzero_pd(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + __m512d sum0 = _mm512_setzero_pd(); + __m512d sum1 = _mm512_setzero_pd(); + __m512d sum2 = _mm512_setzero_pd(); + __m512d sum3 = _mm512_setzero_pd(); - // Deal with remainder first. `dim` is more than 8, so we have at least one 8-double block, - // so mask loading is guaranteed to be safe + // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so + // this is safe for any dimension. if constexpr (residual) { __mmask8 constexpr mask = (1 << residual) - 1; __m512d v1 = _mm512_maskz_loadu_pd(mask, pVect1); @@ -36,13 +40,26 @@ double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t d __m512d v2 = _mm512_maskz_loadu_pd(mask, pVect2); pVect2 += residual; __m512d diff = _mm512_sub_pd(v1, v2); - sum = _mm512_mul_pd(diff, diff); + sum0 = _mm512_mul_pd(diff, diff); } // We dealt with the residual part. We are left with some multiple of 8 doubles. - do { - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 32 doubles per iteration; leftover blocks of 8/16/24 doubles are + // handled after it. The loops may run zero times (dim can be as small as 4). + while (pVect1 + 32 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + if (pVect1 + 16 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + } + if (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum3); + } + __m512d sum = _mm512_add_pd(_mm512_add_pd(sum0, sum1), _mm512_add_pd(sum2, sum3)); return _mm512_reduce_add_pd(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX_FP32.h b/src/VecSim/spaces/L2/L2_AVX_FP32.h index 4751d4726..b4d4d8ab7 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP32.h @@ -26,9 +26,15 @@ float FP32_L2SqrSIMD16_AVX(const void *pVect1v, const void *pVect2v, size_t dime const float *pEnd1 = pVect1 + dimension; - __m256 sum = _mm256_setzero_ps(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds 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(); - // Deal with 1-7 floats with mask loading, if needed + // Deal with 1-7 floats with mask loading, if needed. The full-width load is safe because + // `dim` is at least 8, so the vector spans at least 8 floats. if constexpr (residual % 8) { __mmask8 constexpr mask8 = (1 << (residual % 8)) - 1; __m256 v1 = my_mm256_maskz_loadu_ps(pVect1); @@ -36,20 +42,28 @@ float FP32_L2SqrSIMD16_AVX(const void *pVect1v, const void *pVect2v, size_t dime __m256 v2 = my_mm256_maskz_loadu_ps(pVect2); pVect2 += residual % 8; __m256 diff = _mm256_sub_ps(v1, v2); - sum = _mm256_mul_ps(diff, diff); + sum0 = _mm256_mul_ps(diff, diff); } // If the reminder is >=8, have another step of 8 floats if constexpr (residual >= 8) { - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 32 floats per iteration; a possible leftover block of 16 floats is + // handled after it. The loops may run zero times (dim can be as small as 8). + while (pVect1 + 32 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); return my_mm256_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX_FP64.h b/src/VecSim/spaces/L2/L2_AVX_FP64.h index 09257dca5..e2509afcd 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP64.h @@ -26,9 +26,15 @@ double FP64_L2SqrSIMD8_AVX(const void *pVect1v, const void *pVect2v, size_t dime const double *pEnd1 = pVect1 + dimension; - __m256d sum = _mm256_setzero_pd(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds be in + // flight at once. + __m256d sum0 = _mm256_setzero_pd(); + __m256d sum1 = _mm256_setzero_pd(); + __m256d sum2 = _mm256_setzero_pd(); + __m256d sum3 = _mm256_setzero_pd(); - // Deal with 1-3 doubles with mask loading, if needed + // Deal with 1-3 doubles with mask loading, if needed. The full-width load is safe because + // `dim` is at least 4, so the vector spans at least 4 doubles. if constexpr (residual % 4) { // _mm256_maskz_loadu_pd is not available in AVX __mmask8 constexpr mask4 = (1 << (residual % 4)) - 1; @@ -37,21 +43,29 @@ double FP64_L2SqrSIMD8_AVX(const void *pVect1v, const void *pVect2v, size_t dime __m256d v2 = my_mm256_maskz_loadu_pd(pVect2); pVect2 += residual % 4; __m256d diff = _mm256_sub_pd(v1, v2); - sum = _mm256_mul_pd(diff, diff); + sum0 = _mm256_mul_pd(diff, diff); } // If the reminder is >=4, have another step of 4 doubles if constexpr (residual >= 4) { - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); } // We dealt with the residual part. We are left with some multiple of 8 doubles. - // In each iteration we calculate 8 doubles = 512 bits. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // The main loop handles 16 doubles per iteration; a possible leftover block of 8 doubles is + // handled after it. The loops may run zero times (dim can be as small as 4). + while (pVect1 + 16 <= pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + if (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + __m256d sum = _mm256_add_pd(_mm256_add_pd(sum0, sum1), _mm256_add_pd(sum2, sum3)); double PORTABLE_ALIGN32 TmpRes[4]; _mm256_store_pd(TmpRes, sum); return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3]; diff --git a/src/VecSim/spaces/L2/L2_F16C_FP16.h b/src/VecSim/spaces/L2/L2_F16C_FP16.h index c193b9cec..7f57657d3 100644 --- a/src/VecSim/spaces/L2/L2_F16C_FP16.h +++ b/src/VecSim/spaces/L2/L2_F16C_FP16.h @@ -32,11 +32,15 @@ float FP16_L2SqrSIMD32_F16C(const void *pVect1v, const void *pVect2v, size_t dim const float16 *pEnd1 = pVect1 + dimension; - auto sum = _mm256_setzero_ps(); + // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + auto sum0 = _mm256_setzero_ps(); + auto sum1 = _mm256_setzero_ps(); + auto sum2 = _mm256_setzero_ps(); + auto sum3 = _mm256_setzero_ps(); if constexpr (residual % 8) { - // Deal with remainder first. `dim` is more than 32, so we have at least one block of 32 - // 16-bit float so mask loading is guaranteed to be safe. + // Deal with remainder first. The full-width load of 8 16-bit floats is safe because + // `dim` is at least 8, so the vector spans at least 8 elements. __mmask16 constexpr residuals_mask = (1 << (residual % 8)) - 1; // Convert the first 8 half-floats into floats and store them 256 bits register, // where the floats in the positions corresponding to residuals are zeros. @@ -48,28 +52,30 @@ float FP16_L2SqrSIMD32_F16C(const void *pVect1v, const void *pVect2v, size_t dim residuals_mask); // sum = (v1 * v2)^2 + sum auto c = _mm256_sub_ps(v1, v2); - sum = _mm256_fmadd_ps(c, c, sum); + sum0 = _mm256_fmadd_ps(c, c, sum0); pVect1 += residual % 8; pVect2 += residual % 8; } if constexpr (residual >= 8 && residual < 16) { - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); } else if constexpr (residual >= 16 && residual < 24) { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); } else if constexpr (residual >= 24) { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); } // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // In every iteration we process 4 chunk of 128bit (32 FP16) - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // In every iteration we process 4 chunk of 128bit (32 FP16). The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); return my_mm256_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_SSE_FP32.h b/src/VecSim/spaces/L2/L2_SSE_FP32.h index e04cc4fe5..b2548062f 100644 --- a/src/VecSim/spaces/L2/L2_SSE_FP32.h +++ b/src/VecSim/spaces/L2/L2_SSE_FP32.h @@ -24,10 +24,16 @@ float FP32_L2SqrSIMD16_SSE(const void *pVect1v, const void *pVect2v, size_t dime const float *pEnd1 = pVect1 + dimension; - __m128 sum = _mm_setzero_ps(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds 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(); - // Deal with %4 remainder first. `dim` is >16, so we have at least one 16-float block, - // so loading 4 floats and then masking them is safe. + // Deal with %4 remainder first. All the loads below touch at most `residual % 4` floats, + // except for the reversed load in the residual == 3 case, which reads a full (aligned) + // 16-byte block. This is safe because `dim` is at least 8 whenever residual % 4 == 3. if constexpr (residual % 4) { __m128 v1, v2, diff; if constexpr (residual % 4 == 3) { @@ -48,26 +54,28 @@ float FP32_L2SqrSIMD16_SSE(const void *pVect1v, const void *pVect2v, size_t dime pVect1 += residual % 4; pVect2 += residual % 4; diff = _mm_sub_ps(v1, v2); - sum = _mm_mul_ps(diff, diff); + sum0 = _mm_mul_ps(diff, diff); } // have another 1, 2 or 3 4-floats steps according to residual if constexpr (residual >= 12) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); if constexpr (residual >= 8) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum2); if constexpr (residual >= 4) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum3); // We dealt with the residual part. We are left with some multiple of 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // In each iteration we calculate 16 floats = 512 bits. The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + __m128 sum = _mm_add_ps(_mm_add_ps(sum0, sum1), _mm_add_ps(sum2, sum3)); // TmpRes must be 16 bytes aligned float PORTABLE_ALIGN16 TmpRes[4]; _mm_store_ps(TmpRes, sum); diff --git a/src/VecSim/spaces/L2/L2_SSE_FP64.h b/src/VecSim/spaces/L2/L2_SSE_FP64.h index 4640c1cbd..5054c058c 100644 --- a/src/VecSim/spaces/L2/L2_SSE_FP64.h +++ b/src/VecSim/spaces/L2/L2_SSE_FP64.h @@ -24,7 +24,12 @@ double FP64_L2SqrSIMD8_SSE(const void *pVect1v, const void *pVect2v, size_t dime const double *pEnd1 = pVect1 + dimension; - __m128d sum = _mm_setzero_pd(); + // Four accumulators break the mul->add dependency chain, letting more loads/adds be in + // flight at once. + __m128d sum0 = _mm_setzero_pd(); + __m128d sum1 = _mm_setzero_pd(); + __m128d sum2 = _mm_setzero_pd(); + __m128d sum3 = _mm_setzero_pd(); // If residual is odd, we load 1 double and set the last one to 0 if constexpr (residual % 2 == 1) { @@ -33,26 +38,28 @@ double FP64_L2SqrSIMD8_SSE(const void *pVect1v, const void *pVect2v, size_t dime __m128d v2 = _mm_load_sd(pVect2); pVect2++; __m128d diff = _mm_sub_pd(v1, v2); - sum = _mm_mul_pd(diff, diff); + sum0 = _mm_mul_pd(diff, diff); } // have another 1, 2 or 3 2-double steps according to residual if constexpr (residual >= 6) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); if constexpr (residual >= 4) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum2); if constexpr (residual >= 2) - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum3); // We dealt with the residual part. We are left with some multiple of 8 doubles. - // In each iteration we calculate 8 doubles = 512 bits in total. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (pVect1 < pEnd1); + // In each iteration we calculate 8 doubles = 512 bits in total. The loop may run zero times + // (dim can be as small as 4). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); + L2SqrStep(pVect1, pVect2, sum2); + L2SqrStep(pVect1, pVect2, sum3); + } + __m128d sum = _mm_add_pd(_mm_add_pd(sum0, sum1), _mm_add_pd(sum2, sum3)); // TmpRes must be 16 bytes aligned double PORTABLE_ALIGN16 TmpRes[2]; _mm_store_pd(TmpRes, sum); diff --git a/src/VecSim/spaces/L2_space.cpp b/src/VecSim/spaces/L2_space.cpp index 53fa2d873..07e638cba 100644 --- a/src/VecSim/spaces/L2_space.cpp +++ b/src/VecSim/spaces/L2_space.cpp @@ -68,9 +68,9 @@ dist_func_t L2_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment, #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 16 floats. If we have less, we use the naive implementation. - - if (dim < 16) { + // Optimizations assume at least 8 elements (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 8) { return ret_dist_func; } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. @@ -210,9 +210,9 @@ dist_func_t L2_FP32_GetDistFunc(size_t dim, unsigned char *alignment, con #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 16 floats. If we have less, we use the naive implementation. - - if (dim < 16) { + // Optimizations assume at least 8 floats (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 8) { return ret_dist_func; } #ifdef OPT_AVX512F @@ -269,8 +269,9 @@ dist_func_t L2_FP64_GetDistFunc(size_t dim, unsigned char *alignment, #endif #ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 doubles. If we have less, we use the naive implementation. - if (dim < 8) { + // Optimizations assume at least 4 doubles (see the residual handling in the kernels). + // Below that, the scalar implementation is at least as fast anyway. + if (dim < 4) { return ret_dist_func; } #ifdef OPT_AVX512F @@ -381,28 +382,27 @@ dist_func_t L2_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con #endif // CPU_FEATURES_ARCH_AARCH64 #if defined(CPU_FEATURES_ARCH_X86_64) - // Optimizations assume at least 32 16FPs. If we have less, we use the naive implementation. - if (dim < 32) { - return ret_dist_func; - } + // Each tier has a minimal dimension implied by its residual handling: the AVX512FP16_VL + // kernel loads full 512-bit blocks (32 elements), the AVX512F kernel loads full 256-bit + // blocks (16 elements), and the F16C kernel loads full 128-bit blocks (8 elements). #ifdef OPT_AVX512_FP16_VL // More details about the dimension limitation can be found in this PR's description: // https://github.com/RedisAI/VectorSimilarity/pull/477 - if (features.avx512_fp16 && features.avx512vl) { + if (dim >= 32 && features.avx512_fp16 && features.avx512vl) { if (dim % 32 == 0) // no point in aligning if we have an offsetting residual *alignment = 32 * sizeof(float16); // handles 32 floats return Choose_FP16_L2_implementation_AVX512FP16_VL(dim); } #endif #ifdef OPT_AVX512F - if (features.avx512f) { + if (dim >= 16 && features.avx512f) { if (dim % 32 == 0) // no point in aligning if we have an offsetting residual *alignment = 32 * sizeof(float16); // handles 32 floats return Choose_FP16_L2_implementation_AVX512F(dim); } #endif #ifdef OPT_F16C - if (features.f16c && features.fma3 && features.avx) { + if (dim >= 8 && features.f16c && features.fma3 && features.avx) { if (dim % 16 == 0) // no point in aligning if we have an offsetting residual *alignment = 16 * sizeof(float16); // handles 16 floats return Choose_FP16_L2_implementation_F16C(dim); diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 9a16bf30d..b67ae4ec8 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -582,44 +582,27 @@ TEST_F(SpacesTest, GetDistFuncSQ8FP16Asymmetric) { #ifdef CPU_FEATURES_ARCH_X86_64 TEST_F(SpacesTest, smallDimChooser) { - // Verify that small dimensions gets the no optimization function. - for (size_t dim = 1; dim < 8; dim++) { - ASSERT_EQ(L2_FP32_GetDistFunc(dim), FP32_L2Sqr); + // Verify that dimensions below each type's SIMD threshold get the no optimization function. + // FP64 is optimized from dim >= 4. + for (size_t dim = 1; dim < 4; dim++) { ASSERT_EQ(L2_FP64_GetDistFunc(dim), FP64_L2Sqr); - ASSERT_EQ(L2_BF16_GetDistFunc(dim), BF16_L2Sqr_LittleEndian); - ASSERT_EQ(L2_FP16_GetDistFunc(dim), FP16_L2Sqr); - ASSERT_EQ(L2_INT8_GetDistFunc(dim), INT8_L2Sqr); - ASSERT_EQ(L2_UINT8_GetDistFunc(dim), UINT8_L2Sqr); - ASSERT_EQ(IP_FP32_GetDistFunc(dim), FP32_InnerProduct); ASSERT_EQ(IP_FP64_GetDistFunc(dim), FP64_InnerProduct); - ASSERT_EQ(IP_BF16_GetDistFunc(dim), BF16_InnerProduct_LittleEndian); - ASSERT_EQ(IP_FP16_GetDistFunc(dim), FP16_InnerProduct); - ASSERT_EQ(IP_INT8_GetDistFunc(dim), INT8_InnerProduct); - ASSERT_EQ(IP_UINT8_GetDistFunc(dim), UINT8_InnerProduct); - ASSERT_EQ(Cosine_INT8_GetDistFunc(dim), INT8_Cosine); - ASSERT_EQ(Cosine_UINT8_GetDistFunc(dim), UINT8_Cosine); } - for (size_t dim = 8; dim < 16; dim++) { + // FP32 and FP16 are optimized from dim >= 8 (FP16 only on machines with F16C; for + // 8 <= dim < 32 the chosen FP16 function depends on the available features, so we only + // assert the range that is naive regardless of features). + for (size_t dim = 1; dim < 8; dim++) { ASSERT_EQ(L2_FP32_GetDistFunc(dim), FP32_L2Sqr); - ASSERT_EQ(L2_BF16_GetDistFunc(dim), BF16_L2Sqr_LittleEndian); - ASSERT_EQ(L2_FP16_GetDistFunc(dim), FP16_L2Sqr); - ASSERT_EQ(L2_INT8_GetDistFunc(dim), INT8_L2Sqr); - ASSERT_EQ(L2_UINT8_GetDistFunc(dim), UINT8_L2Sqr); ASSERT_EQ(IP_FP32_GetDistFunc(dim), FP32_InnerProduct); - ASSERT_EQ(IP_BF16_GetDistFunc(dim), BF16_InnerProduct_LittleEndian); + ASSERT_EQ(L2_FP16_GetDistFunc(dim), FP16_L2Sqr); ASSERT_EQ(IP_FP16_GetDistFunc(dim), FP16_InnerProduct); - ASSERT_EQ(IP_INT8_GetDistFunc(dim), INT8_InnerProduct); - ASSERT_EQ(IP_UINT8_GetDistFunc(dim), UINT8_InnerProduct); - ASSERT_EQ(Cosine_INT8_GetDistFunc(dim), INT8_Cosine); - ASSERT_EQ(Cosine_UINT8_GetDistFunc(dim), UINT8_Cosine); } - for (size_t dim = 16; dim < 32; dim++) { + // BF16, INT8 and UINT8 are optimized from dim >= 32. + for (size_t dim = 1; dim < 32; dim++) { ASSERT_EQ(L2_BF16_GetDistFunc(dim), BF16_L2Sqr_LittleEndian); - ASSERT_EQ(L2_FP16_GetDistFunc(dim), FP16_L2Sqr); ASSERT_EQ(L2_INT8_GetDistFunc(dim), INT8_L2Sqr); ASSERT_EQ(L2_UINT8_GetDistFunc(dim), UINT8_L2Sqr); ASSERT_EQ(IP_BF16_GetDistFunc(dim), BF16_InnerProduct_LittleEndian); - ASSERT_EQ(IP_FP16_GetDistFunc(dim), FP16_InnerProduct); ASSERT_EQ(IP_INT8_GetDistFunc(dim), INT8_InnerProduct); ASSERT_EQ(IP_UINT8_GetDistFunc(dim), UINT8_InnerProduct); ASSERT_EQ(Cosine_INT8_GetDistFunc(dim), INT8_Cosine); @@ -830,7 +813,7 @@ TEST_P(FP32SpacesOptimizationTest, FP32InnerProductTest) { } INSTANTIATE_TEST_SUITE_P(FP32OptFuncs, FP32SpacesOptimizationTest, - testing::Range(16UL, 16 * 2UL + 1)); + testing::Range(8UL, 16 * 2UL + 1)); class FP64SpacesOptimizationTest : public testing::TestWithParam {}; @@ -1025,7 +1008,7 @@ TEST_P(FP64SpacesOptimizationTest, FP64InnerProductTest) { } INSTANTIATE_TEST_SUITE_P(FP64OptFuncs, FP64SpacesOptimizationTest, - testing::Range(8UL, 8 * 2UL + 1)); + testing::Range(4UL, 8 * 2UL + 1)); class BF16SpacesOptimizationTest : public testing::TestWithParam {}; @@ -1229,15 +1212,17 @@ TEST_P(FP16SpacesOptimizationTest, FP16InnerProductTest) { // Turn off advanced fp16 flags. They will be tested in the next test. optimization.avx512_fp16 = optimization.avx512vl = 0; #ifdef OPT_AVX512F - if (optimization.avx512f) { + // The AVX512F FP16 tier requires at least 16 elements; below that the dispatcher falls + // through to the next tier. + if (optimization.avx512f && dim >= 16) { unsigned char alignment = 0; arch_opt_func = IP_FP16_GetDistFunc(dim, &alignment, &optimization); ASSERT_EQ(arch_opt_func, Choose_FP16_IP_implementation_AVX512F(dim)) << "Unexpected distance function chosen for dim " << dim; ASSERT_EQ(baseline, arch_opt_func(v1, v2, dim)) << "AVX512 with dim " << dim; ASSERT_EQ(alignment, expected_alignment(512, dim)) << "AVX512 with dim " << dim; - optimization.avx512f = 0; } + optimization.avx512f = 0; #endif #ifdef OPT_F16C if (optimization.f16c && optimization.fma3 && optimization.avx) { @@ -1286,15 +1271,17 @@ TEST_P(FP16SpacesOptimizationTest, FP16L2SqrTest) { // Turn off advanced fp16 flags. They will be tested in the next test. optimization.avx512_fp16 = optimization.avx512vl = 0; #ifdef OPT_AVX512F - if (optimization.avx512f) { + // The AVX512F FP16 tier requires at least 16 elements; below that the dispatcher falls + // through to the next tier. + if (optimization.avx512f && dim >= 16) { unsigned char alignment = 0; arch_opt_func = L2_FP16_GetDistFunc(dim, &alignment, &optimization); ASSERT_EQ(arch_opt_func, Choose_FP16_L2_implementation_AVX512F(dim)) << "Unexpected distance function chosen for dim " << dim; ASSERT_EQ(baseline, arch_opt_func(v1, v2, dim)) << "AVX512 with dim " << dim; ASSERT_EQ(alignment, expected_alignment(512, dim)) << "AVX512 with dim " << dim; - optimization.avx512f = 0; } + optimization.avx512f = 0; #endif #ifdef OPT_F16C if (optimization.f16c && optimization.fma3 && optimization.avx) { @@ -1352,7 +1339,7 @@ TEST_P(FP16SpacesOptimizationTest, FP16L2SqrTest) { } INSTANTIATE_TEST_SUITE_P(FP16OptFuncs, FP16SpacesOptimizationTest, - testing::Range(32UL, 32 * 2UL + 1)); + testing::Range(8UL, 32 * 2UL + 1)); /** Since we are handling floats, the order of summation affect on the final result. * This is very significant when the entries are half precision floats, since the accumulated @@ -2374,7 +2361,7 @@ TEST_P(SQ8_FP32_SpacesOptimizationTest, SQ8_FP32_InnerProductTest) { // Instantiate the test suite with dimensions to test INSTANTIATE_TEST_SUITE_P(SQ8_FP32_Test, SQ8_FP32_SpacesOptimizationTest, - testing::Range(16UL, 16 * 2UL + 1)); + testing::Range(8UL, 16 * 2UL + 1)); TEST_P(SQ8_FP32_SpacesOptimizationTest, SQ8_FP32_CosineTest) { auto optimization = getCpuOptimizationFeatures(); From 1afb93930aec54dc1a575dfbe6eeef1bf2a109ea Mon Sep 17 00:00:00 2001 From: GuyAv46 Date: Thu, 2 Jul 2026 15:38:29 +0300 Subject: [PATCH 2/3] Resolve residual block handling at compile time via wider dispatch modulus Match each kernel's dispatch modulus (CHOOSE_IMPLEMENTATION chunk) to its main-loop stride, so all residual-based conditions - including the leftover full blocks that were previously handled with runtime tail checks after the main loop - are hard-coded through the residual template parameter: - AVX FP32 / AVX2 SQ8_FP32(+FMA): dispatch mod 32 (4x8 per iteration) - AVX FP64: dispatch mod 16 (4x4) - AVX512F FP32 / FP16 / VNNI SQ8_FP32: dispatch mod 64 (4x16) - AVX512F FP64: dispatch mod 32 (4x8) The main loops are now branch-free except for the loop condition itself, and the residual (0..stride-1) is fully unrolled at compile time. Loop/tail coverage re-verified by simulation for all dims up to 4096. Co-Authored-By: Claude Fable 5 --- src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h | 24 +++++++------ src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h | 24 +++++++------ .../IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h | 36 ++++++++++--------- src/VecSim/spaces/IP/IP_AVX512F_FP16.h | 21 ++++++----- src/VecSim/spaces/IP/IP_AVX512F_FP32.h | 34 ++++++++++-------- src/VecSim/spaces/IP/IP_AVX512F_FP64.h | 34 ++++++++++-------- src/VecSim/spaces/IP/IP_AVX_FP32.h | 22 ++++++------ src/VecSim/spaces/IP/IP_AVX_FP64.h | 22 ++++++------ src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h | 2 +- src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h | 2 +- .../L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h | 2 +- src/VecSim/spaces/L2/L2_AVX512F_FP16.h | 21 ++++++----- src/VecSim/spaces/L2/L2_AVX512F_FP32.h | 34 ++++++++++-------- src/VecSim/spaces/L2/L2_AVX512F_FP64.h | 34 ++++++++++-------- src/VecSim/spaces/L2/L2_AVX_FP32.h | 22 ++++++------ src/VecSim/spaces/L2/L2_AVX_FP64.h | 22 ++++++------ src/VecSim/spaces/functions/AVX.cpp | 8 ++--- src/VecSim/spaces/functions/AVX2.cpp | 6 ++-- src/VecSim/spaces/functions/AVX2_FMA.cpp | 6 ++-- src/VecSim/spaces/functions/AVX512F.cpp | 12 +++---- .../spaces/functions/AVX512F_BW_VL_VNNI.cpp | 6 ++-- 21 files changed, 216 insertions(+), 178 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h index a948e6166..8aebf768c 100644 --- a/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h @@ -46,7 +46,7 @@ static inline void InnerProductStepSQ8_FMA(const uint8_t *&pVect1, const float * } // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..15 +template // 0..31 float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, size_t dimension) { const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage const float *pVect2 = static_cast(pVect2v); // FP32 query @@ -79,24 +79,26 @@ float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, siz sum0 = _mm256_mul_ps(v1_f, v2); } - // If the residual is >=8, have another step of 8 floats + // Handle the remaining full 8-element blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { InnerProductStepSQ8_FMA(pVect1, pVect2, sum1); } + if constexpr (residual >= 16) { + InnerProductStepSQ8_FMA(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { + InnerProductStepSQ8_FMA(pVect1, pVect2, sum3); + } - // We are left with some multiple of 16 elements. The main loop handles 32 per iteration; - // a possible leftover block of 16 is handled after it. The loops may run zero times + // 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 + 32 <= pEnd1) { + while (pVect1 < pEnd1) { InnerProductStepSQ8_FMA(pVect1, pVect2, sum0); InnerProductStepSQ8_FMA(pVect1, pVect2, sum1); InnerProductStepSQ8_FMA(pVect1, pVect2, sum2); InnerProductStepSQ8_FMA(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { - InnerProductStepSQ8_FMA(pVect1, pVect2, sum2); - InnerProductStepSQ8_FMA(pVect1, pVect2, sum3); - } // Reduce to get Σ(q_i * y_i) float quantized_dot = @@ -115,13 +117,13 @@ float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, siz return min_val * y_sum + delta * quantized_dot; } -template // 0..15 +template // 0..31 float SQ8_FP32_InnerProductSIMD16_AVX2_FMA(const void *pVect1v, const void *pVect2v, size_t dimension) { return 1.0f - SQ8_FP32_InnerProductImp_FMA(pVect1v, pVect2v, dimension); } -template // 0..15 +template // 0..31 float SQ8_FP32_CosineSIMD16_AVX2_FMA(const void *pVect1v, const void *pVect2v, size_t dimension) { // Cosine distance = 1 - IP (vectors are pre-normalized) return SQ8_FP32_InnerProductSIMD16_AVX2_FMA(pVect1v, pVect2v, dimension); diff --git a/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h index 3342e9e20..5152e4544 100644 --- a/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h @@ -46,7 +46,7 @@ static inline void InnerProductStepSQ8_FP32(const uint8_t *&pVect1, const float } // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..15 +template // 0..31 float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) { const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage const float *pVect2 = static_cast(pVect2v); // FP32 query @@ -79,24 +79,26 @@ float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, si sum0 = _mm256_mul_ps(v1_f, v2); } - // If the residual is >=8, have another step of 8 floats + // Handle the remaining full 8-element blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); } + if constexpr (residual >= 16) { + InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { + InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); + } - // We are left with some multiple of 16 elements. The main loop handles 32 per iteration; - // a possible leftover block of 16 is handled after it. The loops may run zero times + // 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 + 32 <= pEnd1) { + while (pVect1 < pEnd1) { InnerProductStepSQ8_FP32(pVect1, pVect2, sum0); InnerProductStepSQ8_FP32(pVect1, pVect2, sum1); InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { - InnerProductStepSQ8_FP32(pVect1, pVect2, sum2); - InnerProductStepSQ8_FP32(pVect1, pVect2, sum3); - } // Reduce to get Σ(q_i * y_i) float quantized_dot = @@ -115,12 +117,12 @@ float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, si return min_val * y_sum + delta * quantized_dot; } -template // 0..15 +template // 0..31 float SQ8_FP32_InnerProductSIMD16_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) { return 1.0f - SQ8_FP32_InnerProductImp_AVX2(pVect1v, pVect2v, dimension); } -template // 0..15 +template // 0..31 float SQ8_FP32_CosineSIMD16_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) { // Calculate inner product using common implementation with normalization return SQ8_FP32_InnerProductSIMD16_AVX2(pVect1v, pVect2v, dimension); diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h index 75fe35b44..da7527e7c 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h @@ -45,7 +45,7 @@ static inline void SQ8_FP32_InnerProductStep(const uint8_t *&pVec1, const float // Common implementation for both inner product and cosine similarity // pVec1v = SQ8 storage, pVec2v = FP32 query -template // 0..15 +template // 0..63 float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, size_t dimension) { const uint8_t *pVec1 = static_cast(pVec1v); // SQ8 storage const float *pVec2 = static_cast(pVec2v); // FP32 query @@ -58,9 +58,9 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si __m512 sum2 = _mm512_setzero_ps(); __m512 sum3 = _mm512_setzero_ps(); - // Handle residual elements first (0 to 15) - if constexpr (residual > 0) { - __mmask16 mask = (1U << residual) - 1; + // 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 @@ -75,24 +75,28 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si // Compute q_i * y_i with mask (no dequantization) sum0 = _mm512_maskz_mul_ps(mask, v1_f, v2); - pVec1 += residual; - pVec2 += residual; + pVec1 += residual % 16; + pVec2 += residual % 16; } - // Process full chunks of 16 elements. The main loop handles 64 per iteration; leftover - // blocks of 16/32/48 are handled after it. The loops may run zero times (dim can be as - // small as 8). - while (pVec1 + 64 <= pEnd1) { - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum0); + // Handle the remaining full 16-element blocks of the residual (compile-time resolved). + if constexpr (residual >= 16) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); + } + if constexpr (residual >= 32) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); + } + if constexpr (residual >= 48) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); } - if (pVec1 + 32 <= pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 64 elements. + // In each iteration we calculate 64 elements = 4 chunks of 16. The loop may run zero times + // (dim can be as small as 8). + while (pVec1 < pEnd1) { + SQ8_FP32_InnerProductStep(pVec1, pVec2, sum0); SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); - } - if (pVec1 < pEnd1) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); } @@ -115,14 +119,14 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si return min_val * y_sum + delta * quantized_dot; } -template // 0..15 +template // 0..63 float SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI(const void *pVec1v, const void *pVec2v, size_t dimension) { // The inner product similarity is 1 - ip return 1.0f - SQ8_FP32_InnerProductImp_AVX512(pVec1v, pVec2v, dimension); } -template // 0..15 +template // 0..63 float SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI(const void *pVec1v, const void *pVec2v, size_t dimension) { // Cosine distance = 1 - IP (vectors are pre-normalized) diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h index 5f298c303..37d6c8f96 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h @@ -24,7 +24,7 @@ static void InnerProductStep(float16 *&pVect1, float16 *&pVect2, __m512 &sum) { pVect2 += 16; } -template // 0..31 +template // 0..63 float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { auto *pVect1 = (float16 *)pVect1v; auto *pVect2 = (float16 *)pVect2v; @@ -52,20 +52,23 @@ float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, s pVect1 += residual % 16; pVect2 += residual % 16; } + // Handle the remaining full 16-element blocks of the residual (compile-time resolved). if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // The main loop handles 64 elements per iteration; a possible leftover block of 32 is - // handled after it. The loops may run zero times (dim can be as small as 16). - while (pVect1 + 64 <= pEnd1) { - InnerProductStep(pVect1, pVect2, sum0); - InnerProductStep(pVect1, pVect2, sum1); + if constexpr (residual >= 32) { InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 48) { InnerProductStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 64 16-bit floats. + // In each iteration we calculate 64 elements = 4 chunks of 256 bits (converted to 512). + // The loop may run zero times (dim can be as small as 16). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); InnerProductStep(pVect1, pVect2, sum2); InnerProductStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h index 20995ee79..42f5dc188 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h @@ -16,7 +16,7 @@ static inline void InnerProductStep(float *&pVect1, float *&pVect2, __m512 &sum5 sum512 = _mm512_fmadd_ps(v1, v2, sum512); } -template // 0..15 +template // 0..63 float FP32_InnerProductSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; @@ -29,31 +29,35 @@ float FP32_InnerProductSIMD16_AVX512(const void *pVect1v, const void *pVect2v, s __m512 sum2 = _mm512_setzero_ps(); __m512 sum3 = _mm512_setzero_ps(); - // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so - // this is safe for any dimension. - if constexpr (residual) { - __mmask16 constexpr mask = (1 << residual) - 1; + // Deal with the sub-16 remainder first. AVX-512 masked loads suppress faults on masked-out + // lanes, so this is safe for any dimension. + if constexpr (residual % 16) { + __mmask16 constexpr mask = (1 << (residual % 16)) - 1; __m512 v1 = _mm512_maskz_loadu_ps(mask, pVect1); - pVect1 += residual; + pVect1 += residual % 16; __m512 v2 = _mm512_maskz_loadu_ps(mask, pVect2); - pVect2 += residual; + pVect2 += residual % 16; sum0 = _mm512_mul_ps(v1, v2); } - // We dealt with the residual part. We are left with some multiple of 16 floats. - // The main loop handles 64 floats per iteration; leftover blocks of 16/32/48 floats are - // handled after it. The loops may run zero times (dim can be as small as 8). - while (pVect1 + 64 <= pEnd1) { - InnerProductStep(pVect1, pVect2, sum0); + // Handle the remaining full 16-float blocks of the residual (compile-time resolved). + if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 32) { InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 48) { InnerProductStep(pVect1, pVect2, sum3); } - if (pVect1 + 32 <= pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 64 floats. + // In each iteration we calculate 64 floats = 4 chunks of 512 bits. The loop may run zero + // times (dim can be as small as 8). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); InnerProductStep(pVect1, pVect2, sum1); InnerProductStep(pVect1, pVect2, sum2); - } - if (pVect1 < pEnd1) { InnerProductStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h index 8c783b753..d9a4f4e19 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h @@ -16,7 +16,7 @@ static inline void InnerProductStep(double *&pVect1, double *&pVect2, __m512d &s sum512 = _mm512_fmadd_pd(v1, v2, sum512); } -template // 0..7 +template // 0..31 double FP64_InnerProductSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; @@ -29,31 +29,35 @@ double FP64_InnerProductSIMD8_AVX512(const void *pVect1v, const void *pVect2v, s __m512d sum2 = _mm512_setzero_pd(); __m512d sum3 = _mm512_setzero_pd(); - // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so - // this is safe for any dimension. - if constexpr (residual) { - __mmask8 constexpr mask = (1 << residual) - 1; + // Deal with the sub-8 remainder first. AVX-512 masked loads suppress faults on masked-out + // lanes, so this is safe for any dimension. + if constexpr (residual % 8) { + __mmask8 constexpr mask = (1 << (residual % 8)) - 1; __m512d v1 = _mm512_maskz_loadu_pd(mask, pVect1); - pVect1 += residual; + pVect1 += residual % 8; __m512d v2 = _mm512_maskz_loadu_pd(mask, pVect2); - pVect2 += residual; + pVect2 += residual % 8; sum0 = _mm512_mul_pd(v1, v2); } - // We dealt with the residual part. We are left with some multiple of 8 doubles. - // The main loop handles 32 doubles per iteration; leftover blocks of 8/16/24 doubles are - // handled after it. The loops may run zero times (dim can be as small as 4). - while (pVect1 + 32 <= pEnd1) { - InnerProductStep(pVect1, pVect2, sum0); + // Handle the remaining full 8-double blocks of the residual (compile-time resolved). + if constexpr (residual >= 8) { InnerProductStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { InnerProductStep(pVect1, pVect2, sum3); } - if (pVect1 + 16 <= pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 32 doubles. + // In each iteration we calculate 32 doubles = 4 chunks of 512 bits. The loop may run zero + // times (dim can be as small as 4). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); InnerProductStep(pVect1, pVect2, sum1); InnerProductStep(pVect1, pVect2, sum2); - } - if (pVect1 < pEnd1) { InnerProductStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/IP/IP_AVX_FP32.h b/src/VecSim/spaces/IP/IP_AVX_FP32.h index c0dec6057..c07cef722 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP32.h @@ -17,7 +17,7 @@ static inline void InnerProductStep(float *&pVect1, float *&pVect2, __m256 &sum2 sum256 = _mm256_add_ps(sum256, _mm256_mul_ps(v1, v2)); } -template // 0..15 +template // 0..31 float FP32_InnerProductSIMD16_AVX(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; @@ -42,21 +42,23 @@ float FP32_InnerProductSIMD16_AVX(const void *pVect1v, const void *pVect2v, size sum0 = _mm256_mul_ps(v1, v2); } - // If the reminder is >=8, have another step of 8 floats + // Handle the remaining full 8-float blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { InnerProductStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 16 floats. - // The main loop handles 32 floats per iteration; a possible leftover block of 16 floats is - // handled after it. The loops may run zero times (dim can be as small as 8). - while (pVect1 + 32 <= pEnd1) { - InnerProductStep(pVect1, pVect2, sum0); - InnerProductStep(pVect1, pVect2, sum1); + if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { InnerProductStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 32 floats. + // In each iteration we calculate 32 floats = 4 chunks of 256 bits. The loop may run zero + // times (dim can be as small as 8). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); InnerProductStep(pVect1, pVect2, sum2); InnerProductStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/IP/IP_AVX_FP64.h b/src/VecSim/spaces/IP/IP_AVX_FP64.h index 6eb74ed3e..9647c8eba 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP64.h @@ -17,7 +17,7 @@ static inline void InnerProductStep(double *&pVect1, double *&pVect2, __m256d &s sum256 = _mm256_add_pd(sum256, _mm256_mul_pd(v1, v2)); } -template // 0..7 +template // 0..15 double FP64_InnerProductSIMD8_AVX(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; @@ -43,21 +43,23 @@ double FP64_InnerProductSIMD8_AVX(const void *pVect1v, const void *pVect2v, size sum0 = _mm256_mul_pd(v1, v2); } - // If the reminder is >=4, have another step of 4 doubles + // Handle the remaining full 4-double blocks of the residual (compile-time resolved). if constexpr (residual >= 4) { InnerProductStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 8 doubles. - // The main loop handles 16 doubles per iteration; a possible leftover block of 8 doubles is - // handled after it. The loops may run zero times (dim can be as small as 4). - while (pVect1 + 16 <= pEnd1) { - InnerProductStep(pVect1, pVect2, sum0); - InnerProductStep(pVect1, pVect2, sum1); + if constexpr (residual >= 8) { InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 12) { InnerProductStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 16 doubles. + // In each iteration we calculate 16 doubles = 4 chunks of 256 bits. The loop may run zero + // times (dim can be as small as 4). + while (pVect1 < pEnd1) { + InnerProductStep(pVect1, pVect2, sum0); + InnerProductStep(pVect1, pVect2, sum1); InnerProductStep(pVect1, pVect2, sum2); InnerProductStep(pVect1, pVect2, sum3); } 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 46eb4cc6e..6e7a3ca59 100644 --- a/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h @@ -28,7 +28,7 @@ using sq8 = vecsim_types::sq8; */ // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..15 +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); diff --git a/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h index cc1fa4272..96eb0d02c 100644 --- a/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h @@ -28,7 +28,7 @@ using sq8 = vecsim_types::sq8; */ // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..15 +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); 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 57db23fb9..44fa9b706 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 @@ -27,7 +27,7 @@ using sq8 = vecsim_types::sq8; */ // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..15 +template // 0..63 float SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, size_t dimension) { // Get the raw inner product using the common SIMD implementation diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h index 91e0fa527..93f2da6ca 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h @@ -25,7 +25,7 @@ static void L2SqrStep(float16 *&pVect1, float16 *&pVect2, __m512 &sum) { pVect2 += 16; } -template // 0..31 +template // 0..63 float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { auto *pVect1 = (float16 *)pVect1v; auto *pVect2 = (float16 *)pVect2v; @@ -54,20 +54,23 @@ float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t d pVect1 += residual % 16; pVect2 += residual % 16; } + // Handle the remaining full 16-element blocks of the residual (compile-time resolved). if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. - // The main loop handles 64 elements per iteration; a possible leftover block of 32 is - // handled after it. The loops may run zero times (dim can be as small as 16). - while (pVect1 + 64 <= pEnd1) { - L2SqrStep(pVect1, pVect2, sum0); - L2SqrStep(pVect1, pVect2, sum1); + if constexpr (residual >= 32) { L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 48) { L2SqrStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 64 16-bit floats. + // In each iteration we calculate 64 elements = 4 chunks of 256 bits (converted to 512). + // The loop may run zero times (dim can be as small as 16). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); L2SqrStep(pVect1, pVect2, sum2); L2SqrStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h index 9fd276359..79aabfaf2 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h @@ -18,7 +18,7 @@ static inline void L2SqrStep(float *&pVect1, float *&pVect2, __m512 &sum) { sum = _mm512_fmadd_ps(diff, diff, sum); } -template // 0..15 +template // 0..63 float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; @@ -31,32 +31,36 @@ float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t d __m512 sum2 = _mm512_setzero_ps(); __m512 sum3 = _mm512_setzero_ps(); - // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so - // this is safe for any dimension. - if constexpr (residual) { - __mmask16 constexpr mask = (1 << residual) - 1; + // Deal with the sub-16 remainder first. AVX-512 masked loads suppress faults on masked-out + // lanes, so this is safe for any dimension. + if constexpr (residual % 16) { + __mmask16 constexpr mask = (1 << (residual % 16)) - 1; __m512 v1 = _mm512_maskz_loadu_ps(mask, pVect1); - pVect1 += residual; + pVect1 += residual % 16; __m512 v2 = _mm512_maskz_loadu_ps(mask, pVect2); - pVect2 += residual; + pVect2 += residual % 16; __m512 diff = _mm512_sub_ps(v1, v2); sum0 = _mm512_mul_ps(diff, diff); } - // We dealt with the residual part. We are left with some multiple of 16 floats. - // The main loop handles 64 floats per iteration; leftover blocks of 16/32/48 floats are - // handled after it. The loops may run zero times (dim can be as small as 8). - while (pVect1 + 64 <= pEnd1) { - L2SqrStep(pVect1, pVect2, sum0); + // Handle the remaining full 16-float blocks of the residual (compile-time resolved). + if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 32) { L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 48) { L2SqrStep(pVect1, pVect2, sum3); } - if (pVect1 + 32 <= pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 64 floats. + // In each iteration we calculate 64 floats = 4 chunks of 512 bits. The loop may run zero + // times (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); L2SqrStep(pVect1, pVect2, sum1); L2SqrStep(pVect1, pVect2, sum2); - } - if (pVect1 < pEnd1) { L2SqrStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h index 5d3d41633..d1715c012 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h @@ -18,7 +18,7 @@ static inline void L2SqrStep(double *&pVect1, double *&pVect2, __m512d &sum) { sum = _mm512_fmadd_pd(diff, diff, sum); } -template // 0..7 +template // 0..31 double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; @@ -31,32 +31,36 @@ double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t d __m512d sum2 = _mm512_setzero_pd(); __m512d sum3 = _mm512_setzero_pd(); - // Deal with remainder first. AVX-512 masked loads suppress faults on masked-out lanes, so - // this is safe for any dimension. - if constexpr (residual) { - __mmask8 constexpr mask = (1 << residual) - 1; + // Deal with the sub-8 remainder first. AVX-512 masked loads suppress faults on masked-out + // lanes, so this is safe for any dimension. + if constexpr (residual % 8) { + __mmask8 constexpr mask = (1 << (residual % 8)) - 1; __m512d v1 = _mm512_maskz_loadu_pd(mask, pVect1); - pVect1 += residual; + pVect1 += residual % 8; __m512d v2 = _mm512_maskz_loadu_pd(mask, pVect2); - pVect2 += residual; + pVect2 += residual % 8; __m512d diff = _mm512_sub_pd(v1, v2); sum0 = _mm512_mul_pd(diff, diff); } - // We dealt with the residual part. We are left with some multiple of 8 doubles. - // The main loop handles 32 doubles per iteration; leftover blocks of 8/16/24 doubles are - // handled after it. The loops may run zero times (dim can be as small as 4). - while (pVect1 + 32 <= pEnd1) { - L2SqrStep(pVect1, pVect2, sum0); + // Handle the remaining full 8-double blocks of the residual (compile-time resolved). + if constexpr (residual >= 8) { L2SqrStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { L2SqrStep(pVect1, pVect2, sum3); } - if (pVect1 + 16 <= pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 32 doubles. + // In each iteration we calculate 32 doubles = 4 chunks of 512 bits. The loop may run zero + // times (dim can be as small as 4). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); L2SqrStep(pVect1, pVect2, sum1); L2SqrStep(pVect1, pVect2, sum2); - } - if (pVect1 < pEnd1) { L2SqrStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/L2/L2_AVX_FP32.h b/src/VecSim/spaces/L2/L2_AVX_FP32.h index b4d4d8ab7..de7d1cc28 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP32.h @@ -19,7 +19,7 @@ static inline void L2SqrStep(float *&pVect1, float *&pVect2, __m256 &sum) { sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); } -template // 0..15 +template // 0..31 float FP32_L2SqrSIMD16_AVX(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; @@ -45,21 +45,23 @@ float FP32_L2SqrSIMD16_AVX(const void *pVect1v, const void *pVect2v, size_t dime sum0 = _mm256_mul_ps(diff, diff); } - // If the reminder is >=8, have another step of 8 floats + // Handle the remaining full 8-float blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { L2SqrStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 16 floats. - // The main loop handles 32 floats per iteration; a possible leftover block of 16 floats is - // handled after it. The loops may run zero times (dim can be as small as 8). - while (pVect1 + 32 <= pEnd1) { - L2SqrStep(pVect1, pVect2, sum0); - L2SqrStep(pVect1, pVect2, sum1); + if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 24) { L2SqrStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 32 floats. + // In each iteration we calculate 32 floats = 4 chunks of 256 bits. The loop may run zero + // times (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); L2SqrStep(pVect1, pVect2, sum2); L2SqrStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/L2/L2_AVX_FP64.h b/src/VecSim/spaces/L2/L2_AVX_FP64.h index e2509afcd..0d8da9646 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP64.h @@ -19,7 +19,7 @@ static inline void L2SqrStep(double *&pVect1, double *&pVect2, __m256d &sum) { sum = _mm256_add_pd(sum, _mm256_mul_pd(diff, diff)); } -template // 0..7 +template // 0..15 double FP64_L2SqrSIMD8_AVX(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; @@ -46,21 +46,23 @@ double FP64_L2SqrSIMD8_AVX(const void *pVect1v, const void *pVect2v, size_t dime sum0 = _mm256_mul_pd(diff, diff); } - // If the reminder is >=4, have another step of 4 doubles + // Handle the remaining full 4-double blocks of the residual (compile-time resolved). if constexpr (residual >= 4) { L2SqrStep(pVect1, pVect2, sum1); } - - // We dealt with the residual part. We are left with some multiple of 8 doubles. - // The main loop handles 16 doubles per iteration; a possible leftover block of 8 doubles is - // handled after it. The loops may run zero times (dim can be as small as 4). - while (pVect1 + 16 <= pEnd1) { - L2SqrStep(pVect1, pVect2, sum0); - L2SqrStep(pVect1, pVect2, sum1); + if constexpr (residual >= 8) { L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 12) { L2SqrStep(pVect1, pVect2, sum3); } - if (pVect1 < pEnd1) { + + // We dealt with the residual part. We are left with some multiple of 16 doubles. + // In each iteration we calculate 16 doubles = 4 chunks of 256 bits. The loop may run zero + // times (dim can be as small as 4). + while (pVect1 < pEnd1) { + L2SqrStep(pVect1, pVect2, sum0); + L2SqrStep(pVect1, pVect2, sum1); L2SqrStep(pVect1, pVect2, sum2); L2SqrStep(pVect1, pVect2, sum3); } diff --git a/src/VecSim/spaces/functions/AVX.cpp b/src/VecSim/spaces/functions/AVX.cpp index 4b707a5b5..ab6e2b08a 100644 --- a/src/VecSim/spaces/functions/AVX.cpp +++ b/src/VecSim/spaces/functions/AVX.cpp @@ -20,25 +20,25 @@ namespace spaces { dist_func_t Choose_FP32_IP_implementation_AVX(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP32_InnerProductSIMD16_AVX); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP32_InnerProductSIMD16_AVX); return ret_dist_func; } dist_func_t Choose_FP64_IP_implementation_AVX(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 8, FP64_InnerProductSIMD8_AVX); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP64_InnerProductSIMD8_AVX); return ret_dist_func; } dist_func_t Choose_FP32_L2_implementation_AVX(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP32_L2SqrSIMD16_AVX); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP32_L2SqrSIMD16_AVX); return ret_dist_func; } dist_func_t Choose_FP64_L2_implementation_AVX(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 8, FP64_L2SqrSIMD8_AVX); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP64_L2SqrSIMD8_AVX); return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/AVX2.cpp b/src/VecSim/spaces/functions/AVX2.cpp index 322ed0aec..1a9c41b6f 100644 --- a/src/VecSim/spaces/functions/AVX2.cpp +++ b/src/VecSim/spaces/functions/AVX2.cpp @@ -31,19 +31,19 @@ dist_func_t Choose_BF16_L2_implementation_AVX2(size_t dim) { dist_func_t Choose_SQ8_FP32_IP_implementation_AVX2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_InnerProductSIMD16_AVX2); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_InnerProductSIMD16_AVX2); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_Cosine_implementation_AVX2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_CosineSIMD16_AVX2); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_CosineSIMD16_AVX2); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_L2_implementation_AVX2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_L2SqrSIMD16_AVX2); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_L2SqrSIMD16_AVX2); return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/AVX2_FMA.cpp b/src/VecSim/spaces/functions/AVX2_FMA.cpp index c859128b2..0eaacd070 100644 --- a/src/VecSim/spaces/functions/AVX2_FMA.cpp +++ b/src/VecSim/spaces/functions/AVX2_FMA.cpp @@ -16,18 +16,18 @@ namespace spaces { // FMA optimized implementations dist_func_t Choose_SQ8_FP32_IP_implementation_AVX2_FMA(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_InnerProductSIMD16_AVX2_FMA); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_InnerProductSIMD16_AVX2_FMA); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_Cosine_implementation_AVX2_FMA(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_CosineSIMD16_AVX2_FMA); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_CosineSIMD16_AVX2_FMA); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_L2_implementation_AVX2_FMA(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_L2SqrSIMD16_AVX2_FMA); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_L2SqrSIMD16_AVX2_FMA); return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/AVX512F.cpp b/src/VecSim/spaces/functions/AVX512F.cpp index feb261fb4..4c166d5ca 100644 --- a/src/VecSim/spaces/functions/AVX512F.cpp +++ b/src/VecSim/spaces/functions/AVX512F.cpp @@ -24,37 +24,37 @@ namespace spaces { dist_func_t Choose_FP32_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP32_InnerProductSIMD16_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP32_InnerProductSIMD16_AVX512); return ret_dist_func; } dist_func_t Choose_FP64_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 8, FP64_InnerProductSIMD8_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP64_InnerProductSIMD8_AVX512); return ret_dist_func; } dist_func_t Choose_FP32_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP32_L2SqrSIMD16_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP32_L2SqrSIMD16_AVX512); return ret_dist_func; } dist_func_t Choose_FP64_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 8, FP64_L2SqrSIMD8_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP64_L2SqrSIMD8_AVX512); return ret_dist_func; } dist_func_t Choose_FP16_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP16_InnerProductSIMD32_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP16_InnerProductSIMD32_AVX512); return ret_dist_func; } dist_func_t Choose_FP16_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP16_L2SqrSIMD32_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP16_L2SqrSIMD32_AVX512); return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp index 3b8813b89..2198b000a 100644 --- a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp +++ b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp @@ -62,17 +62,17 @@ dist_func_t Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t dist_func_t Choose_SQ8_FP32_IP_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } // SQ8-to-SQ8 distance functions (both vectors are uint8 quantized with precomputed sum) From fd959213af3967835c936cd31bffe9600263dda6 Mon Sep 17 00:00:00 2001 From: GuyAv46 Date: Sun, 5 Jul 2026 14:43:47 +0300 Subject: [PATCH 3/3] [BENCH-ONLY] AVX-512 kernels: 2 accumulators instead of 4 Comparison branch off perf-x86-multi-accumulators-small-dims. Reverts the AVX-512 float/SQ8 kernels (FP32, FP64, FP16, SQ8_FP32 VNNI) from 4 accumulators to 2, narrowing each dispatch modulus to 2x the SIMD chunk accordingly: FP32 / FP16 / SQ8_FP32 64 -> 32 FP64 32 -> 16 Everything else (residual-as-constexpr structure, small-dim thresholds) is unchanged. Purpose: measure 2-acc vs 4-acc on real AVX-512 hardware via CI, since this can't be tested locally (no AVX-512 CPU). Not intended to merge. Co-Authored-By: Claude Fable 5 --- .../IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h | 26 ++++++------------- src/VecSim/spaces/IP/IP_AVX512F_FP16.h | 22 +++++----------- src/VecSim/spaces/IP/IP_AVX512F_FP32.h | 22 +++++----------- src/VecSim/spaces/IP/IP_AVX512F_FP64.h | 22 +++++----------- .../L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h | 2 +- src/VecSim/spaces/L2/L2_AVX512F_FP16.h | 22 +++++----------- src/VecSim/spaces/L2/L2_AVX512F_FP32.h | 22 +++++----------- src/VecSim/spaces/L2/L2_AVX512F_FP64.h | 22 +++++----------- src/VecSim/spaces/functions/AVX512F.cpp | 12 ++++----- .../spaces/functions/AVX512F_BW_VL_VNNI.cpp | 6 ++--- 10 files changed, 54 insertions(+), 124 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h index da7527e7c..c2048d01f 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h @@ -45,18 +45,16 @@ static inline void SQ8_FP32_InnerProductStep(const uint8_t *&pVec1, const float // Common implementation for both inner product and cosine similarity // pVec1v = SQ8 storage, pVec2v = FP32 query -template // 0..63 +template // 0..31 float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, size_t dimension) { const uint8_t *pVec1 = static_cast(pVec1v); // SQ8 storage const float *pVec2 = static_cast(pVec2v); // FP32 query const uint8_t *pEnd1 = pVec1 + dimension; - // Initialize sum accumulators for Σ(q_i * y_i). Four accumulators break the FMA dependency + // Initialize sum accumulators for Σ(q_i * y_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(); - __m512 sum2 = _mm512_setzero_ps(); - __m512 sum3 = _mm512_setzero_ps(); // Handle the sub-16 residual elements first if constexpr (residual % 16) { @@ -79,29 +77,21 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si pVec2 += residual % 16; } - // Handle the remaining full 16-element blocks of the residual (compile-time resolved). + // Handle the remaining full 16-element block of the residual (compile-time resolved). if constexpr (residual >= 16) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); } - if constexpr (residual >= 32) { - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); - } - if constexpr (residual >= 48) { - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 64 elements. - // In each iteration we calculate 64 elements = 4 chunks of 16. The loop may run zero times + // 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) { SQ8_FP32_InnerProductStep(pVec1, pVec2, sum0); SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1); - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum2); - SQ8_FP32_InnerProductStep(pVec1, pVec2, sum3); } // Reduce to get Σ(q_i * y_i) - __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + __m512 sum = _mm512_add_ps(sum0, sum1); float quantized_dot = _mm512_reduce_add_ps(sum); // Get quantization parameters from stored vector (after quantized data) @@ -119,14 +109,14 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si return min_val * y_sum + delta * quantized_dot; } -template // 0..63 +template // 0..31 float SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI(const void *pVec1v, const void *pVec2v, size_t dimension) { // The inner product similarity is 1 - ip return 1.0f - SQ8_FP32_InnerProductImp_AVX512(pVec1v, pVec2v, dimension); } -template // 0..63 +template // 0..31 float SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI(const void *pVec1v, const void *pVec2v, size_t dimension) { // Cosine distance = 1 - IP (vectors are pre-normalized) diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h index 37d6c8f96..440f7be49 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP16.h @@ -24,18 +24,16 @@ static void InnerProductStep(float16 *&pVect1, float16 *&pVect2, __m512 &sum) { pVect2 += 16; } -template // 0..63 +template // 0..31 float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { auto *pVect1 = (float16 *)pVect1v; auto *pVect2 = (float16 *)pVect2v; const float16 *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // Two accumulators break the FMA dependency chain, letting more FMAs be in flight at once. auto sum0 = _mm512_setzero_ps(); auto sum1 = _mm512_setzero_ps(); - auto sum2 = _mm512_setzero_ps(); - auto sum3 = _mm512_setzero_ps(); if constexpr (residual % 16) { // Deal with remainder first. The full-width load of 16 16-bit floats is safe because @@ -52,27 +50,19 @@ float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, s pVect1 += residual % 16; pVect2 += residual % 16; } - // Handle the remaining full 16-element blocks of the residual (compile-time resolved). + // Handle the remaining full 16-element block of the residual (compile-time resolved). if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 32) { - InnerProductStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 48) { - InnerProductStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 64 16-bit floats. - // In each iteration we calculate 64 elements = 4 chunks of 256 bits (converted to 512). + // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. + // In each iteration we calculate 32 elements = 2 chunks of 256 bits (converted to 512). // The loop may run zero times (dim can be as small as 16). while (pVect1 < pEnd1) { InnerProductStep(pVect1, pVect2, sum0); InnerProductStep(pVect1, pVect2, sum1); - InnerProductStep(pVect1, pVect2, sum2); - InnerProductStep(pVect1, pVect2, sum3); } - auto sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + auto sum = _mm512_add_ps(sum0, sum1); return 1.0f - _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h index 42f5dc188..efb8f5cf2 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h @@ -16,18 +16,16 @@ static inline void InnerProductStep(float *&pVect1, float *&pVect2, __m512 &sum5 sum512 = _mm512_fmadd_ps(v1, v2, sum512); } -template // 0..63 +template // 0..31 float FP32_InnerProductSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; const float *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // 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(); - __m512 sum2 = _mm512_setzero_ps(); - __m512 sum3 = _mm512_setzero_ps(); // Deal with the sub-16 remainder first. AVX-512 masked loads suppress faults on masked-out // lanes, so this is safe for any dimension. @@ -40,27 +38,19 @@ float FP32_InnerProductSIMD16_AVX512(const void *pVect1v, const void *pVect2v, s sum0 = _mm512_mul_ps(v1, v2); } - // Handle the remaining full 16-float blocks of the residual (compile-time resolved). + // Handle the remaining full 16-float block of the residual (compile-time resolved). if constexpr (residual >= 16) { InnerProductStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 32) { - InnerProductStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 48) { - InnerProductStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 64 floats. - // In each iteration we calculate 64 floats = 4 chunks of 512 bits. The loop may run zero + // We dealt with the residual part. We are left with some multiple of 32 floats. + // In each iteration we calculate 32 floats = 2 chunks of 512 bits. The loop may run zero // times (dim can be as small as 8). while (pVect1 < pEnd1) { InnerProductStep(pVect1, pVect2, sum0); InnerProductStep(pVect1, pVect2, sum1); - InnerProductStep(pVect1, pVect2, sum2); - InnerProductStep(pVect1, pVect2, sum3); } - __m512 sum512 = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + __m512 sum512 = _mm512_add_ps(sum0, sum1); return 1.0f - _mm512_reduce_add_ps(sum512); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h index d9a4f4e19..5558957f7 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h @@ -16,18 +16,16 @@ static inline void InnerProductStep(double *&pVect1, double *&pVect2, __m512d &s sum512 = _mm512_fmadd_pd(v1, v2, sum512); } -template // 0..31 +template // 0..15 double FP64_InnerProductSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; const double *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // Two accumulators break the FMA dependency chain, letting more FMAs be in flight at once. __m512d sum0 = _mm512_setzero_pd(); __m512d sum1 = _mm512_setzero_pd(); - __m512d sum2 = _mm512_setzero_pd(); - __m512d sum3 = _mm512_setzero_pd(); // Deal with the sub-8 remainder first. AVX-512 masked loads suppress faults on masked-out // lanes, so this is safe for any dimension. @@ -40,27 +38,19 @@ double FP64_InnerProductSIMD8_AVX512(const void *pVect1v, const void *pVect2v, s sum0 = _mm512_mul_pd(v1, v2); } - // Handle the remaining full 8-double blocks of the residual (compile-time resolved). + // Handle the remaining full 8-double block of the residual (compile-time resolved). if constexpr (residual >= 8) { InnerProductStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 16) { - InnerProductStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 24) { - InnerProductStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 32 doubles. - // In each iteration we calculate 32 doubles = 4 chunks of 512 bits. The loop may run zero + // We dealt with the residual part. We are left with some multiple of 16 doubles. + // In each iteration we calculate 16 doubles = 2 chunks of 512 bits. The loop may run zero // times (dim can be as small as 4). while (pVect1 < pEnd1) { InnerProductStep(pVect1, pVect2, sum0); InnerProductStep(pVect1, pVect2, sum1); - InnerProductStep(pVect1, pVect2, sum2); - InnerProductStep(pVect1, pVect2, sum3); } - __m512d sum512 = _mm512_add_pd(_mm512_add_pd(sum0, sum1), _mm512_add_pd(sum2, sum3)); + __m512d sum512 = _mm512_add_pd(sum0, sum1); return 1.0 - _mm512_reduce_add_pd(sum512); } 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 44fa9b706..721789c8e 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 @@ -27,7 +27,7 @@ using sq8 = vecsim_types::sq8; */ // pVect1v = SQ8 storage, pVect2v = FP32 query -template // 0..63 +template // 0..31 float SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, size_t dimension) { // Get the raw inner product using the common SIMD implementation diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h index 93f2da6ca..728b8ebaa 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h @@ -25,18 +25,16 @@ static void L2SqrStep(float16 *&pVect1, float16 *&pVect2, __m512 &sum) { pVect2 += 16; } -template // 0..63 +template // 0..31 float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { auto *pVect1 = (float16 *)pVect1v; auto *pVect2 = (float16 *)pVect2v; const float16 *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // Two accumulators break the FMA dependency chain, letting more FMAs be in flight at once. auto sum0 = _mm512_setzero_ps(); auto sum1 = _mm512_setzero_ps(); - auto sum2 = _mm512_setzero_ps(); - auto sum3 = _mm512_setzero_ps(); if constexpr (residual % 16) { // Deal with remainder first. The full-width load of 16 16-bit floats is safe because @@ -54,27 +52,19 @@ float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t d pVect1 += residual % 16; pVect2 += residual % 16; } - // Handle the remaining full 16-element blocks of the residual (compile-time resolved). + // Handle the remaining full 16-element block of the residual (compile-time resolved). if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 32) { - L2SqrStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 48) { - L2SqrStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 64 16-bit floats. - // In each iteration we calculate 64 elements = 4 chunks of 256 bits (converted to 512). + // We dealt with the residual part. We are left with some multiple of 32 16-bit floats. + // In each iteration we calculate 32 elements = 2 chunks of 256 bits (converted to 512). // The loop may run zero times (dim can be as small as 16). while (pVect1 < pEnd1) { L2SqrStep(pVect1, pVect2, sum0); L2SqrStep(pVect1, pVect2, sum1); - L2SqrStep(pVect1, pVect2, sum2); - L2SqrStep(pVect1, pVect2, sum3); } - auto sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + auto sum = _mm512_add_ps(sum0, sum1); return _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h index 79aabfaf2..19ef0b74e 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h @@ -18,18 +18,16 @@ static inline void L2SqrStep(float *&pVect1, float *&pVect2, __m512 &sum) { sum = _mm512_fmadd_ps(diff, diff, sum); } -template // 0..63 +template // 0..31 float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { float *pVect1 = (float *)pVect1v; float *pVect2 = (float *)pVect2v; const float *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // 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(); - __m512 sum2 = _mm512_setzero_ps(); - __m512 sum3 = _mm512_setzero_ps(); // Deal with the sub-16 remainder first. AVX-512 masked loads suppress faults on masked-out // lanes, so this is safe for any dimension. @@ -43,27 +41,19 @@ float FP32_L2SqrSIMD16_AVX512(const void *pVect1v, const void *pVect2v, size_t d sum0 = _mm512_mul_ps(diff, diff); } - // Handle the remaining full 16-float blocks of the residual (compile-time resolved). + // Handle the remaining full 16-float block of the residual (compile-time resolved). if constexpr (residual >= 16) { L2SqrStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 32) { - L2SqrStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 48) { - L2SqrStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 64 floats. - // In each iteration we calculate 64 floats = 4 chunks of 512 bits. The loop may run zero + // We dealt with the residual part. We are left with some multiple of 32 floats. + // In each iteration we calculate 32 floats = 2 chunks of 512 bits. The loop may run zero // times (dim can be as small as 8). while (pVect1 < pEnd1) { L2SqrStep(pVect1, pVect2, sum0); L2SqrStep(pVect1, pVect2, sum1); - L2SqrStep(pVect1, pVect2, sum2); - L2SqrStep(pVect1, pVect2, sum3); } - __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + __m512 sum = _mm512_add_ps(sum0, sum1); return _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h index d1715c012..6c68da19a 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h @@ -18,18 +18,16 @@ static inline void L2SqrStep(double *&pVect1, double *&pVect2, __m512d &sum) { sum = _mm512_fmadd_pd(diff, diff, sum); } -template // 0..31 +template // 0..15 double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t dimension) { double *pVect1 = (double *)pVect1v; double *pVect2 = (double *)pVect2v; const double *pEnd1 = pVect1 + dimension; - // Four accumulators break the FMA dependency chain, letting more FMAs be in flight at once. + // Two accumulators break the FMA dependency chain, letting more FMAs be in flight at once. __m512d sum0 = _mm512_setzero_pd(); __m512d sum1 = _mm512_setzero_pd(); - __m512d sum2 = _mm512_setzero_pd(); - __m512d sum3 = _mm512_setzero_pd(); // Deal with the sub-8 remainder first. AVX-512 masked loads suppress faults on masked-out // lanes, so this is safe for any dimension. @@ -43,27 +41,19 @@ double FP64_L2SqrSIMD8_AVX512(const void *pVect1v, const void *pVect2v, size_t d sum0 = _mm512_mul_pd(diff, diff); } - // Handle the remaining full 8-double blocks of the residual (compile-time resolved). + // Handle the remaining full 8-double block of the residual (compile-time resolved). if constexpr (residual >= 8) { L2SqrStep(pVect1, pVect2, sum1); } - if constexpr (residual >= 16) { - L2SqrStep(pVect1, pVect2, sum2); - } - if constexpr (residual >= 24) { - L2SqrStep(pVect1, pVect2, sum3); - } - // We dealt with the residual part. We are left with some multiple of 32 doubles. - // In each iteration we calculate 32 doubles = 4 chunks of 512 bits. The loop may run zero + // We dealt with the residual part. We are left with some multiple of 16 doubles. + // In each iteration we calculate 16 doubles = 2 chunks of 512 bits. The loop may run zero // times (dim can be as small as 4). while (pVect1 < pEnd1) { L2SqrStep(pVect1, pVect2, sum0); L2SqrStep(pVect1, pVect2, sum1); - L2SqrStep(pVect1, pVect2, sum2); - L2SqrStep(pVect1, pVect2, sum3); } - __m512d sum = _mm512_add_pd(_mm512_add_pd(sum0, sum1), _mm512_add_pd(sum2, sum3)); + __m512d sum = _mm512_add_pd(sum0, sum1); return _mm512_reduce_add_pd(sum); } diff --git a/src/VecSim/spaces/functions/AVX512F.cpp b/src/VecSim/spaces/functions/AVX512F.cpp index 4c166d5ca..0b68d505e 100644 --- a/src/VecSim/spaces/functions/AVX512F.cpp +++ b/src/VecSim/spaces/functions/AVX512F.cpp @@ -24,37 +24,37 @@ namespace spaces { dist_func_t Choose_FP32_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP32_InnerProductSIMD16_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP32_InnerProductSIMD16_AVX512); return ret_dist_func; } dist_func_t Choose_FP64_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP64_InnerProductSIMD8_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP64_InnerProductSIMD8_AVX512); return ret_dist_func; } dist_func_t Choose_FP32_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP32_L2SqrSIMD16_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP32_L2SqrSIMD16_AVX512); return ret_dist_func; } dist_func_t Choose_FP64_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP64_L2SqrSIMD8_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP64_L2SqrSIMD8_AVX512); return ret_dist_func; } dist_func_t Choose_FP16_IP_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP16_InnerProductSIMD32_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP16_InnerProductSIMD32_AVX512); return ret_dist_func; } dist_func_t Choose_FP16_L2_implementation_AVX512F(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP16_L2SqrSIMD32_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP16_L2SqrSIMD32_AVX512); return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp index 2198b000a..97da55546 100644 --- a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp +++ b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp @@ -62,17 +62,17 @@ dist_func_t Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t dist_func_t Choose_SQ8_FP32_IP_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_InnerProductSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_CosineSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } dist_func_t Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI); return ret_dist_func; } // SQ8-to-SQ8 distance functions (both vectors are uint8 quantized with precomputed sum)