From b52880bf6d098455fef8a18d32a072514bfd03e9 Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 22:14:59 +0200 Subject: [PATCH 1/2] perf: keep the top-k candidates in a heap Every accepted candidate triggered an O(k) rescan of all k slots to find the new worst entry, and the final ordering was an O(k^2) exchange sort. For a scan where the distance itself is cheap - 1-bit Hamming over a preloaded index is 96 bytes of popcount per row - that bookkeeping was the query. The candidate set is now a binary max-heap over (distance, rowid) in the parallel arrays the cursor already owns: the root is the worst entry still in the set, so admitting a candidate is one comparison plus O(log k), and the final ordering is a heapsort in place. 20k rows, dim 768, 1BIT preloaded, ms per query: k=10 0.06 -> 0.05 k=100 0.10 -> 0.06 k=1000 3.17 -> 0.16 20x k=4000 33.46 -> 0.55 61x Small k is unchanged, as expected - there was never much to rescan. This also removes a latent out-of-bounds read. max_index was never reset between filters, so a cursor reused with a smaller k (a correlated subquery, or a join where k varies per row) indexed the reallocated array with a stale offset. A heap root is always slot zero, so there is no index to go stale. Tie-breaking among equal distances changes: neither sort is stable and the old order was not a guarantee. Across 144 top-k queries covering every distance, quantization type, and both scan shapes, all 144 return the same multiset of distances; 136 return the same rows in the same order, and the 8 that differ are cases where several rows tie at the k-th distance and either choice is equally correct. Co-Authored-By: Claude Opus 5 --- src/sqlite-vector.c | 116 ++++++++++++++++++-------------------------- 1 file changed, 46 insertions(+), 70 deletions(-) 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 } } } From be9b8be8b2aee8efcbb8a5ebcf897a6571b0424e Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 22:26:49 +0200 Subject: [PATCH 2/2] perf: rewrite the u8/i8 kernels around the instructions built for them The integer kernels widened every byte to 32 bits before multiplying - twelve or more instructions per 16 bytes - and accumulated on a single dependency chain. Both architectures have instructions for exactly this shape: * NEON: absolute difference, widening multiply, pairwise-accumulate. Five instructions per 16 bytes instead of twelve. * x86: PSADBW sums the absolute differences of a whole vector in one go, which is L1 outright; 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 reuse the unsigned ones where 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 unsigned kernel plus one XOR per vector. Dot and cosine need the true signed values and get their own widening path. AVX2 and AVX-512 cosine also stop making three separate passes over the data. Plus four accumulators, as the f32 kernels already had. NEON, dim 768, cache-resident (Mvec/s): L2 SQ_L2 COSINE DOT L1 u8 18 -> 68 18 -> 73 12 -> 29 15 -> 83 22 -> 98 i8 18 -> 53 18 -> 53 12 -> 29 15 -> 82 25 -> 74 End to end this is the result that actually moves, because u8 rows are a quarter the size of f32 and the scan is compute-bound rather than bandwidth-bound. A quantized scan over 60k rows at dim 768: preloaded 17.4 -> 62.6 Mvec/s 3.6x from disk 13.1 -> 29.2 Mvec/s 2.2x Accuracy improves: the reductions now widen to 64 bits before folding lanes - each lane is itself a running total, and folding them in 32 bits capped the usable dimension well below what the accumulators could hold - and cosine sums exact integers and divides in double rather than accumulating in float. Against a double-precision reference across 25 dimensions including tails, u8/i8 cosine goes from 5.9e-08 to 2.9e-08 worst error, and every other integer kernel is unchanged or better. Verified on CPU, NEON, SSE2 and AVX2 with that reference harness, and end to end across 144 top-k queries covering every distance, quantization type and scan shape: 136 return identical rows, and the 8 that differ are ties where either choice is equally correct. AVX-512 is compile-verified here; the avx512 CI job runs the suite on those kernels under Intel SDE. SSE2 is left alone: it is only selected on a CPU with neither AVX-512 nor AVX2. Co-Authored-By: Claude Opus 5 --- src/distance-avx2.c | 458 +++++++++++++++++---------------------- src/distance-avx512.c | 337 ++++++++++++++--------------- src/distance-neon.c | 486 +++++++++++++++++------------------------- 3 files changed, 561 insertions(+), 720 deletions(-) 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 -