[BENCH-ONLY] AVX-512: 2 accumulators vs 4 (do not merge)

src/VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h

@@ -46,16 +46,21 @@ static inline void InnerProductStepSQ8_FMA(const uint8_t *&pVect1, const float *
 }
 
 // pVect1v = SQ8 storage, pVect2v = FP32 query
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 0..31
 float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, size_t dimension) {
     const uint8_t *pVect1 = static_cast<const uint8_t *>(pVect1v); // SQ8 storage
     const float *pVect2 = static_cast<const float *>(pVect2v);     // FP32 query
     const uint8_t *pEnd1 = pVect1 + dimension;
 
-    // Initialize sum accumulator for Σ(q_i * y_i)
-    __m256 sum256 = _mm256_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.
+    __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)
+    // 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,33 @@ 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
+    // Handle the remaining full 8-element blocks of the residual (compile-time resolved).
     if constexpr (residual >= 8) {
-        InnerProductStepSQ8_FMA(pVect1, pVect2, sum256);
+        InnerProductStepSQ8_FMA(pVect1, pVect2, sum1);
+    }
+    if constexpr (residual >= 16) {
+        InnerProductStepSQ8_FMA(pVect1, pVect2, sum2);
+    }
+    if constexpr (residual >= 24) {
+        InnerProductStepSQ8_FMA(pVect1, pVect2, sum3);
     }
 
