diff --git a/src/distance-avx2.c b/src/distance-avx2.c index 243d772..efd52b6 100644 --- a/src/distance-avx2.c +++ b/src/distance-avx2.c @@ -643,69 +643,78 @@ float bfloat16_distance_cosine_avx2 (const void *v1, const void *v2, int n) { // MARK: - UINT8 - +// The integer kernels used to widen every byte to 32 bits before multiplying, on a single +// accumulator chain. x86 has instructions for exactly this shape: PSADBW sums absolute +// differences of 32 bytes in one go (L1), and PMADDWD multiplies 16-bit pairs and adds +// adjacent products into 32-bit lanes, which is a squared difference or a dot product +// depending on what you feed it. Two byte-sized factors always fit 16 bits, and PMADDWD's +// pairwise add keeps the running value inside 32 bits. +// +// The signed kernels bias by 0x80 where the arithmetic allows it: that maps int8 onto +// uint8 without changing any difference between two elements, so L2 and L1 become the +// unsigned kernel plus one XOR per vector. Dot and cosine need the true signed values. +static inline __m256i abs_diff_epu8 (__m256i a, __m256i b) { + return _mm256_or_si256(_mm256_subs_epu8(a, b), _mm256_subs_epu8(b, a)); +} + +static inline uint64_t hsum256_epi64 (__m256i v) { + __m128i lo = _mm256_castsi256_si128(v); + __m128i hi = _mm256_extracti128_si256(v, 1); + __m128i s = _mm_add_epi64(lo, hi); + return (uint64_t)_mm_cvtsi128_si64(s) + (uint64_t)_mm_extract_epi64(s, 1); +} + +// The reductions widen to 64 bits before summing the lanes. Each lane is itself a running +// total, so folding eight of them in 32 bits would cap the usable dimension far below +// what the per-lane accumulators can hold. +static inline uint64_t hsum256_epi32 (__m256i v) { + const __m256i zero = _mm256_setzero_si256(); + return hsum256_epi64(_mm256_add_epi64(_mm256_unpacklo_epi32(v, zero), _mm256_unpackhi_epi32(v, zero))); +} + +static inline int64_t hsum256_epi32_signed (__m256i v) { + __m256i wl = _mm256_cvtepi32_epi64(_mm256_castsi256_si128(v)); + __m256i wh = _mm256_cvtepi32_epi64(_mm256_extracti128_si256(v, 1)); + __m256i s = _mm256_add_epi64(wl, wh); + __m128i lo = _mm256_castsi256_si128(s); + __m128i hi = _mm256_extracti128_si256(s, 1); + __m128i r = _mm_add_epi64(lo, hi); + return (int64_t)_mm_cvtsi128_si64(r) + (int64_t)_mm_extract_epi64(r, 1); +} + +// sum of squared differences of 32 bytes, accumulated 32 bits wide +static inline __m256i sqdiff_epu8 (__m256i a, __m256i b) { + __m256i d = abs_diff_epu8(a, b); + __m256i lo = _mm256_unpacklo_epi8(d, _mm256_setzero_si256()); + __m256i hi = _mm256_unpackhi_epi8(d, _mm256_setzero_si256()); + return _mm256_add_epi32(_mm256_madd_epi16(lo, lo), _mm256_madd_epi16(hi, hi)); +} + static inline float uint8_distance_l2_impl_avx2 (const void *v1, const void *v2, int n, bool use_sqrt) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; - + + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi32(acc0, sqdiff_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); + acc1 = _mm256_add_epi32(acc1, sqdiff_epu8(_mm256_loadu_si256((const __m256i *)(a + i + 32)), + _mm256_loadu_si256((const __m256i *)(b + i + 32)))); + } for (; i <= n - 32; i += 32) { - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); - - // Split into 2 x 128-bit chunks - __m128i va_lo = _mm256_extracti128_si256(va, 0); - __m128i va_hi = _mm256_extracti128_si256(va, 1); - __m128i vb_lo = _mm256_extracti128_si256(vb, 0); - __m128i vb_hi = _mm256_extracti128_si256(vb, 1); - - // Unpack to 16-bit integers - __m128i va_lo_u16 = _mm_unpacklo_epi8(va_lo, _mm_setzero_si128()); - __m128i va_hi_u16 = _mm_unpackhi_epi8(va_lo, _mm_setzero_si128()); - __m128i vb_lo_u16 = _mm_unpacklo_epi8(vb_lo, _mm_setzero_si128()); - __m128i vb_hi_u16 = _mm_unpackhi_epi8(vb_lo, _mm_setzero_si128()); - - __m128i va_lo_u16_hi = _mm_unpacklo_epi8(va_hi, _mm_setzero_si128()); - __m128i va_hi_u16_hi = _mm_unpackhi_epi8(va_hi, _mm_setzero_si128()); - __m128i vb_lo_u16_hi = _mm_unpacklo_epi8(vb_hi, _mm_setzero_si128()); - __m128i vb_hi_u16_hi = _mm_unpackhi_epi8(vb_hi, _mm_setzero_si128()); - - // Compute diffs - __m128i d0 = _mm_sub_epi16(va_lo_u16, vb_lo_u16); - __m128i d1 = _mm_sub_epi16(va_hi_u16, vb_hi_u16); - __m128i d2 = _mm_sub_epi16(va_lo_u16_hi, vb_lo_u16_hi); - __m128i d3 = _mm_sub_epi16(va_hi_u16_hi, vb_hi_u16_hi); - - // Square diffs - __m128i s0 = _mm_mullo_epi16(d0, d0); - __m128i s1 = _mm_mullo_epi16(d1, d1); - __m128i s2 = _mm_mullo_epi16(d2, d2); - __m128i s3 = _mm_mullo_epi16(d3, d3); - - // Widen to 32-bit and accumulate - __m256i w0 = _mm256_cvtepu16_epi32(s0); - __m256i w1 = _mm256_cvtepu16_epi32(s1); - __m256i w2 = _mm256_cvtepu16_epi32(s2); - __m256i w3 = _mm256_cvtepu16_epi32(s3); - - acc = _mm256_add_epi32(acc, w0); - acc = _mm256_add_epi32(acc, w1); - acc = _mm256_add_epi32(acc, w2); - acc = _mm256_add_epi32(acc, w3); + acc0 = _mm256_add_epi32(acc0, sqdiff_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); } - - // Horizontal sum of 8 x 32-bit integers - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - uint32_t total = temp[0] + temp[1] + temp[2] + temp[3] + temp[4] + temp[5] + temp[6] + temp[7]; - - // Tail loop + + uint64_t total = hsum256_epi32(_mm256_add_epi32(acc0, acc1)); + for (; i < n; ++i) { int d = (int)a[i] - (int)b[i]; - total += d * d; + total += (uint64_t)(d * d); } - + return use_sqrt ? sqrtf((float)total) : (float)total; } @@ -717,98 +726,96 @@ float uint8_distance_l2_squared_avx2 (const void *v1, const void *v2, int n) { return uint8_distance_l2_impl_avx2(v1, v2, n, false); } +// widening dot product of 32 unsigned bytes +static inline __m256i dot_epu8 (__m256i a, __m256i b) { + const __m256i zero = _mm256_setzero_si256(); + __m256i al = _mm256_unpacklo_epi8(a, zero), ah = _mm256_unpackhi_epi8(a, zero); + __m256i bl = _mm256_unpacklo_epi8(b, zero), bh = _mm256_unpackhi_epi8(b, zero); + return _mm256_add_epi32(_mm256_madd_epi16(al, bl), _mm256_madd_epi16(ah, bh)); +} + float uint8_distance_dot_avx2 (const void *v1, const void *v2, int n) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi32(acc0, dot_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); + acc1 = _mm256_add_epi32(acc1, dot_epu8(_mm256_loadu_si256((const __m256i *)(a + i + 32)), + _mm256_loadu_si256((const __m256i *)(b + i + 32)))); + } for (; i <= n - 32; i += 32) { - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); - - __m256i a_lo = _mm256_unpacklo_epi8(va, _mm256_setzero_si256()); - __m256i a_hi = _mm256_unpackhi_epi8(va, _mm256_setzero_si256()); - __m256i b_lo = _mm256_unpacklo_epi8(vb, _mm256_setzero_si256()); - __m256i b_hi = _mm256_unpackhi_epi8(vb, _mm256_setzero_si256()); - - __m256i prod_lo = _mm256_mullo_epi16(a_lo, b_lo); - __m256i prod_hi = _mm256_mullo_epi16(a_hi, b_hi); - - __m256i prod_lo_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(prod_lo, 0)); - __m256i prod_hi_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(prod_lo, 1)); - acc = _mm256_add_epi32(acc, prod_lo_32); - acc = _mm256_add_epi32(acc, prod_hi_32); - - prod_lo_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(prod_hi, 0)); - prod_hi_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(prod_hi, 1)); - acc = _mm256_add_epi32(acc, prod_lo_32); - acc = _mm256_add_epi32(acc, prod_hi_32); + acc0 = _mm256_add_epi32(acc0, dot_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); } - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - uint32_t total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + uint64_t dot = hsum256_epi32(_mm256_add_epi32(acc0, acc1)); - for (; i < n; ++i) { - total += a[i] * b[i]; - } + for (; i < n; ++i) dot += (uint64_t)((uint32_t)a[i] * (uint32_t)b[i]); - return -(float)total; + return -(float)dot; } float uint8_distance_l1_avx2 (const void *v1, const void *v2, int n) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + // PSADBW already sums absolute differences into 64-bit lanes: one instruction per + // 32 bytes, and it cannot overflow for any vector length we support + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi64(acc0, _mm256_sad_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); + acc1 = _mm256_add_epi64(acc1, _mm256_sad_epu8(_mm256_loadu_si256((const __m256i *)(a + i + 32)), + _mm256_loadu_si256((const __m256i *)(b + i + 32)))); + } for (; i <= n - 32; i += 32) { - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); + acc0 = _mm256_add_epi64(acc0, _mm256_sad_epu8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); + } - __m256i a_lo = _mm256_unpacklo_epi8(va, _mm256_setzero_si256()); - __m256i a_hi = _mm256_unpackhi_epi8(va, _mm256_setzero_si256()); - __m256i b_lo = _mm256_unpacklo_epi8(vb, _mm256_setzero_si256()); - __m256i b_hi = _mm256_unpackhi_epi8(vb, _mm256_setzero_si256()); + uint64_t sum = hsum256_epi64(_mm256_add_epi64(acc0, acc1)); - __m256i diff_lo = _mm256_abs_epi16(_mm256_sub_epi16(a_lo, b_lo)); - __m256i diff_hi = _mm256_abs_epi16(_mm256_sub_epi16(a_hi, b_hi)); + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - __m256i diff_lo_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(diff_lo, 0)); - __m256i diff_hi_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(diff_lo, 1)); - acc = _mm256_add_epi32(acc, diff_lo_32); - acc = _mm256_add_epi32(acc, diff_hi_32); + return (float)sum; +} - diff_lo_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(diff_hi, 0)); - diff_hi_32 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(diff_hi, 1)); - acc = _mm256_add_epi32(acc, diff_lo_32); - acc = _mm256_add_epi32(acc, diff_hi_32); - } +float uint8_distance_cosine_avx2 (const void *a, const void *b, int n) { + const uint8_t *x = (const uint8_t *)a; + const uint8_t *y = (const uint8_t *)b; - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - uint32_t total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + // one fused pass rather than three calls to the dot kernel + __m256i dacc = _mm256_setzero_si256(), aacc = dacc, bacc = dacc; + int i = 0; - for (; i < n; ++i) { - total += abs((int)a[i] - (int)b[i]); + for (; i <= n - 32; i += 32) { + __m256i va = _mm256_loadu_si256((const __m256i *)(x + i)); + __m256i vb = _mm256_loadu_si256((const __m256i *)(y + i)); + dacc = _mm256_add_epi32(dacc, dot_epu8(va, vb)); + aacc = _mm256_add_epi32(aacc, dot_epu8(va, va)); + bacc = _mm256_add_epi32(bacc, dot_epu8(vb, vb)); } - return (float)total; -} + uint64_t dot = hsum256_epi32(dacc); + uint64_t norm_a = hsum256_epi32(aacc); + uint64_t norm_b = hsum256_epi32(bacc); -float uint8_distance_cosine_avx2 (const void *a, const void *b, int n) { - float dot = -uint8_distance_dot_avx2(a, b, n); - float norm_a = sqrtf(-uint8_distance_dot_avx2(a, a, n)); - float norm_b = sqrtf(-uint8_distance_dot_avx2(b, b, n)); + for (; i < n; ++i) { + uint32_t p = x[i], q = y[i]; + dot += (uint64_t)(p * q); + norm_a += (uint64_t)(p * p); + norm_b += (uint64_t)(q * q); + } - if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; + if (norm_a == 0 || norm_b == 0) return 1.0f; - float cosine_similarity = dot / (norm_a * norm_b); + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; @@ -816,69 +823,31 @@ float uint8_distance_cosine_avx2 (const void *a, const void *b, int n) { // MARK: - INT8 - +#define S8_TO_BIASED_U8_AVX2(_v) _mm256_xor_si256((_v), _mm256_set1_epi8((char)0x80)) + static inline float int8_distance_l2_impl_avx2 (const void *v1, const void *v2, int n, bool use_sqrt) { const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi32(acc0, sqdiff_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i))))); + acc1 = _mm256_add_epi32(acc1, sqdiff_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i + 32))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i + 32))))); + } for (; i <= n - 32; i += 32) { - // Load 32 int8_t elements from each input - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); - - // Extract 128-bit halves - __m128i va_lo = _mm256_extracti128_si256(va, 0); - __m128i va_hi = _mm256_extracti128_si256(va, 1); - __m128i vb_lo = _mm256_extracti128_si256(vb, 0); - __m128i vb_hi = _mm256_extracti128_si256(vb, 1); - - // Sign-extend int8_t to int16_t - __m128i va_lo_s16 = _mm_cvtepi8_epi16(va_lo); - __m128i va_hi_s16 = _mm_cvtepi8_epi16(_mm_srli_si128(va_lo, 8)); - __m128i vb_lo_s16 = _mm_cvtepi8_epi16(vb_lo); - __m128i vb_hi_s16 = _mm_cvtepi8_epi16(_mm_srli_si128(vb_lo, 8)); - - __m128i va_lo_s16_hi = _mm_cvtepi8_epi16(va_hi); - __m128i va_hi_s16_hi = _mm_cvtepi8_epi16(_mm_srli_si128(va_hi, 8)); - __m128i vb_lo_s16_hi = _mm_cvtepi8_epi16(vb_hi); - __m128i vb_hi_s16_hi = _mm_cvtepi8_epi16(_mm_srli_si128(vb_hi, 8)); - - // Compute differences - __m128i d0 = _mm_sub_epi16(va_lo_s16, vb_lo_s16); - __m128i d1 = _mm_sub_epi16(va_hi_s16, vb_hi_s16); - __m128i d2 = _mm_sub_epi16(va_lo_s16_hi, vb_lo_s16_hi); - __m128i d3 = _mm_sub_epi16(va_hi_s16_hi, vb_hi_s16_hi); - - // Square differences - __m128i s0 = _mm_mullo_epi16(d0, d0); - __m128i s1 = _mm_mullo_epi16(d1, d1); - __m128i s2 = _mm_mullo_epi16(d2, d2); - __m128i s3 = _mm_mullo_epi16(d3, d3); - - // Extend to 32-bit and accumulate - __m256i w0 = _mm256_cvtepu16_epi32(s0); - __m256i w1 = _mm256_cvtepu16_epi32(s1); - __m256i w2 = _mm256_cvtepu16_epi32(s2); - __m256i w3 = _mm256_cvtepu16_epi32(s3); - - acc = _mm256_add_epi32(acc, w0); - acc = _mm256_add_epi32(acc, w1); - acc = _mm256_add_epi32(acc, w2); - acc = _mm256_add_epi32(acc, w3); + acc0 = _mm256_add_epi32(acc0, sqdiff_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i))))); } - // Horizontal sum - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - uint32_t total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + uint64_t total = hsum256_epi32(_mm256_add_epi32(acc0, acc1)); - // Scalar tail for (; i < n; ++i) { int d = (int)a[i] - (int)b[i]; - total += d * d; + total += (uint64_t)(d * d); } return use_sqrt ? sqrtf((float)total) : (float)total; @@ -892,122 +861,95 @@ float int8_distance_l2_squared_avx2 (const void *v1, const void *v2, int n) { return int8_distance_l2_impl_avx2(v1, v2, n, false); } +// widening dot product of 32 signed bytes: sign-extend each half, then PMADDWD +static inline __m256i dot_epi8 (__m256i a, __m256i b) { + __m256i al = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(a)); + __m256i ah = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(a, 1)); + __m256i bl = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(b)); + __m256i bh = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(b, 1)); + return _mm256_add_epi32(_mm256_madd_epi16(al, bl), _mm256_madd_epi16(ah, bh)); +} + float int8_distance_dot_avx2 (const void *v1, const void *v2, int n) { const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi32(acc0, dot_epi8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); + acc1 = _mm256_add_epi32(acc1, dot_epi8(_mm256_loadu_si256((const __m256i *)(a + i + 32)), + _mm256_loadu_si256((const __m256i *)(b + i + 32)))); + } for (; i <= n - 32; i += 32) { - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); - - __m128i va_lo = _mm256_extracti128_si256(va, 0); - __m128i va_hi = _mm256_extracti128_si256(va, 1); - __m128i vb_lo = _mm256_extracti128_si256(vb, 0); - __m128i vb_hi = _mm256_extracti128_si256(vb, 1); - - __m128i a0 = _mm_cvtepi8_epi16(va_lo); - __m128i a1 = _mm_cvtepi8_epi16(_mm_srli_si128(va_lo, 8)); - __m128i b0 = _mm_cvtepi8_epi16(vb_lo); - __m128i b1 = _mm_cvtepi8_epi16(_mm_srli_si128(vb_lo, 8)); - - __m128i a2 = _mm_cvtepi8_epi16(va_hi); - __m128i a3 = _mm_cvtepi8_epi16(_mm_srli_si128(va_hi, 8)); - __m128i b2 = _mm_cvtepi8_epi16(vb_hi); - __m128i b3 = _mm_cvtepi8_epi16(_mm_srli_si128(vb_hi, 8)); - - __m128i p0 = _mm_mullo_epi16(a0, b0); - __m128i p1 = _mm_mullo_epi16(a1, b1); - __m128i p2 = _mm_mullo_epi16(a2, b2); - __m128i p3 = _mm_mullo_epi16(a3, b3); - - __m256i w0 = _mm256_cvtepi16_epi32(p0); - __m256i w1 = _mm256_cvtepi16_epi32(p1); - __m256i w2 = _mm256_cvtepi16_epi32(p2); - __m256i w3 = _mm256_cvtepi16_epi32(p3); - - acc = _mm256_add_epi32(acc, w0); - acc = _mm256_add_epi32(acc, w1); - acc = _mm256_add_epi32(acc, w2); - acc = _mm256_add_epi32(acc, w3); + acc0 = _mm256_add_epi32(acc0, dot_epi8(_mm256_loadu_si256((const __m256i *)(a + i)), + _mm256_loadu_si256((const __m256i *)(b + i)))); } - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - int32_t total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + int64_t dot = hsum256_epi32_signed(_mm256_add_epi32(acc0, acc1)); - for (; i < n; ++i) { - total += (int)a[i] * (int)b[i]; - } + for (; i < n; ++i) dot += (int64_t)((int32_t)a[i] * (int32_t)b[i]); - return -(float)total; + return -(float)dot; } float int8_distance_l1_avx2 (const void *v1, const void *v2, int n) { const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - - __m256i acc = _mm256_setzero_si256(); + + // same biasing trick as L2, so PSADBW applies unchanged + __m256i acc0 = _mm256_setzero_si256(), acc1 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm256_add_epi64(acc0, _mm256_sad_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i))))); + acc1 = _mm256_add_epi64(acc1, _mm256_sad_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i + 32))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i + 32))))); + } for (; i <= n - 32; i += 32) { - __m256i va = _mm256_loadu_si256((const __m256i *)(a + i)); - __m256i vb = _mm256_loadu_si256((const __m256i *)(b + i)); - - __m128i va_lo = _mm256_extracti128_si256(va, 0); - __m128i va_hi = _mm256_extracti128_si256(va, 1); - __m128i vb_lo = _mm256_extracti128_si256(vb, 0); - __m128i vb_hi = _mm256_extracti128_si256(vb, 1); - - __m128i a0 = _mm_cvtepi8_epi16(va_lo); - __m128i a1 = _mm_cvtepi8_epi16(_mm_srli_si128(va_lo, 8)); - __m128i b0 = _mm_cvtepi8_epi16(vb_lo); - __m128i b1 = _mm_cvtepi8_epi16(_mm_srli_si128(vb_lo, 8)); - - __m128i a2 = _mm_cvtepi8_epi16(va_hi); - __m128i a3 = _mm_cvtepi8_epi16(_mm_srli_si128(va_hi, 8)); - __m128i b2 = _mm_cvtepi8_epi16(vb_hi); - __m128i b3 = _mm_cvtepi8_epi16(_mm_srli_si128(vb_hi, 8)); - - __m128i d0 = _mm_abs_epi16(_mm_sub_epi16(a0, b0)); - __m128i d1 = _mm_abs_epi16(_mm_sub_epi16(a1, b1)); - __m128i d2 = _mm_abs_epi16(_mm_sub_epi16(a2, b2)); - __m128i d3 = _mm_abs_epi16(_mm_sub_epi16(a3, b3)); - - __m256i w0 = _mm256_cvtepu16_epi32(d0); - __m256i w1 = _mm256_cvtepu16_epi32(d1); - __m256i w2 = _mm256_cvtepu16_epi32(d2); - __m256i w3 = _mm256_cvtepu16_epi32(d3); - - acc = _mm256_add_epi32(acc, w0); - acc = _mm256_add_epi32(acc, w1); - acc = _mm256_add_epi32(acc, w2); - acc = _mm256_add_epi32(acc, w3); + acc0 = _mm256_add_epi64(acc0, _mm256_sad_epu8(S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(a + i))), + S8_TO_BIASED_U8_AVX2(_mm256_loadu_si256((const __m256i *)(b + i))))); } - uint32_t temp[8]; - _mm256_storeu_si256((__m256i *)temp, acc); - int32_t total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + uint64_t sum = hsum256_epi64(_mm256_add_epi64(acc0, acc1)); - for (; i < n; ++i) { - total += abs((int)a[i] - (int)b[i]); - } + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - return (float)total; + return (float)sum; } float int8_distance_cosine_avx2 (const void *a, const void *b, int n) { - float dot = -int8_distance_dot_avx2(a, b, n); - float norm_a = sqrtf(-int8_distance_dot_avx2(a, a, n)); - float norm_b = sqrtf(-int8_distance_dot_avx2(b, b, n)); + const int8_t *x = (const int8_t *)a; + const int8_t *y = (const int8_t *)b; - if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; + __m256i dacc = _mm256_setzero_si256(), aacc = dacc, bacc = dacc; + int i = 0; + + for (; i <= n - 32; i += 32) { + __m256i va = _mm256_loadu_si256((const __m256i *)(x + i)); + __m256i vb = _mm256_loadu_si256((const __m256i *)(y + i)); + dacc = _mm256_add_epi32(dacc, dot_epi8(va, vb)); + aacc = _mm256_add_epi32(aacc, dot_epi8(va, va)); + bacc = _mm256_add_epi32(bacc, dot_epi8(vb, vb)); + } + + int64_t dot = hsum256_epi32_signed(dacc); + int64_t norm_a = hsum256_epi32_signed(aacc); + int64_t norm_b = hsum256_epi32_signed(bacc); + + for (; i < n; ++i) { + int32_t p = x[i], q = y[i]; + dot += (int64_t)(p * q); + norm_a += (int64_t)(p * p); + norm_b += (int64_t)(q * q); + } + + if (norm_a == 0 || norm_b == 0) return 1.0f; - float cosine_similarity = dot / (norm_a * norm_b); + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; diff --git a/src/distance-avx512.c b/src/distance-avx512.c index aff8be3..f8e2e9e 100644 --- a/src/distance-avx512.c +++ b/src/distance-avx512.c @@ -32,12 +32,6 @@ static inline double hsum512d(__m512d v) { return _mm512_reduce_add_pd(v); } -// Helper: Horizontal sum for __m256i (used in int accumulators if we reduce to 256 first) -// But for AVX512 we usually reduce the full ZMM. -static inline uint32_t hsum512_epi32(__m512i v) { - return _mm512_reduce_add_epi32(v); -} - // per-block Inf mismatch test on 16 lanes (returns true if L1/L2 should be +Inf) static inline bool block_has_l2_inf_mismatch_16(const uint16_t* a, const uint16_t* b) { /* mismatch if (a_inf ^ b_inf) OR (both Inf and signs differ) */ @@ -626,59 +620,73 @@ float bfloat16_distance_cosine_avx512(const void* v1, const void* v2, int n) { // MARK: - UINT8 - +// Same shape as the AVX2 integer kernels, 64 bytes at a time: PSADBW for L1, and PMADDWD +// on widened bytes for squared differences and dot products. The previous versions +// widened every byte to 32 bits before multiplying, on one accumulator chain. +static inline __m512i abs_diff_epu8_512 (__m512i a, __m512i b) { + return _mm512_or_si512(_mm512_subs_epu8(a, b), _mm512_subs_epu8(b, a)); +} + +static inline __m512i sqdiff_epu8_512 (__m512i a, __m512i b) { + __m512i d = abs_diff_epu8_512(a, b); + __m512i lo = _mm512_unpacklo_epi8(d, _mm512_setzero_si512()); + __m512i hi = _mm512_unpackhi_epi8(d, _mm512_setzero_si512()); + return _mm512_add_epi32(_mm512_madd_epi16(lo, lo), _mm512_madd_epi16(hi, hi)); +} + +static inline __m512i dot_epu8_512 (__m512i a, __m512i b) { + const __m512i zero = _mm512_setzero_si512(); + __m512i al = _mm512_unpacklo_epi8(a, zero), ah = _mm512_unpackhi_epi8(a, zero); + __m512i bl = _mm512_unpacklo_epi8(b, zero), bh = _mm512_unpackhi_epi8(b, zero); + return _mm512_add_epi32(_mm512_madd_epi16(al, bl), _mm512_madd_epi16(ah, bh)); +} + +static inline __m512i dot_epi8_512 (__m512i a, __m512i b) { + __m512i al = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(a)); + __m512i ah = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(a, 1)); + __m512i bl = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(b)); + __m512i bh = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(b, 1)); + return _mm512_add_epi32(_mm512_madd_epi16(al, bl), _mm512_madd_epi16(ah, bh)); +} + +// widen to 64 bits before folding the lanes: each lane is a running total, and summing +// sixteen of them in 32 bits would cap the usable dimension +static inline uint64_t hsum512_epu32 (__m512i v) { + const __m512i zero = _mm512_setzero_si512(); + __m512i lo = _mm512_unpacklo_epi32(v, zero); + __m512i hi = _mm512_unpackhi_epi32(v, zero); + return (uint64_t)_mm512_reduce_add_epi64(_mm512_add_epi64(lo, hi)); +} + +static inline int64_t hsum512_epi32_signed (__m512i v) { + __m512i lo = _mm512_cvtepi32_epi64(_mm512_castsi512_si256(v)); + __m512i hi = _mm512_cvtepi32_epi64(_mm512_extracti64x4_epi64(v, 1)); + return _mm512_reduce_add_epi64(_mm512_add_epi64(lo, hi)); +} + static inline float uint8_distance_l2_impl_avx512(const void* v1, const void* v2, int n, bool use_sqrt) { const uint8_t* a = (const uint8_t*)v1; const uint8_t* b = (const uint8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; - // Process 64 elements at a time (64 bytes = 512 bits) + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi32(acc0, sqdiff_epu8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); + acc1 = _mm512_add_epi32(acc1, sqdiff_epu8_512(_mm512_loadu_si512((const void *)(a + i + 64)), + _mm512_loadu_si512((const void *)(b + i + 64)))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); - - // Split 64x u8 into 2x 32x u16 (Low 32 bytes and High 32 bytes of 512 register) - - // 1. Lower 32 bytes -> 32x u16 - __m256i va_half_lo = _mm512_castsi512_si256(va); - __m256i vb_half_lo = _mm512_castsi512_si256(vb); - __m512i va_16_lo = _mm512_cvtepu8_epi16(va_half_lo); - __m512i vb_16_lo = _mm512_cvtepu8_epi16(vb_half_lo); - - // 2. Upper 32 bytes -> 32x u16 - __m256i va_half_hi = _mm512_extracti64x4_epi64(va, 1); - __m256i vb_half_hi = _mm512_extracti64x4_epi64(vb, 1); - __m512i va_16_hi = _mm512_cvtepu8_epi16(va_half_hi); - __m512i vb_16_hi = _mm512_cvtepu8_epi16(vb_half_hi); - - // Compute diffs (16-bit) - __m512i d_lo = _mm512_sub_epi16(va_16_lo, vb_16_lo); - __m512i d_hi = _mm512_sub_epi16(va_16_hi, vb_16_hi); - - // Square diffs (16-bit result) - __m512i s_lo = _mm512_mullo_epi16(d_lo, d_lo); - __m512i s_hi = _mm512_mullo_epi16(d_hi, d_hi); - - // Widen to 32-bit and accumulate. - // Each 512-bit register of 16-bit ints splits into TWO 512-bit registers of 32-bit ints. - // s_lo splits into s_lo_0, s_lo_1 - - __m256i s_lo_half = _mm512_castsi512_si256(s_lo); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(s_lo_half)); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(s_lo, 1))); - - __m256i s_hi_half = _mm512_castsi512_si256(s_hi); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(s_hi_half)); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(s_hi, 1))); + acc0 = _mm512_add_epi32(acc0, sqdiff_epu8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); } - uint32_t total = hsum512_epi32(acc); + uint64_t total = hsum512_epu32(_mm512_add_epi32(acc0, acc1)); - // Tail loop for (; i < n; ++i) { int d = (int)a[i] - (int)b[i]; - total += d * d; + total += (uint64_t)(d * d); } return use_sqrt ? sqrtf((float)total) : (float)total; @@ -696,124 +704,116 @@ float uint8_distance_dot_avx512(const void* v1, const void* v2, int n) { const uint8_t* a = (const uint8_t*)v1; const uint8_t* b = (const uint8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi32(acc0, dot_epu8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); + acc1 = _mm512_add_epi32(acc1, dot_epu8_512(_mm512_loadu_si512((const void *)(a + i + 64)), + _mm512_loadu_si512((const void *)(b + i + 64)))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); - - __m512i va_16_lo = _mm512_cvtepu8_epi16(_mm512_castsi512_si256(va)); - __m512i vb_16_lo = _mm512_cvtepu8_epi16(_mm512_castsi512_si256(vb)); - __m512i va_16_hi = _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(va, 1)); - __m512i vb_16_hi = _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(vb, 1)); - - __m512i p_lo = _mm512_mullo_epi16(va_16_lo, vb_16_lo); - __m512i p_hi = _mm512_mullo_epi16(va_16_hi, vb_16_hi); - - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(p_lo))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(p_lo, 1))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(p_hi))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(p_hi, 1))); + acc0 = _mm512_add_epi32(acc0, dot_epu8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); } - uint32_t total = hsum512_epi32(acc); + uint64_t dot = hsum512_epu32(_mm512_add_epi32(acc0, acc1)); - for (; i < n; ++i) { - total += a[i] * b[i]; - } + for (; i < n; ++i) dot += (uint64_t)((uint32_t)a[i] * (uint32_t)b[i]); - return -(float)total; + return -(float)dot; } float uint8_distance_l1_avx512(const void* v1, const void* v2, int n) { const uint8_t* a = (const uint8_t*)v1; const uint8_t* b = (const uint8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi64(acc0, _mm512_sad_epu8(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); + acc1 = _mm512_add_epi64(acc1, _mm512_sad_epu8(_mm512_loadu_si512((const void *)(a + i + 64)), + _mm512_loadu_si512((const void *)(b + i + 64)))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); - - __m512i va_16_lo = _mm512_cvtepu8_epi16(_mm512_castsi512_si256(va)); - __m512i vb_16_lo = _mm512_cvtepu8_epi16(_mm512_castsi512_si256(vb)); - __m512i va_16_hi = _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(va, 1)); - __m512i vb_16_hi = _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(vb, 1)); - - // abs(a-b) in 16-bit - // Note: AVX512BW has _mm512_abs_epi16 - __m512i d_lo = _mm512_abs_epi16(_mm512_sub_epi16(va_16_lo, vb_16_lo)); - __m512i d_hi = _mm512_abs_epi16(_mm512_sub_epi16(va_16_hi, vb_16_hi)); - - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(d_lo))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(d_lo, 1))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(d_hi))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(d_hi, 1))); + acc0 = _mm512_add_epi64(acc0, _mm512_sad_epu8(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); } - uint32_t total = hsum512_epi32(acc); + uint64_t sum = (uint64_t)_mm512_reduce_add_epi64(_mm512_add_epi64(acc0, acc1)); - for (; i < n; ++i) { - total += abs((int)a[i] - (int)b[i]); - } + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - return (float)total; + return (float)sum; } float uint8_distance_cosine_avx512(const void* a, const void* b, int n) { - float dot = -uint8_distance_dot_avx512(a, b, n); - float norm_a = sqrtf(-uint8_distance_dot_avx512(a, a, n)); - float norm_b = sqrtf(-uint8_distance_dot_avx512(b, b, n)); + const uint8_t* x = (const uint8_t*)a; + const uint8_t* y = (const uint8_t*)b; - if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; + // one fused pass rather than three calls to the dot kernel + __m512i dacc = _mm512_setzero_si512(), aacc = dacc, bacc = dacc; + int i = 0; + + for (; i <= n - 64; i += 64) { + __m512i va = _mm512_loadu_si512((const void *)(x + i)); + __m512i vb = _mm512_loadu_si512((const void *)(y + i)); + dacc = _mm512_add_epi32(dacc, dot_epu8_512(va, vb)); + aacc = _mm512_add_epi32(aacc, dot_epu8_512(va, va)); + bacc = _mm512_add_epi32(bacc, dot_epu8_512(vb, vb)); + } + + uint64_t dot = hsum512_epu32(dacc); + uint64_t norm_a = hsum512_epu32(aacc); + uint64_t norm_b = hsum512_epu32(bacc); + + for (; i < n; ++i) { + uint32_t p = x[i], q = y[i]; + dot += (uint64_t)(p * q); + norm_a += (uint64_t)(p * p); + norm_b += (uint64_t)(q * q); + } - float cosine_similarity = dot / (norm_a * norm_b); + if (norm_a == 0 || norm_b == 0) return 1.0f; + + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; } - // MARK: - INT8 - +// biasing an int8 by 0x80 maps it onto uint8 without changing any difference between two +// elements, so L2 and L1 are the unsigned kernel plus one XOR per vector +#define S8_TO_BIASED_U8_512(_v) _mm512_xor_si512((_v), _mm512_set1_epi8((char)0x80)) + static inline float int8_distance_l2_impl_avx512(const void* v1, const void* v2, int n, bool use_sqrt) { const int8_t* a = (const int8_t*)v1; const int8_t* b = (const int8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi32(acc0, sqdiff_epu8_512(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i))))); + acc1 = _mm512_add_epi32(acc1, sqdiff_epu8_512(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i + 64))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i + 64))))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); - - // Sign extend int8 to int16. - // _mm512_cvtepi8_epi16 behaves exactly like cvtepu8 but for signed. - __m512i va_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(va)); - __m512i vb_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(vb)); - __m512i va_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(va, 1)); - __m512i vb_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(vb, 1)); - - __m512i d_lo = _mm512_sub_epi16(va_16_lo, vb_16_lo); - __m512i d_hi = _mm512_sub_epi16(va_16_hi, vb_16_hi); - - __m512i s_lo = _mm512_mullo_epi16(d_lo, d_lo); - __m512i s_hi = _mm512_mullo_epi16(d_hi, d_hi); - - // Sign extend 16 to 32 and add (results of square are positive, but keeping types consistent) - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_castsi512_si256(s_lo))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(s_lo, 1))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_castsi512_si256(s_hi))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(s_hi, 1))); + acc0 = _mm512_add_epi32(acc0, sqdiff_epu8_512(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i))))); } - uint32_t total = hsum512_epi32(acc); + uint64_t total = hsum512_epu32(_mm512_add_epi32(acc0, acc1)); for (; i < n; ++i) { int d = (int)a[i] - (int)b[i]; - total += d * d; + total += (uint64_t)(d * d); } return use_sqrt ? sqrtf((float)total) : (float)total; @@ -831,78 +831,81 @@ float int8_distance_dot_avx512(const void* v1, const void* v2, int n) { const int8_t* a = (const int8_t*)v1; const int8_t* b = (const int8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi32(acc0, dot_epi8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); + acc1 = _mm512_add_epi32(acc1, dot_epi8_512(_mm512_loadu_si512((const void *)(a + i + 64)), + _mm512_loadu_si512((const void *)(b + i + 64)))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); - - __m512i va_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(va)); - __m512i vb_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(vb)); - __m512i va_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(va, 1)); - __m512i vb_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(vb, 1)); - - __m512i p_lo = _mm512_mullo_epi16(va_16_lo, vb_16_lo); - __m512i p_hi = _mm512_mullo_epi16(va_16_hi, vb_16_hi); - - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_castsi512_si256(p_lo))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(p_lo, 1))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_castsi512_si256(p_hi))); - acc = _mm512_add_epi32(acc, _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(p_hi, 1))); + acc0 = _mm512_add_epi32(acc0, dot_epi8_512(_mm512_loadu_si512((const void *)(a + i)), + _mm512_loadu_si512((const void *)(b + i)))); } - int32_t total = (int32_t)hsum512_epi32(acc); + int64_t dot = hsum512_epi32_signed(_mm512_add_epi32(acc0, acc1)); - for (; i < n; ++i) { - total += (int)a[i] * (int)b[i]; - } + for (; i < n; ++i) dot += (int64_t)((int32_t)a[i] * (int32_t)b[i]); - return -(float)total; + return -(float)dot; } float int8_distance_l1_avx512(const void* v1, const void* v2, int n) { const int8_t* a = (const int8_t*)v1; const int8_t* b = (const int8_t*)v2; - __m512i acc = _mm512_setzero_si512(); + __m512i acc0 = _mm512_setzero_si512(), acc1 = acc0; int i = 0; + for (; i <= n - 128; i += 128) { + acc0 = _mm512_add_epi64(acc0, _mm512_sad_epu8(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i))))); + acc1 = _mm512_add_epi64(acc1, _mm512_sad_epu8(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i + 64))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i + 64))))); + } for (; i <= n - 64; i += 64) { - __m512i va = _mm512_loadu_si512((const void*)(a + i)); - __m512i vb = _mm512_loadu_si512((const void*)(b + i)); + acc0 = _mm512_add_epi64(acc0, _mm512_sad_epu8(S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(a + i))), + S8_TO_BIASED_U8_512(_mm512_loadu_si512((const void *)(b + i))))); + } - __m512i va_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(va)); - __m512i vb_16_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(vb)); - __m512i va_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(va, 1)); - __m512i vb_16_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(vb, 1)); + uint64_t sum = (uint64_t)_mm512_reduce_add_epi64(_mm512_add_epi64(acc0, acc1)); - __m512i d_lo = _mm512_abs_epi16(_mm512_sub_epi16(va_16_lo, vb_16_lo)); - __m512i d_hi = _mm512_abs_epi16(_mm512_sub_epi16(va_16_hi, vb_16_hi)); + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(d_lo))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(d_lo, 1))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(d_hi))); - acc = _mm512_add_epi32(acc, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(d_hi, 1))); - } + return (float)sum; +} - int32_t total = (int32_t)hsum512_epi32(acc); +float int8_distance_cosine_avx512(const void* a, const void* b, int n) { + const int8_t* x = (const int8_t*)a; + const int8_t* y = (const int8_t*)b; - for (; i < n; ++i) { - total += abs((int)a[i] - (int)b[i]); + __m512i dacc = _mm512_setzero_si512(), aacc = dacc, bacc = dacc; + int i = 0; + + for (; i <= n - 64; i += 64) { + __m512i va = _mm512_loadu_si512((const void *)(x + i)); + __m512i vb = _mm512_loadu_si512((const void *)(y + i)); + dacc = _mm512_add_epi32(dacc, dot_epi8_512(va, vb)); + aacc = _mm512_add_epi32(aacc, dot_epi8_512(va, va)); + bacc = _mm512_add_epi32(bacc, dot_epi8_512(vb, vb)); } - return (float)total; -} + int64_t dot = hsum512_epi32_signed(dacc); + int64_t norm_a = hsum512_epi32_signed(aacc); + int64_t norm_b = hsum512_epi32_signed(bacc); -float int8_distance_cosine_avx512(const void* a, const void* b, int n) { - float dot = -int8_distance_dot_avx512(a, b, n); - float norm_a = sqrtf(-int8_distance_dot_avx512(a, a, n)); - float norm_b = sqrtf(-int8_distance_dot_avx512(b, b, n)); + for (; i < n; ++i) { + int32_t p = x[i], q = y[i]; + dot += (int64_t)(p * q); + norm_a += (int64_t)(p * p); + norm_b += (int64_t)(q * q); + } - if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; + if (norm_a == 0 || norm_b == 0) return 1.0f; - float cosine_similarity = dot / (norm_a * norm_b); + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; diff --git a/src/distance-neon.c b/src/distance-neon.c index 3298146..74caed0 100644 --- a/src/distance-neon.c +++ b/src/distance-neon.c @@ -854,50 +854,48 @@ float float16_distance_l1_neon (const void *v1, const void *v2, int n) { // MARK: - UINT8 - +// The integer kernels used to widen every byte to 32 bits and multiply there - twelve or +// more instructions per 16 bytes - on a single accumulator chain. NEON has the whole job +// in five: absolute difference, widening multiply, pairwise-accumulate. Products of two +// bytes fit u16, so the only rule is to widen into the u32 accumulator before a second +// product can overflow the u16 lane. +// +// The signed kernels reuse the unsigned ones wherever the arithmetic allows: biasing an +// int8 by 0x80 maps it onto uint8 without changing any difference between two elements, +// so L2 and L1 are the same computation. Dot and cosine need the true signed values. +#define S8_TO_BIASED_U8(_v) veorq_u8(vreinterpretq_u8_s8(_v), vdupq_n_u8(0x80)) + static inline float uint8_distance_l2_impl_neon(const void *v1, const void *v2, int n, bool use_sqrt) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - uint32x4_t acc = vmovq_n_u32(0); + uint32x4_t acc0 = vmovq_n_u32(0), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + uint8x16_t d0 = vabdq_u8(vld1q_u8(a + i ), vld1q_u8(b + i )); + uint8x16_t d1 = vabdq_u8(vld1q_u8(a + i + 16), vld1q_u8(b + i + 16)); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(d0), vget_low_u8(d0))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(d0), vget_high_u8(d0))); + acc2 = vpadalq_u16(acc2, vmull_u8(vget_low_u8(d1), vget_low_u8(d1))); + acc3 = vpadalq_u16(acc3, vmull_u8(vget_high_u8(d1), vget_high_u8(d1))); + } for (; i <= n - 16; i += 16) { - uint8x16_t va = vld1q_u8(a + i); - uint8x16_t vb = vld1q_u8(b + i); - - // compute 8-bit differences widened to signed 16-bit - int16x8_t diff_lo = (int16x8_t)vsubl_u8(vget_low_u8(va), vget_low_u8(vb)); - int16x8_t diff_hi = (int16x8_t)vsubl_u8(vget_high_u8(va), vget_high_u8(vb)); - - // widen to signed 32-bit and square - int32x4_t diff_lo_0 = vmovl_s16(vget_low_s16(diff_lo)); - int32x4_t diff_lo_1 = vmovl_s16(vget_high_s16(diff_lo)); - int32x4_t diff_hi_0 = vmovl_s16(vget_low_s16(diff_hi)); - int32x4_t diff_hi_1 = vmovl_s16(vget_high_s16(diff_hi)); - - diff_lo_0 = vmulq_s32(diff_lo_0, diff_lo_0); - diff_lo_1 = vmulq_s32(diff_lo_1, diff_lo_1); - diff_hi_0 = vmulq_s32(diff_hi_0, diff_hi_0); - diff_hi_1 = vmulq_s32(diff_hi_1, diff_hi_1); - - // accumulate into uint32_t accumulator - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_lo_0)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_lo_1)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_hi_0)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_hi_1)); + uint8x16_t d = vabdq_u8(vld1q_u8(a + i), vld1q_u8(b + i)); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(d), vget_low_u8(d))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(d), vget_high_u8(d))); } - // horizontal sum + uint32x4_t acc = vaddq_u32(vaddq_u32(acc0, acc1), vaddq_u32(acc2, acc3)); uint64x2_t sum64 = vpaddlq_u32(acc); - uint64_t final_sum = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); + uint64_t total = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); - // tail for (; i < n; ++i) { int diff = (int)a[i] - (int)b[i]; - final_sum += (uint64_t)(diff * diff); + total += (uint64_t)(diff * diff); } - return use_sqrt ? sqrtf((float)final_sum) : (float)final_sum; + return use_sqrt ? sqrtf((float)total) : (float)total; } float uint8_distance_l2_neon (const void *v1, const void *v2, int n) { @@ -908,81 +906,53 @@ float uint8_distance_l2_squared_neon (const void *v1, const void *v2, int n) { return uint8_distance_l2_impl_neon(v1, v2, n, false); } -float uint8_distance_cosine_neon (const void *v1, const void *v2, int n) { - const uint8_t *a = (const uint8_t *)v1; - const uint8_t *b = (const uint8_t *)v2; - - uint32x4_t dot_acc = vmovq_n_u32(0); - uint32x4_t norm_a_acc = vmovq_n_u32(0); - uint32x4_t norm_b_acc = vmovq_n_u32(0); - +// dot, |a|^2 and |b|^2 in one pass, two accumulators each +static inline void uint8_sums_neon (const uint8_t *a, const uint8_t *b, int n, uint64_t *dot_out, uint64_t *na_out, uint64_t *nb_out) { + uint32x4_t dot0 = vmovq_n_u32(0), dot1 = dot0; + uint32x4_t na0 = dot0, na1 = dot0; + uint32x4_t nb0 = dot0, nb1 = dot0; int i = 0; + for (; i <= n - 16; i += 16) { - // Load 16 bytes from each vector - uint8x16_t va_u8 = vld1q_u8(a + i); - uint8x16_t vb_u8 = vld1q_u8(b + i); - - // Convert to uint16x8_t - uint16x8_t va_lo_u16 = vmovl_u8(vget_low_u8(va_u8)); - uint16x8_t va_hi_u16 = vmovl_u8(vget_high_u8(va_u8)); - uint16x8_t vb_lo_u16 = vmovl_u8(vget_low_u8(vb_u8)); - uint16x8_t vb_hi_u16 = vmovl_u8(vget_high_u8(vb_u8)); - - // Multiply for dot product - uint32x4_t dot_lo = vmull_u16(vget_low_u16(va_lo_u16), vget_low_u16(vb_lo_u16)); - uint32x4_t dot_hi = vmull_u16(vget_high_u16(va_lo_u16), vget_high_u16(vb_lo_u16)); - uint32x4_t dot_lo2 = vmull_u16(vget_low_u16(va_hi_u16), vget_low_u16(vb_hi_u16)); - uint32x4_t dot_hi2 = vmull_u16(vget_high_u16(va_hi_u16), vget_high_u16(vb_hi_u16)); - - // Multiply for norms - uint32x4_t a2_lo = vmull_u16(vget_low_u16(va_lo_u16), vget_low_u16(va_lo_u16)); - uint32x4_t a2_hi = vmull_u16(vget_high_u16(va_lo_u16), vget_high_u16(va_lo_u16)); - uint32x4_t a2_lo2 = vmull_u16(vget_low_u16(va_hi_u16), vget_low_u16(va_hi_u16)); - uint32x4_t a2_hi2 = vmull_u16(vget_high_u16(va_hi_u16), vget_high_u16(va_hi_u16)); - - uint32x4_t b2_lo = vmull_u16(vget_low_u16(vb_lo_u16), vget_low_u16(vb_lo_u16)); - uint32x4_t b2_hi = vmull_u16(vget_high_u16(vb_lo_u16), vget_high_u16(vb_lo_u16)); - uint32x4_t b2_lo2 = vmull_u16(vget_low_u16(vb_hi_u16), vget_low_u16(vb_hi_u16)); - uint32x4_t b2_hi2 = vmull_u16(vget_high_u16(vb_hi_u16), vget_high_u16(vb_hi_u16)); - - // Accumulate - dot_acc = vaddq_u32(dot_acc, dot_lo); - dot_acc = vaddq_u32(dot_acc, dot_hi); - dot_acc = vaddq_u32(dot_acc, dot_lo2); - dot_acc = vaddq_u32(dot_acc, dot_hi2); - - norm_a_acc = vaddq_u32(norm_a_acc, a2_lo); - norm_a_acc = vaddq_u32(norm_a_acc, a2_hi); - norm_a_acc = vaddq_u32(norm_a_acc, a2_lo2); - norm_a_acc = vaddq_u32(norm_a_acc, a2_hi2); - - norm_b_acc = vaddq_u32(norm_b_acc, b2_lo); - norm_b_acc = vaddq_u32(norm_b_acc, b2_hi); - norm_b_acc = vaddq_u32(norm_b_acc, b2_lo2); - norm_b_acc = vaddq_u32(norm_b_acc, b2_hi2); + uint8x16_t va = vld1q_u8(a + i); + uint8x16_t vb = vld1q_u8(b + i); + uint8x8_t al = vget_low_u8(va), ah = vget_high_u8(va); + uint8x8_t bl = vget_low_u8(vb), bh = vget_high_u8(vb); + + dot0 = vpadalq_u16(dot0, vmull_u8(al, bl)); + dot1 = vpadalq_u16(dot1, vmull_u8(ah, bh)); + na0 = vpadalq_u16(na0, vmull_u8(al, al)); + na1 = vpadalq_u16(na1, vmull_u8(ah, ah)); + nb0 = vpadalq_u16(nb0, vmull_u8(bl, bl)); + nb1 = vpadalq_u16(nb1, vmull_u8(bh, bh)); } - - // Horizontal sum - uint32_t dot = vgetq_lane_u32(dot_acc, 0) + vgetq_lane_u32(dot_acc, 1) + - vgetq_lane_u32(dot_acc, 2) + vgetq_lane_u32(dot_acc, 3); - - uint32_t norm_a = vgetq_lane_u32(norm_a_acc, 0) + vgetq_lane_u32(norm_a_acc, 1) + - vgetq_lane_u32(norm_a_acc, 2) + vgetq_lane_u32(norm_a_acc, 3); - - uint32_t norm_b = vgetq_lane_u32(norm_b_acc, 0) + vgetq_lane_u32(norm_b_acc, 1) + - vgetq_lane_u32(norm_b_acc, 2) + vgetq_lane_u32(norm_b_acc, 3); - - // Tail loop + + uint64x2_t d64 = vpaddlq_u32(vaddq_u32(dot0, dot1)); + uint64x2_t a64 = vpaddlq_u32(vaddq_u32(na0, na1)); + uint64x2_t b64 = vpaddlq_u32(vaddq_u32(nb0, nb1)); + uint64_t dot = vgetq_lane_u64(d64, 0) + vgetq_lane_u64(d64, 1); + uint64_t na = vgetq_lane_u64(a64, 0) + vgetq_lane_u64(a64, 1); + uint64_t nb = vgetq_lane_u64(b64, 0) + vgetq_lane_u64(b64, 1); + for (; i < n; ++i) { - int ai = a[i]; - int bi = b[i]; - dot += ai * bi; - norm_a += ai * ai; - norm_b += bi * bi; + uint32_t x = a[i], y = b[i]; + dot += (uint64_t)(x * y); + na += (uint64_t)(x * x); + nb += (uint64_t)(y * y); } - + + *dot_out = dot; + *na_out = na; + *nb_out = nb; +} + +float uint8_distance_cosine_neon (const void *v1, const void *v2, int n) { + uint64_t dot, norm_a, norm_b; + uint8_sums_neon((const uint8_t *)v1, (const uint8_t *)v2, n, &dot, &norm_a, &norm_b); + if (norm_a == 0 || norm_b == 0) return 1.0f; - float cosine_similarity = dot / (sqrtf((float)norm_a) * sqrtf((float)norm_b)); + + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; @@ -991,81 +961,56 @@ float uint8_distance_cosine_neon (const void *v1, const void *v2, int n) { float uint8_distance_dot_neon (const void *v1, const void *v2, int n) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - - uint32x4_t dot_acc = vmovq_n_u32(0); // 4-lane accumulator + + uint32x4_t acc0 = vmovq_n_u32(0), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; - - for (; i <= n - 16; i += 16) { - uint8x16_t va_u8 = vld1q_u8(a + i); - uint8x16_t vb_u8 = vld1q_u8(b + i); - - // Widen to 16-bit - uint16x8_t va_lo = vmovl_u8(vget_low_u8(va_u8)); - uint16x8_t vb_lo = vmovl_u8(vget_low_u8(vb_u8)); - uint16x8_t va_hi = vmovl_u8(vget_high_u8(va_u8)); - uint16x8_t vb_hi = vmovl_u8(vget_high_u8(vb_u8)); - - // Multiply low and high halves - uint32x4_t dot_lo = vmull_u16(vget_low_u16(va_lo), vget_low_u16(vb_lo)); - uint32x4_t dot_hi = vmull_u16(vget_high_u16(va_lo), vget_high_u16(vb_lo)); - uint32x4_t dot_lo2 = vmull_u16(vget_low_u16(va_hi), vget_low_u16(vb_hi)); - uint32x4_t dot_hi2 = vmull_u16(vget_high_u16(va_hi), vget_high_u16(vb_hi)); - - // Accumulate - dot_acc = vaddq_u32(dot_acc, dot_lo); - dot_acc = vaddq_u32(dot_acc, dot_hi); - dot_acc = vaddq_u32(dot_acc, dot_lo2); - dot_acc = vaddq_u32(dot_acc, dot_hi2); + + for (; i <= n - 32; i += 32) { + uint8x16_t a0 = vld1q_u8(a + i ), b0 = vld1q_u8(b + i ); + uint8x16_t a1 = vld1q_u8(a + i + 16), b1 = vld1q_u8(b + i + 16); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(a0), vget_low_u8(b0))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(a0), vget_high_u8(b0))); + acc2 = vpadalq_u16(acc2, vmull_u8(vget_low_u8(a1), vget_low_u8(b1))); + acc3 = vpadalq_u16(acc3, vmull_u8(vget_high_u8(a1), vget_high_u8(b1))); } - - // Horizontal add of 4 lanes - uint32_t dot = vgetq_lane_u32(dot_acc, 0) + - vgetq_lane_u32(dot_acc, 1) + - vgetq_lane_u32(dot_acc, 2) + - vgetq_lane_u32(dot_acc, 3); - - // Tail loop - for (; i < n; ++i) { - dot += a[i] * b[i]; + for (; i <= n - 16; i += 16) { + uint8x16_t va = vld1q_u8(a + i), vb = vld1q_u8(b + i); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(va), vget_low_u8(vb))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(va), vget_high_u8(vb))); } - - return -(float)dot; // negative dot product = dot distance + + uint64x2_t sum64 = vpaddlq_u32(vaddq_u32(vaddq_u32(acc0, acc1), vaddq_u32(acc2, acc3))); + uint64_t dot = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); + + for (; i < n; ++i) dot += (uint64_t)((uint32_t)a[i] * (uint32_t)b[i]); + + return -(float)dot; } float uint8_distance_l1_neon (const void *v1, const void *v2, int n) { const uint8_t *a = (const uint8_t *)v1; const uint8_t *b = (const uint8_t *)v2; - uint32x4_t sum_acc = vdupq_n_u32(0); + uint32x4_t acc0 = vmovq_n_u32(0), acc1 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + uint8x16_t d0 = vabdq_u8(vld1q_u8(a + i ), vld1q_u8(b + i )); + uint8x16_t d1 = vabdq_u8(vld1q_u8(a + i + 16), vld1q_u8(b + i + 16)); + acc0 = vpadalq_u16(acc0, vpaddlq_u8(d0)); + acc1 = vpadalq_u16(acc1, vpaddlq_u8(d1)); + } for (; i <= n - 16; i += 16) { - uint8x16_t va = vld1q_u8(a + i); - uint8x16_t vb = vld1q_u8(b + i); + uint8x16_t d = vabdq_u8(vld1q_u8(a + i), vld1q_u8(b + i)); + acc0 = vpadalq_u16(acc0, vpaddlq_u8(d)); + } - // Compute absolute difference - uint8x16_t abs_diff = vabdq_u8(va, vb); + uint64x2_t sum64 = vpaddlq_u32(vaddq_u32(acc0, acc1)); + uint64_t sum = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); - // Widen to 16-bit then accumulate into 32-bit - uint16x8_t abs_lo = vmovl_u8(vget_low_u8(abs_diff)); - uint16x8_t abs_hi = vmovl_u8(vget_high_u8(abs_diff)); + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - sum_acc = vaddq_u32(sum_acc, vmovl_u16(vget_low_u16(abs_lo))); - sum_acc = vaddq_u32(sum_acc, vmovl_u16(vget_high_u16(abs_lo))); - sum_acc = vaddq_u32(sum_acc, vmovl_u16(vget_low_u16(abs_hi))); - sum_acc = vaddq_u32(sum_acc, vmovl_u16(vget_high_u16(abs_hi))); - } - - // Horizontal sum - uint64x2_t sum64 = vpaddlq_u32(sum_acc); - uint32_t total = (uint32_t)(vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1)); - - // Tail loop - for (; i < n; ++i) { - total += (uint32_t)abs((int)a[i] - (int)b[i]); - } - - return (float)total; + return (float)sum; } // MARK: - INT8 - @@ -1074,46 +1019,38 @@ static inline float int8_distance_l2_neon_imp (const void *v1, const void *v2, i const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - uint32x4_t acc = vmovq_n_u32(0); + // biasing both sides by 0x80 leaves every |a[i] - b[i]| untouched, so this is the + // unsigned kernel with two extra XORs per vector + uint32x4_t acc0 = vmovq_n_u32(0), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + uint8x16_t a0 = S8_TO_BIASED_U8(vld1q_s8(a + i )); + uint8x16_t b0 = S8_TO_BIASED_U8(vld1q_s8(b + i )); + uint8x16_t a1 = S8_TO_BIASED_U8(vld1q_s8(a + i + 16)); + uint8x16_t b1 = S8_TO_BIASED_U8(vld1q_s8(b + i + 16)); + uint8x16_t d0 = vabdq_u8(a0, b0); + uint8x16_t d1 = vabdq_u8(a1, b1); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(d0), vget_low_u8(d0))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(d0), vget_high_u8(d0))); + acc2 = vpadalq_u16(acc2, vmull_u8(vget_low_u8(d1), vget_low_u8(d1))); + acc3 = vpadalq_u16(acc3, vmull_u8(vget_high_u8(d1), vget_high_u8(d1))); + } for (; i <= n - 16; i += 16) { - int8x16_t va = vld1q_s8(a + i); - int8x16_t vb = vld1q_s8(b + i); - - // signed widening subtraction: int8 → int16 - int16x8_t diff_lo = vsubl_s8(vget_low_s8(va), vget_low_s8(vb)); - int16x8_t diff_hi = vsubl_s8(vget_high_s8(va), vget_high_s8(vb)); - - // widen to int32 and square - int32x4_t diff_lo_0 = vmovl_s16(vget_low_s16(diff_lo)); - int32x4_t diff_lo_1 = vmovl_s16(vget_high_s16(diff_lo)); - int32x4_t diff_hi_0 = vmovl_s16(vget_low_s16(diff_hi)); - int32x4_t diff_hi_1 = vmovl_s16(vget_high_s16(diff_hi)); - - diff_lo_0 = vmulq_s32(diff_lo_0, diff_lo_0); - diff_lo_1 = vmulq_s32(diff_lo_1, diff_lo_1); - diff_hi_0 = vmulq_s32(diff_hi_0, diff_hi_0); - diff_hi_1 = vmulq_s32(diff_hi_1, diff_hi_1); - - // accumulate, cast to uint32 to match accumulator type - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_lo_0)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_lo_1)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_hi_0)); - acc = vaddq_u32(acc, vreinterpretq_u32_s32(diff_hi_1)); + uint8x16_t d = vabdq_u8(S8_TO_BIASED_U8(vld1q_s8(a + i)), S8_TO_BIASED_U8(vld1q_s8(b + i))); + acc0 = vpadalq_u16(acc0, vmull_u8(vget_low_u8(d), vget_low_u8(d))); + acc1 = vpadalq_u16(acc1, vmull_u8(vget_high_u8(d), vget_high_u8(d))); } - // horizontal sum - uint64x2_t sum64 = vpaddlq_u32(acc); - uint64_t final_sum = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); + uint64x2_t sum64 = vpaddlq_u32(vaddq_u32(vaddq_u32(acc0, acc1), vaddq_u32(acc2, acc3))); + uint64_t total = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); - // tail for (; i < n; ++i) { int diff = (int)a[i] - (int)b[i]; - final_sum += (uint64_t)(diff * diff); + total += (uint64_t)(diff * diff); } - return use_sqrt ? sqrtf((float)final_sum) : (float)final_sum; + return use_sqrt ? sqrtf((float)total) : (float)total; } float int8_distance_l2_neon (const void *v1, const void *v2, int n) { @@ -1124,75 +1061,53 @@ float int8_distance_l2_squared_neon (const void *v1, const void *v2, int n) { return int8_distance_l2_neon_imp(v1, v2, n, false); } -float int8_distance_cosine_neon (const void *v1, const void *v2, int n) { - const int8_t *a = (const int8_t *)v1; - const int8_t *b = (const int8_t *)v2; - - int32x4_t acc_dot = vdupq_n_s32(0); - int32x4_t acc_a2 = vdupq_n_s32(0); - int32x4_t acc_b2 = vdupq_n_s32(0); +// signed products need the real values, so no biasing here +static inline void int8_sums_neon (const int8_t *a, const int8_t *b, int n, int64_t *dot_out, int64_t *na_out, int64_t *nb_out) { + int32x4_t dot0 = vmovq_n_s32(0), dot1 = dot0; + int32x4_t na0 = dot0, na1 = dot0; + int32x4_t nb0 = dot0, nb1 = dot0; int i = 0; for (; i <= n - 16; i += 16) { int8x16_t va = vld1q_s8(a + i); int8x16_t vb = vld1q_s8(b + i); - - int16x8_t lo_a = vmovl_s8(vget_low_s8(va)); - int16x8_t hi_a = vmovl_s8(vget_high_s8(va)); - int16x8_t lo_b = vmovl_s8(vget_low_s8(vb)); - int16x8_t hi_b = vmovl_s8(vget_high_s8(vb)); - - // Dot product - int32x4_t dot_lo = vmull_s16(vget_low_s16(lo_a), vget_low_s16(lo_b)); - int32x4_t dot_hi = vmull_s16(vget_high_s16(lo_a), vget_high_s16(lo_b)); - int32x4_t dot_lo2 = vmull_s16(vget_low_s16(hi_a), vget_low_s16(hi_b)); - int32x4_t dot_hi2 = vmull_s16(vget_high_s16(hi_a), vget_high_s16(hi_b)); - - acc_dot = vaddq_s32(acc_dot, dot_lo); - acc_dot = vaddq_s32(acc_dot, dot_hi); - acc_dot = vaddq_s32(acc_dot, dot_lo2); - acc_dot = vaddq_s32(acc_dot, dot_hi2); - - // Norm a² - int32x4_t a2_lo = vmull_s16(vget_low_s16(lo_a), vget_low_s16(lo_a)); - int32x4_t a2_hi = vmull_s16(vget_high_s16(lo_a), vget_high_s16(lo_a)); - int32x4_t a2_lo2 = vmull_s16(vget_low_s16(hi_a), vget_low_s16(hi_a)); - int32x4_t a2_hi2 = vmull_s16(vget_high_s16(hi_a), vget_high_s16(hi_a)); - - acc_a2 = vaddq_s32(acc_a2, a2_lo); - acc_a2 = vaddq_s32(acc_a2, a2_hi); - acc_a2 = vaddq_s32(acc_a2, a2_lo2); - acc_a2 = vaddq_s32(acc_a2, a2_hi2); - - // Norm b² - int32x4_t b2_lo = vmull_s16(vget_low_s16(lo_b), vget_low_s16(lo_b)); - int32x4_t b2_hi = vmull_s16(vget_high_s16(lo_b), vget_high_s16(lo_b)); - int32x4_t b2_lo2 = vmull_s16(vget_low_s16(hi_b), vget_low_s16(hi_b)); - int32x4_t b2_hi2 = vmull_s16(vget_high_s16(hi_b), vget_high_s16(hi_b)); - - acc_b2 = vaddq_s32(acc_b2, b2_lo); - acc_b2 = vaddq_s32(acc_b2, b2_hi); - acc_b2 = vaddq_s32(acc_b2, b2_lo2); - acc_b2 = vaddq_s32(acc_b2, b2_hi2); + int8x8_t al = vget_low_s8(va), ah = vget_high_s8(va); + int8x8_t bl = vget_low_s8(vb), bh = vget_high_s8(vb); + + dot0 = vpadalq_s16(dot0, vmull_s8(al, bl)); + dot1 = vpadalq_s16(dot1, vmull_s8(ah, bh)); + na0 = vpadalq_s16(na0, vmull_s8(al, al)); + na1 = vpadalq_s16(na1, vmull_s8(ah, ah)); + nb0 = vpadalq_s16(nb0, vmull_s8(bl, bl)); + nb1 = vpadalq_s16(nb1, vmull_s8(bh, bh)); } - int32_t dot = vgetq_lane_s32(acc_dot, 0) + vgetq_lane_s32(acc_dot, 1) - + vgetq_lane_s32(acc_dot, 2) + vgetq_lane_s32(acc_dot, 3); - int32_t norm_a = vgetq_lane_s32(acc_a2, 0) + vgetq_lane_s32(acc_a2, 1) - + vgetq_lane_s32(acc_a2, 2) + vgetq_lane_s32(acc_a2, 3); - int32_t norm_b = vgetq_lane_s32(acc_b2, 0) + vgetq_lane_s32(acc_b2, 1) - + vgetq_lane_s32(acc_b2, 2) + vgetq_lane_s32(acc_b2, 3); + int64x2_t d64 = vpaddlq_s32(vaddq_s32(dot0, dot1)); + int64x2_t a64 = vpaddlq_s32(vaddq_s32(na0, na1)); + int64x2_t b64 = vpaddlq_s32(vaddq_s32(nb0, nb1)); + int64_t dot = vgetq_lane_s64(d64, 0) + vgetq_lane_s64(d64, 1); + int64_t na = vgetq_lane_s64(a64, 0) + vgetq_lane_s64(a64, 1); + int64_t nb = vgetq_lane_s64(b64, 0) + vgetq_lane_s64(b64, 1); for (; i < n; ++i) { - int ai = a[i]; - int bi = b[i]; - dot += ai * bi; - norm_a += ai * ai; - norm_b += bi * bi; + int32_t x = a[i], y = b[i]; + dot += (int64_t)(x * y); + na += (int64_t)(x * x); + nb += (int64_t)(y * y); } + *dot_out = dot; + *na_out = na; + *nb_out = nb; +} + +float int8_distance_cosine_neon (const void *v1, const void *v2, int n) { + int64_t dot, norm_a, norm_b; + int8_sums_neon((const int8_t *)v1, (const int8_t *)v2, n, &dot, &norm_a, &norm_b); + if (norm_a == 0 || norm_b == 0) return 1.0f; - float cosine_similarity = dot / (sqrtf((float)norm_a) * sqrtf((float)norm_b)); + + float cosine_similarity = (float)((double)dot / (sqrt((double)norm_a) * sqrt((double)norm_b))); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; @@ -1201,76 +1116,57 @@ float int8_distance_cosine_neon (const void *v1, const void *v2, int n) { float int8_distance_dot_neon (const void *v1, const void *v2, int n) { const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - - int32x4_t acc = vdupq_n_s32(0); + + int32x4_t acc0 = vmovq_n_s32(0), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + int8x16_t a0 = vld1q_s8(a + i ), b0 = vld1q_s8(b + i ); + int8x16_t a1 = vld1q_s8(a + i + 16), b1 = vld1q_s8(b + i + 16); + acc0 = vpadalq_s16(acc0, vmull_s8(vget_low_s8(a0), vget_low_s8(b0))); + acc1 = vpadalq_s16(acc1, vmull_s8(vget_high_s8(a0), vget_high_s8(b0))); + acc2 = vpadalq_s16(acc2, vmull_s8(vget_low_s8(a1), vget_low_s8(b1))); + acc3 = vpadalq_s16(acc3, vmull_s8(vget_high_s8(a1), vget_high_s8(b1))); + } for (; i <= n - 16; i += 16) { - int8x16_t va = vld1q_s8(a + i); - int8x16_t vb = vld1q_s8(b + i); - - int16x8_t lo_a = vmovl_s8(vget_low_s8(va)); - int16x8_t hi_a = vmovl_s8(vget_high_s8(va)); - int16x8_t lo_b = vmovl_s8(vget_low_s8(vb)); - int16x8_t hi_b = vmovl_s8(vget_high_s8(vb)); - - int32x4_t prod_lo = vmull_s16(vget_low_s16(lo_a), vget_low_s16(lo_b)); - int32x4_t prod_hi = vmull_s16(vget_high_s16(lo_a), vget_high_s16(lo_b)); - int32x4_t prod_lo2 = vmull_s16(vget_low_s16(hi_a), vget_low_s16(hi_b)); - int32x4_t prod_hi2 = vmull_s16(vget_high_s16(hi_a), vget_high_s16(hi_b)); - - acc = vaddq_s32(acc, prod_lo); - acc = vaddq_s32(acc, prod_hi); - acc = vaddq_s32(acc, prod_lo2); - acc = vaddq_s32(acc, prod_hi2); + int8x16_t va = vld1q_s8(a + i), vb = vld1q_s8(b + i); + acc0 = vpadalq_s16(acc0, vmull_s8(vget_low_s8(va), vget_low_s8(vb))); + acc1 = vpadalq_s16(acc1, vmull_s8(vget_high_s8(va), vget_high_s8(vb))); } - int32_t dot = vgetq_lane_s32(acc, 0) + vgetq_lane_s32(acc, 1) - + vgetq_lane_s32(acc, 2) + vgetq_lane_s32(acc, 3); + int64x2_t sum64 = vpaddlq_s32(vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3))); + int64_t dot = vgetq_lane_s64(sum64, 0) + vgetq_lane_s64(sum64, 1); - for (; i < n; ++i) { - dot += a[i] * b[i]; - } + for (; i < n; ++i) dot += (int64_t)((int32_t)a[i] * (int32_t)b[i]); - return -(float)dot; // negative dot product + return -(float)dot; } float int8_distance_l1_neon(const void *v1, const void *v2, int n) { const int8_t *a = (const int8_t *)v1; const int8_t *b = (const int8_t *)v2; - uint32x4_t acc = vdupq_n_u32(0); + // same biasing trick as L2: absolute differences are unaffected + uint32x4_t acc0 = vmovq_n_u32(0), acc1 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + uint8x16_t d0 = vabdq_u8(S8_TO_BIASED_U8(vld1q_s8(a + i )), S8_TO_BIASED_U8(vld1q_s8(b + i ))); + uint8x16_t d1 = vabdq_u8(S8_TO_BIASED_U8(vld1q_s8(a + i + 16)), S8_TO_BIASED_U8(vld1q_s8(b + i + 16))); + acc0 = vpadalq_u16(acc0, vpaddlq_u8(d0)); + acc1 = vpadalq_u16(acc1, vpaddlq_u8(d1)); + } for (; i <= n - 16; i += 16) { - int8x16_t va = vld1q_s8(a + i); - int8x16_t vb = vld1q_s8(b + i); - - // Widen to 16-bit signed - int16x8_t diff_lo = vsubl_s8(vget_low_s8(va), vget_low_s8(vb)); - int16x8_t diff_hi = vsubl_s8(vget_high_s8(va), vget_high_s8(vb)); - - // Absolute values (safe for -128) - int16x8_t abs_lo = vabsq_s16(diff_lo); - int16x8_t abs_hi = vabsq_s16(diff_hi); - - // Widen to 32-bit and accumulate - acc = vaddq_u32(acc, vmovl_u16(vget_low_u16(vreinterpretq_u16_s16(abs_lo)))); - acc = vaddq_u32(acc, vmovl_u16(vget_high_u16(vreinterpretq_u16_s16(abs_lo)))); - acc = vaddq_u32(acc, vmovl_u16(vget_low_u16(vreinterpretq_u16_s16(abs_hi)))); - acc = vaddq_u32(acc, vmovl_u16(vget_high_u16(vreinterpretq_u16_s16(abs_hi)))); + uint8x16_t d = vabdq_u8(S8_TO_BIASED_U8(vld1q_s8(a + i)), S8_TO_BIASED_U8(vld1q_s8(b + i))); + acc0 = vpadalq_u16(acc0, vpaddlq_u8(d)); } - // Horizontal sum - uint64x2_t sum64 = vpaddlq_u32(acc); - uint64_t final = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); + uint64x2_t sum64 = vpaddlq_u32(vaddq_u32(acc0, acc1)); + uint64_t sum = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1); - // Tail loop - for (; i < n; ++i) { - final += (uint32_t)abs((int)a[i] - (int)b[i]); - } + for (; i < n; ++i) sum += (uint64_t)abs((int)a[i] - (int)b[i]); - return (float)final; + return (float)sum; } // MARK: - BIT - diff --git a/src/sqlite-vector.c b/src/sqlite-vector.c index 81b9347..0f7c8c6 100644 --- a/src/sqlite-vector.c +++ b/src/sqlite-vector.c @@ -210,7 +210,6 @@ typedef struct { int64_t *rowids; double *distance; int size; - int max_index; int row_index; int row_count; } vFullScanCursor; @@ -3113,52 +3112,49 @@ static int vFullScanCursorRowid (sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) return SQLITE_OK; } -static inline int vFullScanFindMaxIndex (double *values, int n) { - int max_idx = 0; - if (n <= 32) { - // use simple version - for (int i = 1; i < n; ++i) { - if (values[i] > values[max_idx]) {max_idx = i;} - } - return max_idx; - } - - // use unrolled version - double max_val = values[0]; - int i = 1; - - // unroll loop in blocks of 4 - for (; i + 3 < n; i += 4) { - if (values[i] > max_val) {max_val = values[i]; max_idx = i;} - if (values[i + 1] > max_val) {max_val = values[i + 1]; max_idx = i + 1;} - if (values[i + 2] > max_val) {max_val = values[i + 2]; max_idx = i + 2;} - if (values[i + 3] > max_val) {max_val = values[i + 3]; max_idx = i + 3;} - } - - // process remaining elements - for (; i < n; ++i) { - if (values[i] > max_val) {max_val = values[i]; max_idx = i;} +// The candidate set is a binary max-heap over (distance, rowid), kept in the parallel +// arrays the cursor already owns. The root is the worst entry still in the set, so +// admitting a candidate costs one comparison and O(log k) to reinsert - the previous +// version rescanned all k slots for a new maximum on every improvement, and then paid a +// second O(k^2) to sort. It also removes a stale-index hazard: the old max_index was +// never reset between filters, so a cursor reused with a smaller k read past the end of +// the reallocated array. The root of a heap is always slot zero. +static inline void vTopKSiftDown (double *distance, int64_t *rowids, int n, int root) { + for (;;) { + int child = 2 * root + 1; + if (child >= n) break; + if (child + 1 < n && distance[child + 1] > distance[child]) ++child; + if (distance[child] <= distance[root]) break; + + SWAP(double, distance[root], distance[child]); + SWAP(int64_t, rowids[root], rowids[child]); + root = child; } - return max_idx; } +// evicts the worst entry in favour of a better one; callers check distance[0] first +static inline void vTopKReplaceWorst (vFullScanCursor *c, double distance, int64_t rowid) { + c->distance[0] = distance; + c->rowids[0] = rowid; + vTopKSiftDown(c->distance, c->rowids, c->row_count, 0); +} + +// Heapsort in place: repeatedly move the largest entry past the end of the heap, which +// leaves the array ascending - the order the cursor emits rows in. Returns how many +// trailing slots were never filled so the caller can trim them. static int vFullScanSortSlots (vFullScanCursor *c) { - int counter = 0; - int row_count = c->row_count; double *distance = c->distance; int64_t *rowids = c->rowids; - - for (int i = 0; i < row_count - 1; ++i) { - if (distance[i] == INFINITY) ++counter; - for (int j = i + 1; j < row_count; ++j) { - if (distance[j] < distance[i]) { - SWAP(double, distance[i], distance[j]); - SWAP(int64_t, rowids[i], rowids[j]); - } - } + int n = c->row_count; + + for (int end = n - 1; end > 0; --end) { + SWAP(double, distance[0], distance[end]); + SWAP(int64_t, rowids[0], rowids[end]); + vTopKSiftDown(distance, rowids, end, 0); } - - if (distance[row_count-1] == INFINITY) ++counter; + + int counter = 0; + while (counter < n && distance[n - 1 - counter] == INFINITY) ++counter; return counter; } @@ -3203,11 +3199,7 @@ static int vFullScanRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1 if (nearly_zero_float32(distance)) distance = 0.0; VECTOR_PRINT((void*)v2, vt, dimension); - if (distance < c->distance[c->max_index]) { - c->distance[c->max_index] = distance; - c->rowids[c->max_index] = (int64_t)sqlite3_column_int64(vm, 0); - c->max_index = vFullScanFindMaxIndex(c->distance, c->row_count); - } + if (distance < c->distance[0]) vTopKReplaceWorst(c, distance, (int64_t)sqlite3_column_int64(vm, 0)); } cleanup: @@ -3237,10 +3229,7 @@ static int vQuantRunMemory(vFullScanCursor *c, const preload_index *idx, uint8_t return SQLITE_CORRUPT; } - double *distance = c->distance; - int64_t *rowids = (int64_t *)c->rowids; - int max_index = c->max_index; - double current_max = distance[max_index]; + double current_max = c->distance[0]; // compute distance function vector_distance vd = c->table->options.v_distance; @@ -3260,16 +3249,11 @@ static int vQuantRunMemory(vFullScanCursor *c, const preload_index *idx, uint8_t if (nearly_zero_float32(dist)) dist = 0.0; if (dist < current_max) { - distance[max_index] = dist; - rowids[max_index] = INT64_FROM_INT8PTR(current_data); - - // Recompute max index efficiently - max_index = vFullScanFindMaxIndex(distance, c->row_count); - current_max = distance[max_index]; + vTopKReplaceWorst(c, dist, INT64_FROM_INT8PTR(current_data)); + current_max = c->distance[0]; } } - c->max_index = max_index; return SQLITE_OK; } @@ -3312,10 +3296,7 @@ static int vTurboRunPackedRows (vFullScanCursor *c, const uint8_t *data, sqlite3 return SQLITE_CORRUPT; } - double *distance = c->distance; - int64_t *rowids = c->rowids; - int max_index = c->max_index; - double current_max = distance[max_index]; + double current_max = c->distance[0]; for (int i = 0; i < counter; ++i) { const uint8_t *current = data + ((size_t)i * total_stride); @@ -3340,14 +3321,11 @@ static int vTurboRunPackedRows (vFullScanCursor *c, const uint8_t *data, sqlite3 if (nearly_zero_float32(dist)) dist = 0.0f; if (dist < current_max) { - distance[max_index] = dist; - rowids[max_index] = INT64_FROM_INT8PTR(current); - max_index = vFullScanFindMaxIndex(distance, c->row_count); - current_max = distance[max_index]; + vTopKReplaceWorst(c, dist, INT64_FROM_INT8PTR(current)); + current_max = c->distance[0]; } } - c->max_index = max_index; return SQLITE_OK; } @@ -3508,7 +3486,7 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz } // cache the maximum value to avoid repeated memory accesses - double current_max_distance = c->distance[c->max_index]; + double current_max_distance = c->distance[0]; for (int i=0; idistance[c->max_index] = distance; - c->rowids[c->max_index] = INT64_FROM_INT8PTR(current_data); - c->max_index = vFullScanFindMaxIndex(c->distance, c->row_count); - current_max_distance = c->distance[c->max_index]; // update cached max + vTopKReplaceWorst(c, distance, INT64_FROM_INT8PTR(current_data)); + current_max_distance = c->distance[0]; // update cached max } } }