Skip to content

Commit be9b8be

Browse files
marcobambiniclaude
andcommitted
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 <noreply@anthropic.com>
1 parent b52880b commit be9b8be

3 files changed

Lines changed: 561 additions & 720 deletions

File tree

0 commit comments

Comments
 (0)