-    // 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 dealt with the residual part. We are left with some multiple of 32 elements.
+    // In each iteration we calculate 32 elements = 4 chunks of 8. The loop may run zero times
+    // (dim can be as small as 8).
+    while (pVect1 < pEnd1) {
+        InnerProductStepSQ8_FMA(pVect1, pVect2, sum0);
+        InnerProductStepSQ8_FMA(pVect1, pVect2, sum1);
+        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<const uint8_t *>(pVect1v);
@@ -102,13 +117,13 @@ float SQ8_FP32_InnerProductImp_FMA(const void *pVect1v, const void *pVect2v, siz
     return min_val * y_sum + delta * quantized_dot;
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 0..31
 float SQ8_FP32_InnerProductSIMD16_AVX2_FMA(const void *pVect1v, const void *pVect2v,
                                            size_t dimension) {
     return 1.0f - SQ8_FP32_InnerProductImp_FMA<residual>(pVect1v, pVect2v, dimension);
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 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<residual>(pVect1v, pVect2v, dimension);

src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h

@@ -46,16 +46,21 @@ static inline void InnerProductStepSQ8_FP32(const uint8_t *&pVect1, const float
 }
 
 // pVect1v = SQ8 storage, pVect2v = FP32 query
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 0..31
 float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) {
     const uint8_t *pVect1 = static_cast<const uint8_t *>(pVect1v); // SQ8 storage
     const float *pVect2 = static_cast<const float *>(pVect2v);     // FP32 query
     const uint8_t *pEnd1 = pVect1 + dimension;
 
-    // Initialize sum accumulator for Σ(q_i * y_i)
-    __m256 sum256 = _mm256_setzero_ps();
+    // 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)
+    // 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,33 @@ 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
+    // Handle the remaining full 8-element blocks of the residual (compile-time resolved).
     if constexpr (residual >= 8) {
-        InnerProductStepSQ8_FP32(pVect1, pVect2, sum256);
+        InnerProductStepSQ8_FP32(pVect1, pVect2, sum1);
+    }
+    if constexpr (residual >= 16) {
+        InnerProductStepSQ8_FP32(pVect1, pVect2, sum2);
+    }
+    if constexpr (residual >= 24) {
+        InnerProductStepSQ8_FP32(pVect1, pVect2, sum3);
     }
 
-    // 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 dealt with the residual part. We are left with some multiple of 32 elements.
+    // In each iteration we calculate 32 elements = 4 chunks of 8. The loop may run zero times
+    // (dim can be as small as 8).
+    while (pVect1 < pEnd1) {
+        InnerProductStepSQ8_FP32(pVect1, pVect2, sum0);
+        InnerProductStepSQ8_FP32(pVect1, pVect2, sum1);
+        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<const uint8_t *>(pVect1v);
@@ -102,12 +117,12 @@ float SQ8_FP32_InnerProductImp_AVX2(const void *pVect1v, const void *pVect2v, si
     return min_val * y_sum + delta * quantized_dot;
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 0..31
 float SQ8_FP32_InnerProductSIMD16_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) {
     return 1.0f - SQ8_FP32_InnerProductImp_AVX2<residual>(pVect1v, pVect2v, dimension);
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 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<residual>(pVect1v, pVect2v, dimension);

src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h

@@ -45,20 +45,24 @@ 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 <unsigned char residual> // 0..15
+template <unsigned char residual> // 0..31
 float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, size_t dimension) {
     const uint8_t *pVec1 = static_cast<const uint8_t *>(pVec1v); // SQ8 storage
     const float *pVec2 = static_cast<const float *>(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). Two accumulators break the FMA dependency
+    // chain, letting more FMAs be in flight at once.
+    __m512 sum0 = _mm512_setzero_ps();
+    __m512 sum1 = _mm512_setzero_ps();
 
-    // Handle 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 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<const __m128i *>(pVec1));
         __m512i v1_512 = _mm512_cvtepu8_epi32(v1_128);
         __m512 v1_f = _mm512_cvtepi32_ps(v1_512);
@@ -67,19 +71,27 @@ 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;
+        pVec1 += residual % 16;
+        pVec2 += residual % 16;
     }
 
-    // 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);
+    // Handle the remaining full 16-element block of the residual (compile-time resolved).
+    if constexpr (residual >= 16) {
+        SQ8_FP32_InnerProductStep(pVec1, pVec2, sum1);
+    }
+
+    // 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);
+    }
 
     // Reduce to get Σ(q_i * y_i)
+    __m512 sum = _mm512_add_ps(sum0, sum1);
     float quantized_dot = _mm512_reduce_add_ps(sum);
 
     // Get quantization parameters from stored vector (after quantized data)
@@ -97,14 +109,14 @@ float SQ8_FP32_InnerProductImp_AVX512(const void *pVec1v, const void *pVec2v, si
     return min_val * y_sum + delta * quantized_dot;
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 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<residual>(pVec1v, pVec2v, dimension);
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 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)

src/VecSim/spaces/IP/IP_AVX512F_FP16.h

@@ -31,11 +31,13 @@ float FP16_InnerProductSIMD32_AVX512(const void *pVect1v, const void *pVect2v, s
 
     const float16 *pEnd1 = pVect1 + dimension;
 
-    auto sum = _mm512_setzero_ps();
+    // 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();
 
     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 +46,23 @@ 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;
     }
+    // Handle the remaining full 16-element block of the residual (compile-time resolved).
     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);
+    // 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);
+    }
 
+    auto sum = _mm512_add_ps(sum0, sum1);
     return 1.0f - _mm512_reduce_add_ps(sum);
 }

src/VecSim/spaces/IP/IP_AVX512F_FP32.h

@@ -16,30 +16,41 @@ static inline void InnerProductStep(float *&pVect1, float *&pVect2, __m512 &sum5
     sum512 = _mm512_fmadd_ps(v1, v2, sum512);
 }
 
-template <unsigned char residual> // 0..15
+template <unsigned char residual> // 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;
 
-    __m512 sum512 = _mm512_setzero_ps();
+    // 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();
 
-    // 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
-    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;
-        sum512 = _mm512_mul_ps(v1, v2);
+        pVect2 += residual % 16;
+        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);
+    // Handle the remaining full 16-float block 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 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);
+    }
 
+    __m512 sum512 = _mm512_add_ps(sum0, sum1);
     return 1.0f - _mm512_reduce_add_ps(sum512);
 }