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..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,16 +46,21 @@ 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 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(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 // 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 dea167eb3..5152e4544 100644 --- a/src/VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h +++ b/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 // 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 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(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 // 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 76a590519..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,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 // 0..15 +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 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(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 // 0..15 +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..15 +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 b09352533..440f7be49 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP16.h +++ b/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); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h b/src/VecSim/spaces/IP/IP_AVX512F_FP32.h index 88421fd39..efb8f5cf2 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP32.h +++ b/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 // 0..15 +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; - __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); } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h index e6eebcc44..5558957f7 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_FP64.h @@ -16,30 +16,41 @@ static inline void InnerProductStep(double *&pVect1, double *&pVect2, __m512d &s sum512 = _mm512_fmadd_pd(v1, v2, sum512); } -template // 0..7 +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; - __m512d sum512 = _mm512_setzero_pd(); + // 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(); - // 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 - 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; - sum512 = _mm512_mul_pd(v1, v2); + pVect2 += residual % 8; + 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); + // Handle the remaining full 8-double block 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 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); + } + __m512d sum512 = _mm512_add_pd(sum0, sum1); 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..c07cef722 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP32.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP32.h @@ -17,37 +17,52 @@ 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; 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 + // Handle the remaining full 8-float blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { - InnerProductStep(pVect1, pVect2, sum256); + 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 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - InnerProductStep(pVect1, pVect2, sum256); - InnerProductStep(pVect1, pVect2, sum256); - } while (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); + } - 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..9647c8eba 100644 --- a/src/VecSim/spaces/IP/IP_AVX_FP64.h +++ b/src/VecSim/spaces/IP/IP_AVX_FP64.h @@ -17,17 +17,22 @@ 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; 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,31 @@ 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 + // Handle the remaining full 4-double blocks of the residual (compile-time resolved). if constexpr (residual >= 4) { - InnerProductStep(pVect1, pVect2, sum256); + InnerProductStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 8) { + InnerProductStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 12) { + 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. - do { - InnerProductStep(pVect1, pVect2, sum256); - InnerProductStep(pVect1, pVect2, sum256); - } while (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); + } + __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_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..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..15 +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 e2a21414f..728b8ebaa 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP16.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP16.h @@ -32,11 +32,13 @@ float FP16_L2SqrSIMD32_AVX512(const void *pVect1v, const void *pVect2v, size_t d 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 @@ -46,20 +48,23 @@ 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; } + // Handle the remaining full 16-element block of the residual (compile-time resolved). 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); + // 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); + } + 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 0100e4264..19ef0b74e 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP32.h @@ -18,31 +18,42 @@ static inline void L2SqrStep(float *&pVect1, float *&pVect2, __m512 &sum) { sum = _mm512_fmadd_ps(diff, diff, sum); } -template // 0..15 +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; - __m512 sum = _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; + pVect2 += residual % 16; __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); + // Handle the remaining full 16-float block 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 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); + } + __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 1a54c7048..6c68da19a 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_FP64.h @@ -18,31 +18,42 @@ static inline void L2SqrStep(double *&pVect1, double *&pVect2, __m512d &sum) { sum = _mm512_fmadd_pd(diff, diff, sum); } -template // 0..7 +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; - __m512d sum = _mm512_setzero_pd(); + // 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(); - // 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 - 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); - 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); + // Handle the remaining full 8-double block 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 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); + } + __m512d sum = _mm512_add_pd(sum0, sum1); 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..de7d1cc28 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP32.h @@ -19,16 +19,22 @@ 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; 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,30 @@ 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 + // Handle the remaining full 8-float blocks of the residual (compile-time resolved). if constexpr (residual >= 8) { - L2SqrStep(pVect1, pVect2, sum); + 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 16 floats. - // In each iteration we calculate 16 floats = 512 bits. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (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); + } + __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..0d8da9646 100644 --- a/src/VecSim/spaces/L2/L2_AVX_FP64.h +++ b/src/VecSim/spaces/L2/L2_AVX_FP64.h @@ -19,16 +19,22 @@ 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; 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,31 @@ 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 + // Handle the remaining full 4-double blocks of the residual (compile-time resolved). if constexpr (residual >= 4) { - L2SqrStep(pVect1, pVect2, sum); + L2SqrStep(pVect1, pVect2, sum1); + } + if constexpr (residual >= 8) { + L2SqrStep(pVect1, pVect2, sum2); + } + if constexpr (residual >= 12) { + 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. - do { - L2SqrStep(pVect1, pVect2, sum); - L2SqrStep(pVect1, pVect2, sum); - } while (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); + } + __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/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..0b68d505e 100644 --- a/src/VecSim/spaces/functions/AVX512F.cpp +++ b/src/VecSim/spaces/functions/AVX512F.cpp @@ -24,25 +24,25 @@ 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, 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, 8, 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, 16, 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, 8, FP64_L2SqrSIMD8_AVX512); + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 16, FP64_L2SqrSIMD8_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..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, 16, 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, 16, 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, 16, 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) 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();