Top-k in a heap, and integer kernels that use the instructions built for them

[marcobambini]

Stacked on #53 — review that one first, and this diff shows only the two commits
here. The two changes are independent of each other; both were measured before
and after, and both were checked for result equivalence, not just for speed.

Top-k 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²) exchange sort. 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. 20k rows, dim 768, ms per
query:

k	before	after
10	0.06	0.05	—
100	0.10	0.06	1.7×
1000	3.17	0.16	20×
4000	33.46	0.55	61×

Small k is unchanged, as expected — there was never much to rescan.

It 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.

Integer kernels

The u8/i8 kernels widened every byte to 32 bits before multiplying — twelve or
more instructions per 16 bytes — on a single accumulator chain. Both
architectures have instructions for exactly this shape. NEON: absolute
difference, widening multiply, pairwise-accumulate, five instructions per 16
bytes. x86: PSADBW sums the absolute differences of a whole vector in one go,
which is L1 outright, 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.

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 become the unsigned kernel plus one XOR per vector. Dot
and cosine need the true signed values. AVX2 and AVX-512 cosine also stop making
three separate passes over the data.

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 change that moves, because u8 rows are a quarter the size
of f32 and the scan is compute-bound rather than bandwidth-bound. Quantized scan
over 60k rows at dim 768: 17.4 → 62.6 Mvec/s preloaded (3.6×), 13.1 → 29.2
from disk. For comparison, the f32 kernel work in #53 only reached 1.6×
end to end, because there the scan is waiting on memory.

Accuracy improves rather than degrades. The reductions now widen to 64 bits
before folding lanes — each lane is itself a running total, and folding eight or
sixteen of them in 32 bits capped the usable dimension well below what the
accumulators could hold — and cosine sums exact integers and divides in double
instead of 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.

Verification

• Every kernel against a double-precision reference on CPU, NEON, SSE2 and AVX2.
• AVX-512: the avx512 job ran the full suite on these kernels under Intel SDE —
  distance backend: AVX512, 1447/1447. Both runs so far landed on AMD EPYC
  runners with no AVX-512, so the emulated path is the one that got exercised.
• End to end: 144 top-k queries across every distance, quantization type and both
  scan shapes. 136 return identical rows; the 8 that differ are ties where either
  choice is equally correct. All 144 equivalent.
• 1447/1447 under ASan. Full CI matrix green, 17/17.

Notes for review

• Tie-breaking among equal distances changes. Neither sort is stable and the
  old order was not a guarantee, but it is observable.
• SSE2 is deliberately untouched — it is only selected on a CPU with neither
  AVX-512 nor AVX2.
• The scalar CPU cosine still accumulates u8 norms in uint32_t, which wraps
  above dimension ≈ 66,000. That is now the only backend with the limit; it is
  pre-existing and not addressed here.
• P5 from the audit — four database rows per query load, plus prefetch — is
  not in this PR. Measured: the prefetch does nothing (the loop is a
  sequential bandwidth-bound stream, the hardware prefetcher already has it), and
  batching is worth 1.3× at the cost of a batched-kernel API across five
  backends. The 3.2× that looked like batching in the prototype turned out to be
  the badly written kernel this PR replaces.

🤖 Generated with Claude Code