diff --git a/src/VecSim/spaces/IP/IP.cpp b/src/VecSim/spaces/IP/IP.cpp index 2140c2345..83be88044 100644 --- a/src/VecSim/spaces/IP/IP.cpp +++ b/src/VecSim/spaces/IP/IP.cpp @@ -238,11 +238,15 @@ float FP16_InnerProduct(const void *pVect1, const void *pVect2, size_t dimension } // Return type for the inner product functions. -// The type should be able to hold `dimension * MAX_VAL(int_elem_t) * MAX_VAL(int_elem_t)`. -// To support dimension up to 2^16, we need the difference between the type and int_elem_t to be at -// least 2 bytes. We assert that in the implementation. +// The type must hold `dimension * MAX_VAL(int_elem_t) * MAX_VAL(int_elem_t)`. For uint8 that is +// 65025 * dimension, which overflows a 32-bit int from dimension 33,026 -- the alias was previously +// `int` for any 1-byte element, so UINT8_InnerProduct executed signed-overflow UB there. +// +// Signedness follows the element type, which matters for the wrappers below: UINT8_InnerProduct +// converts to float before subtracting from 1, so an unsigned accumulator is fine there, while +// INT8_InnerProduct computes `1 - ip` in integer arithmetic and needs a signed one. template -using ret_t = std::conditional_t; +using ret_t = std::conditional_t, uint64_t, int64_t>; template static inline ret_t @@ -273,7 +277,11 @@ float INT8_Cosine(const void *pVect1v, const void *pVect2v, size_t dimension) { float UINT8_InnerProduct(const void *pVect1v, const void *pVect2v, size_t dimension) { const auto *pVect1 = static_cast(pVect1v); const auto *pVect2 = static_cast(pVect2v); - return 1 - INTEGER_InnerProductImp(pVect1, pVect2, dimension); + // Integer subtract then a single conversion, matching INT8_InnerProduct above. The cast to a + // signed type is required because ret_t is unsigned for uint8: 1 - an unsigned total would + // wrap. The SIMD kernels do the same, so their results stay bit-identical to this one. + const auto ip = static_cast(INTEGER_InnerProductImp(pVect1, pVect2, dimension)); + return static_cast(1 - ip); } float UINT8_Cosine(const void *pVect1v, const void *pVect2v, size_t dimension) { diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_SQ8.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_SQ8.h index ae6f96ea2..61806504d 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_SQ8.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_SQ8.h @@ -40,7 +40,15 @@ template // 0..63 float SQ8_SQ8_InnerProductImp(const void *pVec1v, const void *pVec2v, size_t dimension) { // Compute raw dot product using efficient UINT8 AVX512 VNNI implementation // UINT8_InnerProductImp uses _mm512_dpwssd_epi32 for native integer dot product - int dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); + // uint32_t, matching what the helper returns. This kernel is reachable at any dimension: unlike + // the plain uint8 choosers, the SQ8_SQ8 choosers have no dimension guard, so the previous int + // narrowed and wrapped past 33,025 and the float ones lost exactness past 258. + // + // Note this calls the helper directly rather than through a uint8 chooser, so it does not get + // the chunked accumulation those choosers select past spaces::UINT8_CHUNK_ELEMENTS: the total + // here is still a single 32-bit reduce. SQ8 is capped well below that by its uint32 + // q_sum_squares metadata slot, so the fence belongs with SQ8 index creation (#1007), not here. + const uint32_t dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); // Get dequantization parameters and precomputed values from the end of vectors // Layout: [data (dim)] [min (float)] [delta (float)] [sum (float)] diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_UINT8.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_UINT8.h index bd43bc901..71123af65 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_UINT8.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_UINT8.h @@ -8,6 +8,19 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" + +// uint8 inner product: Imp returns the raw integer total and the wrappers convert it. The chooser +// picks plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h +// carries the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. +// The inner product subtracts in integer and converts once, signed because the total is not. static inline void InnerProductStep(uint8_t *&pVect1, uint8_t *&pVect2, __m512i &sum) { __m512i va = _mm512_loadu_epi8(pVect1); // AVX512BW @@ -30,9 +43,12 @@ static inline void InnerProductStep(uint8_t *&pVect1, uint8_t *&pVect2, __m512i // with the corresponding 32-bit integer in src, and store the packed 32-bit results in dst. } +// always_inline, not merely inline: the chunked wrapper below calls this twice, and without the +// attribute GCC outlines it once it has several callers, which also costs the plain wrapper its +// inlining. Measured: the plain residual-0 wrapper went from 33 instructions to 9 plus a call. template // 0..63 -static inline int UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, - size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -87,19 +103,60 @@ static inline int UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v } while (pVect1 < pEnd1); } - return _mm512_reduce_add_epi32(sum); + // Unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return static_cast(_mm512_reduce_add_epi32(sum)); } template // 0..63 float UINT8_InnerProductSIMD64_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, size_t dimension) { - return 1 - UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + const auto ip = + static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); } template // 0..63 float UINT8_CosineSIMD64_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, size_t dimension) { - float ip = UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + float ip = static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); + const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); + return 1.0f - ip / (norm_v1 * norm_v2); +} + +__attribute__((noinline)) static uint32_t +UINT8_InnerProductFullChunk_AVX512F_BW_VL_VNNI(const uint8_t *pVect1, const uint8_t *pVect2, + size_t dimension) { + return UINT8_InnerProductImp<0>(pVect1, pVect2, dimension); +} + +template // 0..63 +struct UINT8_IPChunkKernel_AVX512F_BW_VL_VNNI { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductFullChunk_AVX512F_BW_VL_VNNI(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_InnerProductSIMD64_AVX512F_BW_VL_VNNI_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const auto ip = static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); +} + +template // 0..63 +float UINT8_CosineSIMD64_AVX512F_BW_VL_VNNI_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const float ip = static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); return 1.0f - ip / (norm_v1 * norm_v2); diff --git a/src/VecSim/spaces/IP/IP_NEON_DOTPROD_SQ8_SQ8.h b/src/VecSim/spaces/IP/IP_NEON_DOTPROD_SQ8_SQ8.h index d7f5b444e..a05f932ee 100644 --- a/src/VecSim/spaces/IP/IP_NEON_DOTPROD_SQ8_SQ8.h +++ b/src/VecSim/spaces/IP/IP_NEON_DOTPROD_SQ8_SQ8.h @@ -41,7 +41,15 @@ float SQ8_SQ8_InnerProductSIMD64_NEON_DOTPROD_IMP(const void *pVec1v, const void size_t dimension) { // Compute raw dot product using efficient UINT8 DOTPROD implementation // UINT8_InnerProductImp uses vdotq_u32 for native uint8 dot product - float dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); + // uint32_t, matching what the helper returns. This kernel is reachable at any dimension: unlike + // the plain uint8 choosers, the SQ8_SQ8 choosers have no dimension guard, so the previous int + // narrowed and wrapped past 33,025 and the float ones lost exactness past 258. + // + // Note this calls the helper directly rather than through a uint8 chooser, so it does not get + // the chunked accumulation those choosers select past spaces::UINT8_CHUNK_ELEMENTS: the total + // here is still a single 32-bit reduce. SQ8 is capped well below that by its uint32 + // q_sum_squares metadata slot, so the fence belongs with SQ8 index creation (#1007), not here. + const uint32_t dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); // Get dequantization parameters and precomputed values from the end of vectors // Layout: [data (dim)] [min (float)] [delta (float)] [sum (float)] diff --git a/src/VecSim/spaces/IP/IP_NEON_DOTPROD_UINT8.h b/src/VecSim/spaces/IP/IP_NEON_DOTPROD_UINT8.h index 3abbd9bba..504fd912c 100644 --- a/src/VecSim/spaces/IP/IP_NEON_DOTPROD_UINT8.h +++ b/src/VecSim/spaces/IP/IP_NEON_DOTPROD_UINT8.h @@ -8,8 +8,22 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 inner product: Imp returns the raw integer total and the wrappers convert it. The chooser +// picks plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h +// carries the chunk-size argument. +// +// Imp is static because IP_NEON_UINT8.h defines the same name with a different body and +// aarch64 gcc 12.3 outlines it, so shared linkage lets a NEON call site execute udot and fault +// where asimddp is absent. always_inline keeps the plain wrapper's codegen unchanged. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. +// The inner product subtracts in integer and converts once, signed because the total is not. + __attribute__((always_inline)) static inline void InnerProductOp(uint8x16_t &v1, uint8x16_t &v2, uint32x4_t &sum) { sum = vdotq_u32(sum, v1, v2); @@ -27,7 +41,8 @@ InnerProductStep(uint8_t *&pVect1, uint8_t *&pVect2, uint32x4_t &sum) { } template // 0..63 -float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -97,20 +112,59 @@ float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dim uint32x4_t total_sum = vaddq_u32(sum0, sum1); - int32_t result = vaddvq_u32(total_sum); - - return static_cast(result); + // ADDV, unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return vaddvq_u32(total_sum); } template // 0..63 float UINT8_InnerProductSIMD16_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, size_t dimension) { - return 1.0f - UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + const auto ip = + static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); } template // 0..63 float UINT8_CosineSIMD_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, size_t dimension) { - float ip = UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + float ip = static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); + const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); + return 1.0f - ip / (norm_v1 * norm_v2); +} + +__attribute__((noinline)) static uint32_t +UINT8_InnerProductFullChunk_NEON_DOTPROD(const uint8_t *pVect1, const uint8_t *pVect2, + size_t dimension) { + return UINT8_InnerProductImp<0>(pVect1, pVect2, dimension); +} + +template // 0..63 +struct UINT8_IPChunkKernel_NEON_DOTPROD { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductFullChunk_NEON_DOTPROD(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_InnerProductSIMD16_NEON_DOTPROD_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const auto ip = static_cast( + spaces::uint8_chunked_total>(pVect1v, pVect2v, + dimension)); + return static_cast(1 - ip); +} + +template // 0..63 +float UINT8_CosineSIMD_NEON_DOTPROD_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + float ip = + static_cast(spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); return 1.0f - ip / (norm_v1 * norm_v2); diff --git a/src/VecSim/spaces/IP/IP_NEON_SQ8_SQ8.h b/src/VecSim/spaces/IP/IP_NEON_SQ8_SQ8.h index 3e931ee0e..215d20e5e 100644 --- a/src/VecSim/spaces/IP/IP_NEON_SQ8_SQ8.h +++ b/src/VecSim/spaces/IP/IP_NEON_SQ8_SQ8.h @@ -41,7 +41,15 @@ float SQ8_SQ8_InnerProductSIMD64_NEON_IMP(const void *pVec1v, const void *pVec2v size_t dimension) { // Compute raw dot product using efficient UINT8 implementation // UINT8_InnerProductImp processes 16 elements at a time using native uint8 instructions - float dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); + // uint32_t, matching what the helper returns. This kernel is reachable at any dimension: unlike + // the plain uint8 choosers, the SQ8_SQ8 choosers have no dimension guard, so the previous int + // narrowed and wrapped past 33,025 and the float ones lost exactness past 258. + // + // Note this calls the helper directly rather than through a uint8 chooser, so it does not get + // the chunked accumulation those choosers select past spaces::UINT8_CHUNK_ELEMENTS: the total + // here is still a single 32-bit reduce. SQ8 is capped well below that by its uint32 + // q_sum_squares metadata slot, so the fence belongs with SQ8 index creation (#1007), not here. + const uint32_t dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); // Get dequantization parameters and precomputed values from the end of pVec1 // Layout: [data (dim)] [min (float)] [delta (float)] [sum (float)] diff --git a/src/VecSim/spaces/IP/IP_NEON_UINT8.h b/src/VecSim/spaces/IP/IP_NEON_UINT8.h index 2d2b3f555..d012e47e1 100644 --- a/src/VecSim/spaces/IP/IP_NEON_UINT8.h +++ b/src/VecSim/spaces/IP/IP_NEON_UINT8.h @@ -8,8 +8,22 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 inner product: Imp returns the raw integer total and the wrappers convert it. The chooser +// picks plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h +// carries the chunk-size argument. +// +// Imp is static because IP_NEON_DOTPROD_UINT8.h defines the same name with a different body and +// aarch64 gcc 12.3 outlines it, so shared linkage lets a NEON call site execute udot and fault +// where asimddp is absent. always_inline keeps the plain wrapper's codegen unchanged. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. +// The inner product subtracts in integer and converts once, signed because the total is not. + __attribute__((always_inline)) static inline void InnerProductOp(uint8x16_t &v1, uint8x16_t &v2, uint32x4_t &sum) { // Multiply and accumulate low 8 elements (first half) @@ -35,7 +49,8 @@ InnerProductStep(uint8_t *&pVect1, uint8_t *&pVect2, uint32x4_t &sum) { } template // 0..63 -float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -105,20 +120,55 @@ float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dim uint32x4_t total_sum = vaddq_u32(sum0, sum1); - // Horizontal sum of the 4 elements in the combined sum register - int32_t result = vaddvq_u32(total_sum); - - return static_cast(result); + // ADDV, unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return vaddvq_u32(total_sum); } template // 0..15 float UINT8_InnerProductSIMD16_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { - return 1.0f - UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + const auto ip = + static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); } template // 0..63 float UINT8_CosineSIMD_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { - float ip = UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + float ip = static_cast(UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); + const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); + return 1.0f - ip / (norm_v1 * norm_v2); +} + +__attribute__((noinline)) static uint32_t +UINT8_InnerProductFullChunk_NEON(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp<0>(pVect1, pVect2, dimension); +} + +template // 0..63 +struct UINT8_IPChunkKernel_NEON { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductFullChunk_NEON(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_InnerProductSIMD16_NEON_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const auto ip = + static_cast(spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); +} + +template // 0..63 +float UINT8_CosineSIMD_NEON_Chunked(const void *pVect1v, const void *pVect2v, size_t dimension) { + float ip = static_cast(spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); return 1.0f - ip / (norm_v1 * norm_v2); diff --git a/src/VecSim/spaces/IP/IP_SVE_SQ8_SQ8.h b/src/VecSim/spaces/IP/IP_SVE_SQ8_SQ8.h index 93ddb76cb..e13a88823 100644 --- a/src/VecSim/spaces/IP/IP_SVE_SQ8_SQ8.h +++ b/src/VecSim/spaces/IP/IP_SVE_SQ8_SQ8.h @@ -40,7 +40,15 @@ template float SQ8_SQ8_InnerProductSIMD_SVE_IMP(const void *pVec1v, const void *pVec2v, size_t dimension) { // Compute raw dot product using efficient UINT8 SVE implementation // UINT8_InnerProductImp uses svdot_u32 for native uint8 dot product - float dot_product = + // uint32_t, matching what the helper returns. This kernel is reachable at any dimension: unlike + // the plain uint8 choosers, the SQ8_SQ8 choosers have no dimension guard, so the previous int + // narrowed and wrapped past 33,025 and the float ones lost exactness past 258. + // + // Note this calls the helper directly rather than through a uint8 chooser, so it does not get + // the chunked accumulation those choosers select past spaces::UINT8_CHUNK_ELEMENTS: the total + // here is still a single 32-bit reduce. SQ8 is capped well below that by its uint32 + // q_sum_squares metadata slot, so the fence belongs with SQ8 index creation (#1007), not here. + const uint32_t dot_product = UINT8_InnerProductImp(pVec1v, pVec2v, dimension); // Get dequantization parameters and precomputed values from the end of vectors diff --git a/src/VecSim/spaces/IP/IP_SVE_UINT8.h b/src/VecSim/spaces/IP/IP_SVE_UINT8.h index f6c6af3b8..7dd963903 100644 --- a/src/VecSim/spaces/IP/IP_SVE_UINT8.h +++ b/src/VecSim/spaces/IP/IP_SVE_UINT8.h @@ -8,8 +8,21 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 inner product: Imp returns the raw integer total and the wrappers convert it. The chooser +// picks plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h +// carries the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk keeps this instantiation's residual shape, clamped to the +// dimension; the vector length is a runtime value so that split is computed rather than folded. +// Later chunks share one out-of-line copy of the kernel. +// The inner product subtracts in integer and converts once, signed because the total is not. + inline void InnerProductStep(const uint8_t *&pVect1, const uint8_t *&pVect2, size_t &offset, svuint32_t &sum, const size_t chunk) { svbool_t pg = svptrue_b8(); @@ -24,7 +37,8 @@ inline void InnerProductStep(const uint8_t *&pVect1, const uint8_t *&pVect2, siz } template -float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dimension) { const uint8_t *pVect1 = reinterpret_cast(pVect1v); const uint8_t *pVect2 = reinterpret_cast(pVect2v); @@ -82,21 +96,57 @@ float UINT8_InnerProductImp(const void *pVect1v, const void *pVect2v, size_t dim sum0 = svadd_u32_x(svptrue_b32(), sum0, sum1); sum2 = svadd_u32_x(svptrue_b32(), sum2, sum3); - // Perform vector addition in parallel and Horizontal sum - int32_t sum_all = svaddv_u32(svptrue_b32(), svadd_u32_x(svptrue_b32(), sum0, sum2)); - - return sum_all; + // Exact for up to spaces::UINT8_CHUNK_ELEMENTS elements; narrowed from svaddv_u32. + return static_cast(svaddv_u32(svptrue_b32(), svadd_u32_x(svptrue_b32(), sum0, sum2))); } template float UINT8_InnerProductSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { - return 1.0f - - UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + const auto ip = static_cast( + UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); } template float UINT8_CosineSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { - float ip = UINT8_InnerProductImp(pVect1v, pVect2v, dimension); + float ip = static_cast( + UINT8_InnerProductImp(pVect1v, pVect2v, dimension)); + const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); + const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); + return 1.0f - ip / (norm_v1 * norm_v2); +} + +__attribute__((noinline)) static uint32_t +UINT8_InnerProductFullChunk_SVE(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp(pVect1, pVect2, dimension); +} + +template +struct UINT8_IPChunkKernel_SVE { + static size_t granule() { return 4 * svcntb(); } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductImp(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_InnerProductFullChunk_SVE(pVect1, pVect2, dimension); + } +}; + +template +float UINT8_InnerProductSIMD_SVE_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const auto ip = static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); + return static_cast(1 - ip); +} + +template +float UINT8_CosineSIMD_SVE_Chunked(const void *pVect1v, const void *pVect2v, size_t dimension) { + float ip = static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); const float norm_v1 = load_unaligned(static_cast(pVect1v) + dimension); const float norm_v2 = load_unaligned(static_cast(pVect2v) + dimension); return 1.0f - ip / (norm_v1 * norm_v2); diff --git a/src/VecSim/spaces/L2/L2.cpp b/src/VecSim/spaces/L2/L2.cpp index 015f2200d..9a2589789 100644 --- a/src/VecSim/spaces/L2/L2.cpp +++ b/src/VecSim/spaces/L2/L2.cpp @@ -133,11 +133,13 @@ float FP16_L2Sqr(const void *pVect1, const void *pVect2, size_t dimension) { } // Return type for the L2 functions. -// The type should be able to hold `dimension * MAX_VAL(int_elem_t) * MAX_VAL(int_elem_t)`. -// To support dimension up to 2^16, we need the difference between the type and int_elem_t to be at -// least 2 bytes. We assert that in the implementation. +// The type must hold `dimension * MAX_VAL(int_elem_t) * MAX_VAL(int_elem_t)`. For uint8 that is +// 65025 * dimension, which overflows a 32-bit int from dimension 33,026 -- the alias was previously +// `int` for any 1-byte element, so UINT8_L2Sqr executed signed-overflow UB there. Signedness +// follows the element type, matching the inner product; both wrappers convert to float before +// returning, so either would be safe here. diff_t stays signed, which the assert below enforces. template -using ret_t = std::conditional_t; +using ret_t = std::conditional_t, uint64_t, int64_t>; // Difference type for the L2 functions. // The type should be able to hold `MIN_VAL(int_elem_t)-MAX_VAL(int_elem_t)`, and should be signed diff --git a/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_UINT8.h b/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_UINT8.h index 350b759ea..6ee1a1777 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_UINT8.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_UINT8.h @@ -7,6 +7,18 @@ * GNU Affero General Public License v3 (AGPLv3). */ #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" + +// uint8 L2: Imp returns the raw integer total and the wrappers convert it. The chooser picks +// plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h carries +// the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. static inline void L2SqrStep(uint8_t *&pVect1, uint8_t *&pVect2, __m512i &sum) { __m512i va = _mm512_loadu_epi8(pVect1); // AVX512BW @@ -32,8 +44,8 @@ static inline void L2SqrStep(uint8_t *&pVect1, uint8_t *&pVect2, __m512i &sum) { } template // 0..63 -float UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, - size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_L2SqrImp_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -92,5 +104,39 @@ float UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVec } while (pVect1 < pEnd1); } - return _mm512_reduce_add_epi32(sum); + // Unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return static_cast(_mm512_reduce_add_epi32(sum)); +} + +template // 0..63 +float UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, + size_t dimension) { + return static_cast( + UINT8_L2SqrImp_AVX512F_BW_VL_VNNI(pVect1v, pVect2v, dimension)); +} + +__attribute__((noinline)) static uint32_t +UINT8_L2SqrFullChunk_AVX512F_BW_VL_VNNI(const uint8_t *pVect1, const uint8_t *pVect2, + size_t dimension) { + return UINT8_L2SqrImp_AVX512F_BW_VL_VNNI<0>(pVect1, pVect2, dimension); +} + +template // 0..63 +struct UINT8_L2ChunkKernel_AVX512F_BW_VL_VNNI { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_AVX512F_BW_VL_VNNI(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrFullChunk_AVX512F_BW_VL_VNNI(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + return static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); } diff --git a/src/VecSim/spaces/L2/L2_NEON_DOTPROD_UINT8.h b/src/VecSim/spaces/L2/L2_NEON_DOTPROD_UINT8.h index 654c0b3b1..a0e5a9ebb 100644 --- a/src/VecSim/spaces/L2/L2_NEON_DOTPROD_UINT8.h +++ b/src/VecSim/spaces/L2/L2_NEON_DOTPROD_UINT8.h @@ -7,8 +7,20 @@ * GNU Affero General Public License v3 (AGPLv3). */ #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 L2: Imp returns the raw integer total and the wrappers convert it. The chooser picks +// plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h carries +// the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. + __attribute__((always_inline)) static inline void L2SquareOp(const uint8x16_t &v1, const uint8x16_t &v2, uint32x4_t &sum) { // Explicitly reinterpret the int8 vectors as uint8 for vabdq_u8 @@ -51,7 +63,8 @@ L2SquareStep32(uint8_t *&pVect1, uint8_t *&pVect2, uint32x4_t &sum1, uint32x4_t } template // 0..63 -float UINT8_L2SqrSIMD16_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_L2SqrImp_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -121,9 +134,36 @@ float UINT8_L2SqrSIMD16_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, s total_sum = vaddq_u32(total_sum, sum2); total_sum = vaddq_u32(total_sum, sum3); - // Horizontal sum of the 4 elements in the combined sum register - uint32_t result = vaddvq_u32(total_sum); + // Unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return vaddvq_u32(total_sum); +} + +template // 0..63 +float UINT8_L2SqrSIMD16_NEON_DOTPROD(const void *pVect1v, const void *pVect2v, size_t dimension) { + return static_cast(UINT8_L2SqrImp_NEON_DOTPROD(pVect1v, pVect2v, dimension)); +} + +__attribute__((noinline)) static uint32_t +UINT8_L2SqrFullChunk_NEON_DOTPROD(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_NEON_DOTPROD<0>(pVect1, pVect2, dimension); +} - // Return the L2 squared distance as a float - return static_cast(result); +template // 0..63 +struct UINT8_L2ChunkKernel_NEON_DOTPROD { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_NEON_DOTPROD(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrFullChunk_NEON_DOTPROD(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_L2SqrSIMD16_NEON_DOTPROD_Chunked(const void *pVect1v, const void *pVect2v, + size_t dimension) { + return static_cast( + spaces::uint8_chunked_total>(pVect1v, pVect2v, + dimension)); } diff --git a/src/VecSim/spaces/L2/L2_NEON_UINT8.h b/src/VecSim/spaces/L2/L2_NEON_UINT8.h index aa3769867..f1e461c18 100644 --- a/src/VecSim/spaces/L2/L2_NEON_UINT8.h +++ b/src/VecSim/spaces/L2/L2_NEON_UINT8.h @@ -7,8 +7,20 @@ * GNU Affero General Public License v3 (AGPLv3). */ #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 L2: Imp returns the raw integer total and the wrappers convert it. The chooser picks +// plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h carries +// the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk absorbs the residual and its length is a runtime min against +// the dimension, because a compile-time trip count cost 8-9.5% in accumulator copies; later chunks +// share one out-of-line copy of the kernel. + __attribute__((always_inline)) static inline void L2SquareOp(const uint8x16_t &v1, const uint8x16_t &v2, uint32x4_t &sum) { // Compute absolute differences and widen to 16-bit in one step @@ -53,7 +65,8 @@ L2SquareStep32(uint8_t *&pVect1, uint8_t *&pVect2, uint32x4_t &sum1, uint32x4_t } template // 0..63 -float UINT8_L2SqrSIMD16_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_L2SqrImp_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { uint8_t *pVect1 = (uint8_t *)pVect1v; uint8_t *pVect2 = (uint8_t *)pVect2v; @@ -125,9 +138,34 @@ float UINT8_L2SqrSIMD16_NEON(const void *pVect1v, const void *pVect2v, size_t di total_sum = vaddq_u32(total_sum, sum2); total_sum = vaddq_u32(total_sum, sum3); - // Horizontal sum of the 4 elements in the combined sum register - int32_t result = vaddvq_u32(total_sum); + // Unsigned, and exact for up to spaces::UINT8_CHUNK_ELEMENTS elements. + return vaddvq_u32(total_sum); +} + +template // 0..63 +float UINT8_L2SqrSIMD16_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { + return static_cast(UINT8_L2SqrImp_NEON(pVect1v, pVect2v, dimension)); +} + +__attribute__((noinline)) static uint32_t +UINT8_L2SqrFullChunk_NEON(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_NEON<0>(pVect1, pVect2, dimension); +} - // Return the L2 squared distance as a float - return static_cast(result); +template // 0..63 +struct UINT8_L2ChunkKernel_NEON { + static constexpr size_t granule() { return 64; } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_NEON(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrFullChunk_NEON(pVect1, pVect2, dimension); + } +}; + +template // 0..63 +float UINT8_L2SqrSIMD16_NEON_Chunked(const void *pVect1v, const void *pVect2v, size_t dimension) { + return static_cast(spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); } diff --git a/src/VecSim/spaces/L2/L2_SVE_UINT8.h b/src/VecSim/spaces/L2/L2_SVE_UINT8.h index 553db2169..8f05a7253 100644 --- a/src/VecSim/spaces/L2/L2_SVE_UINT8.h +++ b/src/VecSim/spaces/L2/L2_SVE_UINT8.h @@ -7,8 +7,20 @@ * GNU Affero General Public License v3 (AGPLv3). */ #include "VecSim/spaces/space_includes.h" +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS +#include "VecSim/spaces/uint8_chunking.h" #include +// uint8 L2: Imp returns the raw integer total and the wrappers convert it. The chooser picks +// plain up to spaces::UINT8_CHUNK_ELEMENTS and chunked above it, once per index; spaces.h carries +// the chunk-size argument. +// +// Imp is static and always_inline so the plain wrapper's codegen is unchanged now that Imp has +// several callers. +// The chunked wrapper's first chunk keeps this instantiation's residual shape, clamped to the +// dimension; the vector length is a runtime value so that split is computed rather than folded. +// Later chunks share one out-of-line copy of the kernel. + // Aligned step using svptrue_b8() inline void L2SquareStep(const uint8_t *&pVect1, const uint8_t *&pVect2, size_t &offset, svuint32_t &sum, const size_t chunk) { @@ -27,7 +39,8 @@ inline void L2SquareStep(const uint8_t *&pVect1, const uint8_t *&pVect2, size_t } template -float UINT8_L2SqrSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { +__attribute__((always_inline)) static inline uint32_t +UINT8_L2SqrImp_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { const uint8_t *pVect1 = reinterpret_cast(pVect1v); const uint8_t *pVect2 = reinterpret_cast(pVect2v); @@ -85,5 +98,36 @@ float UINT8_L2SqrSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimen sum0 = svadd_u32_x(all, sum0, sum1); sum2 = svadd_u32_x(all, sum2, sum3); svuint32_t sum_all = svadd_u32_x(all, sum0, sum2); - return svaddv_u32(svptrue_b32(), sum_all); + // Exact for up to spaces::UINT8_CHUNK_ELEMENTS elements; narrowed from svaddv_u32. + return static_cast(svaddv_u32(svptrue_b32(), sum_all)); +} + +template +float UINT8_L2SqrSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { + return static_cast( + UINT8_L2SqrImp_SVE(pVect1v, pVect2v, dimension)); +} + +__attribute__((noinline)) static uint32_t +UINT8_L2SqrFullChunk_SVE(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_SVE(pVect1, pVect2, dimension); +} + +template +struct UINT8_L2ChunkKernel_SVE { + static size_t granule() { return 4 * svcntb(); } + __attribute__((always_inline)) static inline uint32_t + first(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrImp_SVE(pVect1, pVect2, dimension); + } + static uint32_t rest(const uint8_t *pVect1, const uint8_t *pVect2, size_t dimension) { + return UINT8_L2SqrFullChunk_SVE(pVect1, pVect2, dimension); + } +}; + +template +float UINT8_L2SqrSIMD_SVE_Chunked(const void *pVect1v, const void *pVect2v, size_t dimension) { + return static_cast( + spaces::uint8_chunked_total>( + pVect1v, pVect2v, dimension)); } diff --git a/src/VecSim/spaces/L2_space.cpp b/src/VecSim/spaces/L2_space.cpp index 07e638cba..b4ca3669a 100644 --- a/src/VecSim/spaces/L2_space.cpp +++ b/src/VecSim/spaces/L2_space.cpp @@ -467,6 +467,7 @@ dist_func_t L2_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, } dist_func_t ret_dist_func = UINT8_L2Sqr; + // Optimizations assume at least 32 uint8. If we have less, we use the naive implementation. [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); diff --git a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp index 97da55546..c9f73c9ed 100644 --- a/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp +++ b/src/VecSim/spaces/functions/AVX512F_BW_VL_VNNI.cpp @@ -42,21 +42,43 @@ dist_func_t Choose_INT8_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t d return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant. Chosen here, once per +// index, so the plain kernel carries no branch and stays what it was. dist_func_t Choose_UINT8_L2_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD64_AVX512F_BW_VL_VNNI); + } return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant, which folds each chunk's +// exact 32-bit total into 64 bits. Chosen here, once per index, so the plain kernel below carries +// no branch and stays what it was. dist_func_t Choose_UINT8_IP_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD64_AVX512F_BW_VL_VNNI); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, + UINT8_InnerProductSIMD64_AVX512F_BW_VL_VNNI_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD64_AVX512F_BW_VL_VNNI); + } return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant, which folds each chunk's +// exact 32-bit total into 64 bits. Chosen here, once per index, so the plain kernel below carries +// no branch and stays what it was. dist_func_t Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD64_AVX512F_BW_VL_VNNI); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, + UINT8_CosineSIMD64_AVX512F_BW_VL_VNNI_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD64_AVX512F_BW_VL_VNNI); + } return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/NEON.cpp b/src/VecSim/spaces/functions/NEON.cpp index 0c9a286e3..d50bc28d5 100644 --- a/src/VecSim/spaces/functions/NEON.cpp +++ b/src/VecSim/spaces/functions/NEON.cpp @@ -30,9 +30,15 @@ dist_func_t Choose_INT8_IP_implementation_NEON(size_t dim) { return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant. Chosen here, once per +// index, so the plain kernel carries no branch and stays exactly what it was. dist_func_t Choose_UINT8_IP_implementation_NEON(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD16_NEON); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD16_NEON_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD16_NEON); + } return ret_dist_func; } @@ -56,7 +62,11 @@ dist_func_t Choose_INT8_Cosine_implementation_NEON(size_t dim) { dist_func_t Choose_UINT8_Cosine_implementation_NEON(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON); + } return ret_dist_func; } @@ -73,7 +83,11 @@ dist_func_t Choose_INT8_L2_implementation_NEON(size_t dim) { dist_func_t Choose_UINT8_L2_implementation_NEON(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON); + } return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/NEON_DOTPROD.cpp b/src/VecSim/spaces/functions/NEON_DOTPROD.cpp index 12f762093..bcf5b8d59 100644 --- a/src/VecSim/spaces/functions/NEON_DOTPROD.cpp +++ b/src/VecSim/spaces/functions/NEON_DOTPROD.cpp @@ -24,9 +24,16 @@ dist_func_t Choose_INT8_IP_implementation_NEON_DOTPROD(size_t dim) { return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant. Chosen here, once per +// index, so the plain kernel carries no branch and stays exactly what it was. dist_func_t Choose_UINT8_IP_implementation_NEON_DOTPROD(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD16_NEON_DOTPROD); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, + UINT8_InnerProductSIMD16_NEON_DOTPROD_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_InnerProductSIMD16_NEON_DOTPROD); + } return ret_dist_func; } @@ -38,7 +45,11 @@ dist_func_t Choose_INT8_Cosine_implementation_NEON_DOTPROD(size_t dim) { dist_func_t Choose_UINT8_Cosine_implementation_NEON_DOTPROD(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON_DOTPROD); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON_DOTPROD_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_CosineSIMD_NEON_DOTPROD); + } return ret_dist_func; } @@ -50,7 +61,11 @@ dist_func_t Choose_INT8_L2_implementation_NEON_DOTPROD(size_t dim) { dist_func_t Choose_UINT8_L2_implementation_NEON_DOTPROD(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON_DOTPROD); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON_DOTPROD_Chunked); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, UINT8_L2SqrSIMD16_NEON_DOTPROD); + } return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/SVE.cpp b/src/VecSim/spaces/functions/SVE.cpp index bd197c84c..5f8c1b625 100644 --- a/src/VecSim/spaces/functions/SVE.cpp +++ b/src/VecSim/spaces/functions/SVE.cpp @@ -86,21 +86,35 @@ dist_func_t Choose_INT8_Cosine_implementation_SVE(size_t dim) { return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant. Chosen here, once per +// index, so the plain kernel carries no branch and stays exactly what it was. dist_func_t Choose_UINT8_L2_implementation_SVE(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE, dim, svcntb); + } return ret_dist_func; } dist_func_t Choose_UINT8_IP_implementation_SVE(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE, dim, svcntb); + } return ret_dist_func; } dist_func_t Choose_UINT8_Cosine_implementation_SVE(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE, dim, svcntb); + } return ret_dist_func; } diff --git a/src/VecSim/spaces/functions/SVE2.cpp b/src/VecSim/spaces/functions/SVE2.cpp index 9eea81523..7c1a662ab 100644 --- a/src/VecSim/spaces/functions/SVE2.cpp +++ b/src/VecSim/spaces/functions/SVE2.cpp @@ -82,21 +82,35 @@ dist_func_t Choose_INT8_Cosine_implementation_SVE2(size_t dim) { return ret_dist_func; } +// Dimensions past spaces::UINT8_CHUNK_ELEMENTS use the chunked variant. Chosen here, once per +// index, so the plain kernel carries no branch and stays exactly what it was. dist_func_t Choose_UINT8_L2_implementation_SVE2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_L2SqrSIMD_SVE, dim, svcntb); + } return ret_dist_func; } dist_func_t Choose_UINT8_IP_implementation_SVE2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_InnerProductSIMD_SVE, dim, svcntb); + } return ret_dist_func; } dist_func_t Choose_UINT8_Cosine_implementation_SVE2(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE, dim, svcntb); + if (dim > spaces::UINT8_CHUNK_ELEMENTS) { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE_Chunked, dim, svcntb); + } else { + CHOOSE_SVE_IMPLEMENTATION(ret_dist_func, UINT8_CosineSIMD_SVE, dim, svcntb); + } return ret_dist_func; } diff --git a/src/VecSim/spaces/spaces.h b/src/VecSim/spaces/spaces.h index 11b0f9801..6419f9772 100644 --- a/src/VecSim/spaces/spaces.h +++ b/src/VecSim/spaces/spaces.h @@ -52,6 +52,25 @@ static int inline is_little_endian() { return *(char *)&x; } +// A full-range uint8 product is at most 255 * 255 = 65,025, so a 32-bit accumulator is exact +// through floor(UINT32_MAX / 65,025) = 66,051 terms. Twice the old signed limit of 33,025: the +// accumulation was always fine, the top bit was being read as a sign. +// +// Rather than cap the dimension there, the kernels accumulate in chunks of this many elements and +// fold each chunk's exact 32-bit total into a 64-bit scalar, which makes them exact at any +// dimension. 65,536 is chosen because it is under 66,051, so the existing 32-bit reduce needs no +// change, and because it is a whole number of 64-byte blocks, so a chunk boundary always lands on +// one. +// +// The margins are deliberately loose, because the tight version was wrong. A previous attempt +// widened the reduce and bounded the dimension at 4 * 66,051, on the assumption that products +// spread evenly across NEON's four lanes after its 32-bit vaddq_u32 merge. The even case already +// sat within 1,020 of UINT32_MAX while a masked residual load can put 1,040,400 into a single lane, +// so lanes wrapped before the widened reduce saw them. At this chunk size the per-chunk total has +// 33 million to spare and a NEON lane has 3.2 billion, so neither constraint is close and no +// per-ISA lane audit is needed. +static constexpr size_t UINT8_CHUNK_ELEMENTS = 65536; + static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) { #if defined(CPU_FEATURES_ARCH_AARCH64) diff --git a/src/VecSim/spaces/uint8_chunking.h b/src/VecSim/spaces/uint8_chunking.h new file mode 100644 index 000000000..cd3ea3c20 --- /dev/null +++ b/src/VecSim/spaces/uint8_chunking.h @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +// Shared chunked-accumulation driver for the uint8 SIMD kernels. It splits a distance +// computation into chunks of at most UINT8_CHUNK_ELEMENTS elements so each chunk's 32-bit +// SIMD total stays exact, and folds the per-chunk totals into a 64-bit scalar. The same +// formula serves both fixed-width kernels (granule 64) and SVE (granule 4 * svcntb()); only +// the granule differs. The caller supplies a Kernel adapter with: +// static size_t granule() - the kernel's block size +// static uint32_t first(const uint8_t *, const uint8_t *, size_t) - the residual-bearing +// kernel, shape already +// bound +// static uint32_t rest(const uint8_t *, const uint8_t *, size_t) - the out-of-line +// residual-0 kernel +// Invariants below hold for any Kernel whose granule() is in (0, UINT8_CHUNK_ELEMENTS]; all +// current adapters return 64 (fixed-width) or 4 * svcntb() (SVE, 64 to 1024 for a 16 to 256 byte +// vector length), so both stay within that range. Given that precondition: first <= dimension +// always, so this is correct at any dimension, including ones below the chunk size (the loop +// then does not execute). first is congruent to dimension modulo granule, so Kernel::first's +// residual shape still describes it. remaining is therefore a whole multiple of granule, and so +// is every step, which is Kernel::rest's precondition. No single call ever gets more than +// UINT8_CHUNK_ELEMENTS elements, which is what keeps each chunk's 32-bit total exact. + +#include "VecSim/spaces/spaces.h" // spaces::UINT8_CHUNK_ELEMENTS + +#include +#include +#include +#include + +namespace spaces { + +template +static inline uint64_t uint8_chunked_total(const void *pVect1v, const void *pVect2v, + size_t dimension) { + const auto *pVect1 = static_cast(pVect1v); + const auto *pVect2 = static_cast(pVect2v); + + constexpr size_t chunk = UINT8_CHUNK_ELEMENTS; + // Enforce the granule precondition at compile time when the adapter can express it, which is + // every fixed-width kernel. A plain assert would vanish under NDEBUG, so it is the fallback + // only for SVE, whose granule depends on the runtime vector length and cannot be constant. + if constexpr (requires { std::integral_constant{}; }) { + static_assert(Kernel::granule() > 0 && Kernel::granule() <= UINT8_CHUNK_ELEMENTS, + "Kernel::granule() must be in (0, UINT8_CHUNK_ELEMENTS]"); + } + // SVE is the only adapter whose granule cannot be constant, so it keeps the runtime assert and + // is unprotected under NDEBUG. That is acceptable because the architecture bounds it: an SVE + // vector is 16 to 256 bytes, so 4 * svcntb() is 64 to 1024, three orders of magnitude below the + // 65,536 limit. Only a change to that multiplier, or to the chunk size, could approach it. + const size_t granule = Kernel::granule(); + assert(granule > 0 && granule <= chunk); + const size_t tail = dimension % granule; + const size_t first_chunk = tail + ((chunk - tail) / granule) * granule; + const size_t first = dimension < first_chunk ? dimension : first_chunk; + const size_t max_step = (chunk / granule) * granule; + + uint64_t total = Kernel::first(pVect1, pVect2, first); + pVect1 += first; + pVect2 += first; + size_t remaining = dimension - first; + + while (remaining) { + const size_t step = remaining < max_step ? remaining : max_step; + total += Kernel::rest(pVect1, pVect2, step); + pVect1 += step; + pVect2 += step; + remaining -= step; + } + return total; +} + +} // namespace spaces diff --git a/tests/benchmark/spaces_benchmarks/bm_spaces_uint8.cpp b/tests/benchmark/spaces_benchmarks/bm_spaces_uint8.cpp index 602fff719..33f819936 100644 --- a/tests/benchmark/spaces_benchmarks/bm_spaces_uint8.cpp +++ b/tests/benchmark/spaces_benchmarks/bm_spaces_uint8.cpp @@ -31,12 +31,15 @@ class BM_VecSimSpaces_Integers_UINT8 : public benchmark::Fixture { test_utils::populate_uint8_vec(v2, dim, 1234); // Store the norm in the extra space for cosine calculations - *(float *)(v1 + dim) = test_utils::integral_compute_norm(v1, dim); - *(float *)(v2 + dim) = test_utils::integral_compute_norm(v2, dim); + // memcpy because v1 + dim is not guaranteed to be 4-byte aligned for arbitrary dim. + const float norm1 = test_utils::integral_compute_norm(v1, dim); + const float norm2 = test_utils::integral_compute_norm(v2, dim); + memcpy(v1 + dim, &norm1, sizeof(norm1)); + memcpy(v2 + dim, &norm2, sizeof(norm2)); } void TearDown(const ::benchmark::State &state) { - delete v1; - delete v2; + delete[] v1; + delete[] v2; } }; diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 8e53c83d1..ba8eb5776 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -13,6 +13,11 @@ #include #include #include +#include +#include +#include +#include +#include #include "gtest/gtest.h" #include "VecSim/spaces/space_includes.h" @@ -46,6 +51,7 @@ #include "VecSim/spaces/functions/SVE.h" #include "VecSim/spaces/functions/SVE_BF16.h" #include "VecSim/spaces/functions/SVE2.h" +#include "VecSim/spaces/uint8_chunking.h" #include "tests_utils.h" using bfloat16 = vecsim_types::bfloat16; @@ -2186,6 +2192,392 @@ TEST_P(UINT8SpacesOptimizationTest, UINT8_full_range_test) { INSTANTIATE_TEST_SUITE_P(UINT8OptFuncs, UINT8SpacesOptimizationTest, testing::Range(32UL, 64 * 2UL + 1)); +// The accumulated total is 255 * 255 * dim, which passes INT_MAX from dimension 33,026: the scalar +// path was signed-overflow UB there, and the AVX512 and NEON L2 reduces read their unsigned total +// back as a signed int and went negative. All-255 bytes are the worst case and make the expected +// value an exact integer. The existing UINT8 suites stop at dim 128, which is why this went unseen. +TEST_F(SpacesTest, UINT8_L2Sqr_and_InnerProduct_are_exact_past_int32) { + for (const size_t dim : {33026UL, 40000UL}) { + std::vector v1(dim + sizeof(float), 255); + std::vector v2(dim + sizeof(float), 0); + + // L2 between all-255 and all-0 is 255^2 * dim. + const double expected_l2 = 255.0 * 255.0 * static_cast(dim); + const float l2 = UINT8_L2Sqr(v1.data(), v2.data(), dim); + EXPECT_GT(l2, 0.0f) << "dim " << dim << ": squared distance went negative"; + EXPECT_LT(std::abs(static_cast(l2) - expected_l2) / expected_l2, 1e-6) + << "scalar L2, dim " << dim; + + unsigned char alignment = 0; + auto dispatched_l2 = L2_UINT8_GetDistFunc(dim, &alignment, nullptr); + const float l2_simd = dispatched_l2(v1.data(), v2.data(), dim); + EXPECT_GT(l2_simd, 0.0f) << "dim " << dim << ": SIMD squared distance went negative"; + EXPECT_LT(std::abs(static_cast(l2_simd) - expected_l2) / expected_l2, 1e-6) + << "dispatched L2, dim " << dim; + + // IP between two all-255 vectors is 255^2 * dim, and the kernel returns 1 - IP. + const double expected_ip = 1.0 - 255.0 * 255.0 * static_cast(dim); + const double ip = static_cast(UINT8_InnerProduct(v1.data(), v1.data(), dim)); + EXPECT_LT(std::abs(ip - expected_ip) / std::abs(expected_ip), 1e-6) + << "scalar IP, dim " << dim; + + auto dispatched_ip = IP_UINT8_GetDistFunc(dim, &alignment, nullptr); + const double ip_simd = static_cast(dispatched_ip(v1.data(), v1.data(), dim)); + EXPECT_LT(std::abs(ip_simd - expected_ip) / std::abs(expected_ip), 1e-6) + << "dispatched IP, dim " << dim; + } +} + +// Past spaces::UINT8_CHUNK_ELEMENTS the uint8 SIMD kernels accumulate in chunks: each chunk's total +// still fits the 32-bit accumulators (65025 * 65536 <= UINT32_MAX), and the per-chunk totals are +// folded in 64 bits. The dispatched kernel must therefore agree exactly with the scalar kernel, +// which accumulates the whole vector into a 64-bit ret_t. Exact equality is the right assertion +// because both paths convert the same integer total to float once, at the end. +// +// All-255 against all-0 is the worst case and puts the L2 total past UINT32_MAX from dimension +// 66,052, so the multi-chunk dimensions below genuinely exercise the 64-bit fold. The existing +// UINT8 suites stop at dim 128, which is why the wrap went unseen. +TEST_F(SpacesTest, UINT8_dispatched_kernels_are_exact_across_the_chunk_boundary) { + // Below the boundary, on it, one past it (whose last chunk is a single 64-element block), an + // exact multiple of it, and dimensions spanning two and three chunks. + for (const size_t dim : {65535UL, 65536UL, 65537UL, 65600UL, 131072UL, 131109UL, 200000UL}) { + // The cosine kernels read a float norm from just past the payload, so size for both. + std::vector ones(dim + sizeof(float), 255); + std::vector zeros(dim + sizeof(float), 0); + std::vector ramp(dim + sizeof(float)); + for (size_t i = 0; i < dim; i++) { + ramp[i] = static_cast(i % 256); + } + const float norm = std::sqrt(255.0f * 255.0f * static_cast(dim)); + memcpy(ones.data() + dim, &norm, sizeof(float)); + memcpy(ramp.data() + dim, &norm, sizeof(float)); + + unsigned char alignment = 0; + auto l2 = L2_UINT8_GetDistFunc(dim, &alignment, nullptr); + auto ip = IP_UINT8_GetDistFunc(dim, &alignment, nullptr); + auto cosine = Cosine_UINT8_GetDistFunc(dim, &alignment, nullptr); + + // Worst case: the largest total the byte range allows. + EXPECT_EQ(UINT8_L2Sqr(ones.data(), zeros.data(), dim), l2(ones.data(), zeros.data(), dim)) + << "L2 all-255 vs all-0, dim " << dim; + EXPECT_EQ(UINT8_InnerProduct(ones.data(), ones.data(), dim), + ip(ones.data(), ones.data(), dim)) + << "IP all-255, dim " << dim; + EXPECT_EQ(UINT8_Cosine(ones.data(), ones.data(), dim), + cosine(ones.data(), ones.data(), dim)) + << "Cosine all-255, dim " << dim; + + // A varying pattern, so the residual and chunk seams have to line up element for element + // rather than merely produce the right sum of identical values. + EXPECT_EQ(UINT8_L2Sqr(ramp.data(), ones.data(), dim), l2(ramp.data(), ones.data(), dim)) + << "L2 ramp vs all-255, dim " << dim; + EXPECT_EQ(UINT8_InnerProduct(ramp.data(), ones.data(), dim), + ip(ramp.data(), ones.data(), dim)) + << "IP ramp vs all-255, dim " << dim; + EXPECT_EQ(UINT8_Cosine(ramp.data(), ones.data(), dim), + cosine(ramp.data(), ones.data(), dim)) + << "Cosine ramp vs all-255, dim " << dim; + } +} + +// The boundary test above samples dimensions; this sweeps every residual instantiation. 65,600 and +// 196,608 are both multiples of 64, so base + r has residual r: one chunk past the boundary, then +// three chunks past it, so the seam between the residual-bearing first chunk and the residual-0 +// chunks after it is exercised for all 64 shapes. A ramp against all-255 is position sensitive, so +// a seam that double-counts or skips elements changes the total rather than cancelling out, and the +// total still passes UINT32_MAX (about 32,500 * dim) so the 64-bit fold is under test throughout. +TEST_F(SpacesTest, UINT8_dispatched_kernels_are_exact_at_every_residual_past_the_chunk_boundary) { + constexpr size_t max_dim = 196608 + 63; + std::vector ones(max_dim + sizeof(float), 255); + std::vector ramp(max_dim + sizeof(float)); + for (size_t i = 0; i < max_dim; i++) { + ramp[i] = static_cast(i % 256); + } + + for (const size_t base : {65600UL, 196608UL}) { + for (size_t r = 0; r < 64; r++) { + const size_t dim = base + r; + // The cosine kernels read a float norm from just past the payload, which moves with + // dim. + const float norm = std::sqrt(255.0f * 255.0f * static_cast(dim)); + memcpy(ones.data() + dim, &norm, sizeof(float)); + memcpy(ramp.data() + dim, &norm, sizeof(float)); + + unsigned char alignment = 0; + const void *a = ramp.data(); + const void *b = ones.data(); + + EXPECT_EQ(UINT8_L2Sqr(a, b, dim), + L2_UINT8_GetDistFunc(dim, &alignment, nullptr)(a, b, dim)) + << "L2 at dim " << dim << " (residual " << r << ")"; + EXPECT_EQ(UINT8_InnerProduct(a, b, dim), + IP_UINT8_GetDistFunc(dim, &alignment, nullptr)(a, b, dim)) + << "IP at dim " << dim << " (residual " << r << ")"; + EXPECT_EQ(UINT8_Cosine(a, b, dim), + Cosine_UINT8_GetDistFunc(dim, &alignment, nullptr)(a, b, dim)) + << "Cosine at dim " << dim << " (residual " << r << ")"; + } + } +} + +// Every uint8 SIMD tier this host can actually execute, with its three dispatched kernels. Both +// the per-tier exactness test and the independent-oracle test below iterate this list, so a tier +// cannot be covered by one and silently missed by the other. +struct UInt8TierFuncs { + const char *name; + dist_func_t l2; + dist_func_t ip; + dist_func_t cosine; +}; + +static std::vector AvailableUInt8Tiers(size_t dim) { + std::vector tiers; + [[maybe_unused]] const auto opt = getCpuOptimizationFeatures(); +#ifdef OPT_AVX512_F_BW_VL_VNNI + if (opt.avx512f && opt.avx512bw && opt.avx512vl && opt.avx512vnni) { + tiers.push_back({"AVX512F_BW_VL_VNNI", + Choose_UINT8_L2_implementation_AVX512F_BW_VL_VNNI(dim), + Choose_UINT8_IP_implementation_AVX512F_BW_VL_VNNI(dim), + Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI(dim)}); + } +#endif +#ifdef OPT_SVE2 + if (opt.sve2) { + tiers.push_back({"SVE2", Choose_UINT8_L2_implementation_SVE2(dim), + Choose_UINT8_IP_implementation_SVE2(dim), + Choose_UINT8_Cosine_implementation_SVE2(dim)}); + } +#endif +#ifdef OPT_SVE + if (opt.sve) { + tiers.push_back({"SVE", Choose_UINT8_L2_implementation_SVE(dim), + Choose_UINT8_IP_implementation_SVE(dim), + Choose_UINT8_Cosine_implementation_SVE(dim)}); + } +#endif +#ifdef OPT_NEON_DOTPROD + if (opt.asimddp) { + tiers.push_back({"NEON_DOTPROD", Choose_UINT8_L2_implementation_NEON_DOTPROD(dim), + Choose_UINT8_IP_implementation_NEON_DOTPROD(dim), + Choose_UINT8_Cosine_implementation_NEON_DOTPROD(dim)}); + } +#endif +#ifdef OPT_NEON + if (opt.asimd) { + tiers.push_back({"NEON", Choose_UINT8_L2_implementation_NEON(dim), + Choose_UINT8_IP_implementation_NEON(dim), + Choose_UINT8_Cosine_implementation_NEON(dim)}); + } +#endif + return tiers; +} + +// Worst-case inputs against an oracle computed here in 64-bit integers, rather than by calling the +// scalar kernel. +// +// Why not the scalar kernel: scalar and SIMD share conventions, and this series changed the scalar +// and SIMD inner product epilogues together so they would stay bit-identical. A test asserting only +// scalar == SIMD cannot catch that shared convention being wrong, in either sign or width. Here the +// expectation is derived from the inputs alone, and the scalar kernel is asserted against it on the +// same footing as every SIMD tier. +// +// Why these inputs: all-255 against all-255 puts 65,025 into the inner product accumulator for +// every element, and all-255 against all-0 does the same for L2. Those are the maxima the byte +// range allows, so they are where a 32-bit accumulator wraps first. A ramp against all-255 is +// carried alongside because constant data lets a gap and an overlap of equal size cancel, which a +// position-dependent pattern does not. +// +// Why these dimensions: 65024 is 1016*64, so 65024+r has residual r and stays at or below the +// 65,536 chunk size, exercising the plain kernel right up against the limit of its 32-bit reduce +// (65025 * 65087 is about 4.23e9, just under UINT32_MAX). 131072 is 2048*64, so 131072+r has +// residual r and its total is about 8.5e9, which only a 64-bit fold can carry. Every residual is +// swept at both. +TEST_F(SpacesTest, UINT8_worst_case_matches_an_independent_64bit_oracle) { + constexpr size_t max_dim = 131072 + 63; + std::vector ones(max_dim + sizeof(float), 255); + std::vector zeros(max_dim + sizeof(float), 0); + std::vector ramp(max_dim + sizeof(float)); + for (size_t i = 0; i < max_dim; i++) { + ramp[i] = static_cast(i % 256); + } + + for (const size_t base : {65024UL, 131072UL}) { + for (size_t r = 0; r < 64; r++) { + const size_t dim = base + r; + SCOPED_TRACE("dim " + std::to_string(dim) + " residual " + std::to_string(r)); + + // Norms live just past the payload and move with dim, so rewrite them per dimension. + const float norm_ones = std::sqrt(255.0f * 255.0f * static_cast(dim)); + float norm_ramp = 0.0f; + for (size_t i = 0; i < dim; i++) { + norm_ramp += static_cast(ramp[i]) * static_cast(ramp[i]); + } + norm_ramp = std::sqrt(norm_ramp); + memcpy(ones.data() + dim, &norm_ones, sizeof(float)); + memcpy(ramp.data() + dim, &norm_ramp, sizeof(float)); + + struct Pair { + const char *name; + const uint8_t *a; + const uint8_t *b; + float norm_a; + float norm_b; + bool check_cosine; + }; + const Pair pairs[] = { + {"all-255 vs all-255", ones.data(), ones.data(), norm_ones, norm_ones, true}, + {"all-255 vs all-0", ones.data(), zeros.data(), norm_ones, 0.0f, false}, + {"ramp vs all-255", ramp.data(), ones.data(), norm_ramp, norm_ones, true}, + }; + + for (const auto &pr : pairs) { + // The oracle: plain 64-bit integer accumulation over the inputs. + uint64_t ip_total = 0; + uint64_t l2_total = 0; + for (size_t i = 0; i < dim; i++) { + const uint64_t x = pr.a[i]; + const uint64_t y = pr.b[i]; + ip_total += x * y; + const int64_t diff = static_cast(x) - static_cast(y); + l2_total += static_cast(diff * diff); + } + // Both worst-case pairs must exceed a 32-bit accumulator at the multi-chunk base, + // otherwise this test would not be reaching the case it exists for. + if (base == 131072 && pr.b != ramp.data() && pr.a != ramp.data()) { + EXPECT_GT(std::max(ip_total, l2_total), + static_cast(std::numeric_limits::max())) + << "worst case no longer exceeds UINT32_MAX, test is not exercising the " + "fold"; + } + + // Expected returns, formed with the same operations the kernels use so the + // comparison can be exact rather than approximate. + const float want_ip = static_cast(1 - static_cast(ip_total)); + const float want_l2 = static_cast(l2_total); + const float want_cos = + 1.0f - static_cast(ip_total) / (pr.norm_a * pr.norm_b); + + EXPECT_EQ(want_l2, UINT8_L2Sqr(pr.a, pr.b, dim)) << "scalar L2, " << pr.name; + EXPECT_EQ(want_ip, UINT8_InnerProduct(pr.a, pr.b, dim)) << "scalar IP, " << pr.name; + if (pr.check_cosine) { + EXPECT_EQ(want_cos, UINT8_Cosine(pr.a, pr.b, dim)) + << "scalar cosine, " << pr.name; + } + + for (const auto &tier : AvailableUInt8Tiers(dim)) { + EXPECT_EQ(want_l2, tier.l2(pr.a, pr.b, dim)) + << "L2 " << tier.name << ", " << pr.name; + EXPECT_EQ(want_ip, tier.ip(pr.a, pr.b, dim)) + << "IP " << tier.name << ", " << pr.name; + if (pr.check_cosine) { + EXPECT_EQ(want_cos, tier.cosine(pr.a, pr.b, dim)) + << "cosine " << tier.name << ", " << pr.name; + } + } + } + } + } +} + +// The tests above go through the generic dispatcher, which only ever returns the best tier this +// host supports, so on an ARM machine with SVE the NEON and NEON_DOTPROD chunked kernels are never +// executed. Reach every compiled-in tier directly instead. Each tier is still gated on the CPU +// actually supporting it, since calling an unsupported kernel faults. +TEST_F(SpacesTest, UINT8_every_tier_is_exact_past_the_chunk_boundary) { + // Boundary (plain family), one past it with residuals 0/1/63, two chunks, and a ragged + // multiple. + const std::vector dims = {65536, 65600, 65601, 65663, 131072, 200000}; + std::set all_tiers; + + constexpr size_t max_dim = 200000; + std::vector ones(max_dim + sizeof(float), 255); + std::vector ramp(max_dim + sizeof(float)); + for (size_t i = 0; i < max_dim; i++) { + ramp[i] = static_cast(i % 256); + } + + for (const size_t dim : dims) { + const float norm = std::sqrt(255.0f * 255.0f * static_cast(dim)); + memcpy(ones.data() + dim, &norm, sizeof(float)); + memcpy(ramp.data() + dim, &norm, sizeof(float)); + const void *a = ramp.data(); + const void *b = ones.data(); + + const float want_l2 = UINT8_L2Sqr(a, b, dim); + const float want_ip = UINT8_InnerProduct(a, b, dim); + const float want_cos = UINT8_Cosine(a, b, dim); + + std::string covered; + for (const auto &tier : AvailableUInt8Tiers(dim)) { + EXPECT_EQ(want_l2, tier.l2(a, b, dim)) << "L2 " << tier.name << " dim " << dim; + EXPECT_EQ(want_ip, tier.ip(a, b, dim)) << "IP " << tier.name << " dim " << dim; + EXPECT_EQ(want_cos, tier.cosine(a, b, dim)) << "Cosine " << tier.name << " dim " << dim; + all_tiers.insert(tier.name); + covered += covered.empty() ? tier.name : std::string(", ") + tier.name; + } + RecordProperty("tiers_at_dim_" + std::to_string(dim), covered); + std::cout << " dim " << dim << " covered tiers: " << (covered.empty() ? "" : covered) + << std::endl; + } + + // Order matters. A stated requirement must be able to FAIL, so it is checked before the skip: + // hardware-specific CI sets VECSIM_REQUIRE_UINT8_TIER to the tier that job exists to cover + // (AVX512F_BW_VL_VNNI, SVE2, SVE, NEON_DOTPROD or NEON), and a mislabeled or silently + // downgraded runner then fails instead of quietly skipping. + // An empty value counts as unset. A CI expression that expands to nothing still puts the + // variable in the environment, so getenv returns a pointer to "" rather than nullptr, and + // treating that as a requirement fails every job that did not ask for one. + const char *required = std::getenv("VECSIM_REQUIRE_UINT8_TIER"); + if (required != nullptr && *required != '\0') { + EXPECT_TRUE(all_tiers.count(required) > 0) + << "VECSIM_REQUIRE_UINT8_TIER=" << required << " but that tier was not exercised. " + << "This host reached " << all_tiers.size() << " tier(s), so the run proves nothing " + << "about " << required; + return; + } + + // With no requirement stated, a run that exercised no tier proves nothing. Report it skipped + // rather than passed, because passing reads as coverage on a host that has none. + if (all_tiers.empty()) { + GTEST_SKIP() << "no uint8 SIMD tier on this host, no chunked kernel was executed"; + } +} + +// The chooser picks the chunked kernel once per index rather than branching per call, so assert the +// switch actually happens. Both dimensions are a multiple of 64, so they map to the same residual +// instantiation: any difference in the returned pointer can only come from the chunked family being +// chosen. Only meaningful where a uint8 SIMD tier exists, since otherwise both are the scalar +// kernel. +TEST_F(SpacesTest, UINT8_choosers_switch_to_the_chunked_kernel_past_the_chunk_size) { + const auto features = getCpuOptimizationFeatures(); + const bool has_uint8_simd = +#ifdef CPU_FEATURES_ARCH_X86_64 + features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni; +#else + features.sve2 || features.sve || features.asimddp || features.asimd; +#endif + if (!has_uint8_simd) { + GTEST_SKIP() << "no uint8 SIMD tier on this host"; + } + + constexpr size_t plain = spaces::UINT8_CHUNK_ELEMENTS; // on the boundary, not chunked + constexpr size_t chunked = spaces::UINT8_CHUNK_ELEMENTS * 2; // same residual, chunked + unsigned char alignment = 0; + + EXPECT_NE(L2_UINT8_GetDistFunc(plain, &alignment, nullptr), + L2_UINT8_GetDistFunc(chunked, &alignment, nullptr)); + EXPECT_NE(IP_UINT8_GetDistFunc(plain, &alignment, nullptr), + IP_UINT8_GetDistFunc(chunked, &alignment, nullptr)); + EXPECT_NE(Cosine_UINT8_GetDistFunc(plain, &alignment, nullptr), + Cosine_UINT8_GetDistFunc(chunked, &alignment, nullptr)); + + // And the SIMD kernel is still what gets chosen past the boundary: the chunked variant replaces + // the plain one, it does not fall back to scalar. + EXPECT_NE(L2_UINT8_GetDistFunc(chunked, &alignment, nullptr), UINT8_L2Sqr); + EXPECT_NE(IP_UINT8_GetDistFunc(chunked, &alignment, nullptr), UINT8_InnerProduct); + EXPECT_NE(Cosine_UINT8_GetDistFunc(chunked, &alignment, nullptr), UINT8_Cosine); +} + class SQ8_FP32_SpacesOptimizationTest : public testing::TestWithParam {}; TEST_P(SQ8_FP32_SpacesOptimizationTest, SQ8_FP32_L2SqrTest) { @@ -4684,6 +5076,174 @@ TEST(SQ8_SQ8_EdgeCases, L2ExtremeValuesTest) { ASSERT_NEAR(result, baseline, 0.01f) << "Extreme values L2 should match baseline"; } +// spaces::uint8_chunked_total (uint8_chunking.h) is the chunked-accumulation driver shared by +// every uint8 SIMD kernel: it tiles a vector into segments no larger than UINT8_CHUNK_ELEMENTS +// so each segment's 32-bit SIMD partial sum stays exact, then folds the per-segment totals into +// a 64-bit scalar. The driver only ever calls Kernel::granule/first/rest, so it can be exercised +// directly with a mock kernel instead of a real SIMD kernel, which means this test needs no +// architecture-specific build flag and runs the same way on every host: an x86 box without +// AVX512 and an ARM box exercise identical logic here, closing the gap where this driver was +// previously only reachable through whichever hardware-specific kernel happened to be present. +// Real coverage of "did the driver visit every element exactly once, correctly" comes from the +// value check below: the mock's first()/rest() return the sum of the bytes in the slice they +// were handed (read from a position-dependent, non-constant fill), and the total the driver +// returns is compared against an independent, trivially-correct sum over the whole buffer. A +// skipped element, a double-counted element, or a mis-sized chunk changes that sum; it cannot +// cancel out the way it could with a constant fill or a return value of 0. The mock also records +// the (offset, length) of every call; the tiling check on those recordings does not by itself +// prove the driver visited the right elements (offset is derived from the same length the driver +// just advanced its pointer by, so "no gap/overlap" holds by construction), but it does guard the +// coupling between the length passed to the kernel and the distance the pointer is advanced, plus +// the length-shape properties below (granule multiples, chunk-size cap, congruence). Together the +// two checks cover a sweep of granules (64, 128, 192, 256 and 1024, standing in for fixed-width +// kernels and SVE vector lengths of 32/48/64/256 bytes) and dimensions chosen to cover every +// residue class modulo the granule across one-, two- and three-segment cases, plus the boundary +// around UINT8_CHUNK_ELEMENTS itself. +TEST_F(SpacesTest, UINT8_chunked_driver_tiles_the_vector_exactly) { + struct RecordedCall { + size_t offset; + size_t length; + }; + + // Local mock adapter matching the Kernel contract from uint8_chunking.h. All state lives in + // function-local statics reached through static member functions (a local class cannot have + // static data members), so `reset` must be called before each driver invocation. + struct RecordingKernel { + static size_t granule() { return granule_ref(); } + + // Returns the sum of the bytes in [v1, v1 + length), not 0: combined with a + // position-dependent fill, this makes the driver's return value an end-to-end proof + // that every element was visited exactly once, not just a coupling check on lengths. + static uint32_t first(const uint8_t *v1, const uint8_t *, size_t length) { + record(v1, length); + return sum_of(v1, length); + } + + static uint32_t rest(const uint8_t *v1, const uint8_t *, size_t length) { + record(v1, length); + return sum_of(v1, length); + } + + static void reset(const uint8_t *base, size_t granule) { + calls_ref().clear(); + base_ref() = base; + granule_ref() = granule; + } + + static const std::vector &calls_seen() { return calls_ref(); } + + private: + static uint32_t sum_of(const uint8_t *v1, size_t length) { + uint32_t sum = 0; + for (size_t i = 0; i < length; i++) { + sum += v1[i]; + } + return sum; + } + + static void record(const uint8_t *v1, size_t length) { + calls_ref().push_back({static_cast(v1 - base_ref()), length}); + } + + static std::vector &calls_ref() { + static std::vector calls; + return calls; + } + + static const uint8_t *&base_ref() { + static const uint8_t *base = nullptr; + return base; + } + + static size_t &granule_ref() { + static size_t granule = 0; + return granule; + } + }; + + constexpr size_t chunk = spaces::UINT8_CHUNK_ELEMENTS; + // Large enough for the biggest dimension exercised below (first segment plus two full + // max-size segments, bounded by chunk), with slack. + constexpr size_t buffer_size = 3 * chunk + 4096; + // v1 is filled with a position-dependent, non-constant pattern so a skipped, duplicated or + // mis-sized element changes the summed value rather than cancelling out (a constant fill, or + // returning 0 from the mock, would not catch that). Values stay under 251 and dimensions + // stay well under 600,000, so the accumulated uint64_t sum cannot overflow. v2 is unused by + // the mock kernel and left zero-filled. + std::vector v1(buffer_size); + for (size_t i = 0; i < buffer_size; i++) { + v1[i] = static_cast((i * 31 + 7) % 251); + } + std::vector v2(buffer_size, 0); + + const std::array granules = {64, 128, 192, 256, 1024}; + + for (size_t granule : granules) { + const size_t max_step = (chunk / granule) * granule; + auto first_chunk_for = [&](size_t tail) { + return tail + ((chunk - tail) / granule) * granule; + }; + + std::vector dims; + for (size_t r = 0; r < granule; r++) { + const size_t fc = first_chunk_for(r); + dims.push_back(r == 0 ? granule : r); // one segment: dim <= chunk + dims.push_back(fc + max_step); // two segments + dims.push_back(fc + 2 * max_step); // three segments + } + dims.push_back(chunk - 1); + dims.push_back(chunk); + dims.push_back(chunk + 1); + + for (size_t dimension : dims) { + ASSERT_LE(dimension, buffer_size) + << "granule=" << granule << " dimension=" << dimension; + + RecordingKernel::reset(v1.data(), granule); + const uint64_t total = + spaces::uint8_chunked_total(v1.data(), v2.data(), dimension); + + SCOPED_TRACE("granule=" + std::to_string(granule) + + " dimension=" + std::to_string(dimension)); + + // Value check: independently sum the same byte range the driver was asked to cover. + // This is what actually proves every element was visited exactly once (a skipped, + // duplicated or mis-sized chunk changes this sum); the tiling check below only + // proves the length passed to the kernel matches how far the pointer advanced. + uint64_t expected = 0; + for (size_t i = 0; i < dimension; i++) { + expected += v1[i]; + } + EXPECT_EQ(total, expected) << "driver total does not match independent byte sum"; + + const auto &calls = RecordingKernel::calls_seen(); + ASSERT_FALSE(calls.empty()); + + size_t sum = 0; + size_t expected_offset = 0; + for (size_t i = 0; i < calls.size(); i++) { + EXPECT_EQ(calls[i].offset, expected_offset) + << "call " << i << " does not tile contiguously (gap or overlap)"; + EXPECT_LE(calls[i].length, chunk) + << "call " << i << " exceeds UINT8_CHUNK_ELEMENTS"; + if (i > 0) { + EXPECT_EQ(calls[i].length % granule, 0u) + << "call " << i << " length is not a whole multiple of granule"; + } + sum += calls[i].length; + expected_offset += calls[i].length; + } + EXPECT_EQ(sum, dimension) << "recorded lengths do not sum to the dimension"; + EXPECT_EQ(calls[0].length % granule, dimension % granule) + << "first call length is not congruent to dimension modulo granule"; + if (dimension <= chunk) { + EXPECT_EQ(calls.size(), 1u) + << "dimension at or below UINT8_CHUNK_ELEMENTS should need exactly one call"; + } + } + } +} + // Assert the exact alignment-hint values published by the SQ8 distance dispatchers. // The hint refers to the SQ8 (first / storage) operand per the GetDistFunc contract documented // in spaces/spaces.h. These tests guard against silent regressions of the per-kernel hints used