Skip to content

Use NEON for cosine distance - #5

Closed
MayCXC wants to merge 1 commit into
photostructure:mainfrom
MayCXC:neon-cosine-photostructure
Closed

MayCXC wants to merge 1 commit into
photostructure:mainfrom
MayCXC:neon-cosine-photostructure

Conversation

@MayCXC

@MayCXC MayCXC commented Oct 22, 2026 •

Copy link
Copy Markdown

distance_cosine_float() and distance_cosine_int8() have no NEON kernel, so cosine KNN queries in NEON builds, which include the darwin-arm64 binary, run the scalar loop. On Linux arm64 (GCC 14, -O3), with KNN times the medians of 20 queries on a 100,000-row float[768] table and the distance time the best of 7 runs over 20,000 vectors:

scalar this change
one 768-dimension float32 distance 895 ns 155 ns
KNN, k = 10 94.4 ms 25.9 ms
KNN, k = 100 98.9 ms 30.1 ms

Upstream already has both kernels: asg017/sqlite-vec@0de765f added cosine_float_neon() and cosine_int8_neon(), dispatched above 16 and 15 dimensions. This ports them with two changes:

The int8 kernel is otherwise asg017's as it is, summing products exactly in 32-bit lanes.

The float32 kernel's distances match the scalar loop's on 12,000 random pairs of 16 to 1,536 dimensions, except that a vector's distance to itself, 0 or about 2e-16 in either, can change between the two. The int8 scalar loop sums in float, which is exact only while a sum stays below 2^24, so the two match up to 1,024 dimensions, and past that the kernel's distances are closer to the exact value for vectors whose squares sum past 2^24 (a mean error of 2.3e-8 instead of 2.5e-7 at 2,048 dimensions). The KNN queries above return the same rows at every rank.

The cosine tests used vectors of 2 and 3 elements, below both kernels' thresholds, and pass against asg017's kernels unchanged. test_vec_distance_cosine now checks lengths 17, 31, 32, 40, 127 and 768 against a float64 reference (numpy sums int8 products in int8, which overflows at these lengths), and the zero-vector and extreme-magnitude tests repeat at 40 dimensions. asg017's float32 kernel fails both of those; with a zero-vector check added it still fails the extreme-magnitude test. Rebased onto main at f793711, the full suite passes on the NEON build and on the scalar build (449 passed, 3 skipped), and make test-unit passes. The commit adds a ### Changed entry under [Unreleased] in CHANGELOG.md.

#6 turns NEON on for the Linux arm64 binaries. The two are independent, and each conflicts with #2, #7, #8 and the other only in CHANGELOG.md's [Unreleased] section. The same port is open on vlasky/sqlite-vec as vlasky#20, and asg017#338 sums asg017's own kernel and scalar loop in double.

distance_cosine_float() and distance_cosine_int8() had no NEON kernel,
so cosine KNN queries in NEON builds, which include the darwin-arm64
binary, ran the scalar loop.

asg017/sqlite-vec 0de765f added cosine_float_neon() and
cosine_int8_neon(), dispatched above 16 and 15 dimensions. This ports
them with two changes. asg017's float32 kernel sums in float32, the
loss asg017#324 reported and 2.1.0 fixed in the scalar loop by summing
in double: vectors with elements below about 3e-23 or above about 2e19
get wrong distances. Here each float32x4 load is widened to two
float64x2 halves before it is multiplied and added, in asg017's
16-element blocks with four accumulators per sum. And both kernels
return 1.0 for a zero vector, as the scalar loops do, where asg017's
return NaN.

On Linux arm64 (GCC 14, -O3), one 768-dimension float32 distance takes
about 155 ns instead of 895 ns, and a KNN query on a 100,000-row
float[768] table takes 25.9 ms instead of 94.4 ms for k = 10 and
30.1 ms instead of 98.9 ms for k = 100 (medians of 20 queries),
returning the same rows. The float32 distances match the scalar loop's
on 12,000 random pairs of 16 to 1,536 dimensions, except that a
vector's distance to itself, 0 or about 2e-16 in either, can change
between the two. The int8 kernel sums products exactly in 32-bit
lanes, and the scalar loop sums them in float, which is exact while a
sum stays below 2^24: their distances match up to 1,024 dimensions,
and past that the kernel's are closer to the exact value for vectors
whose squares sum past 2^24 (a mean error of 2.3e-8 instead of 2.5e-7
at 2,048 dimensions).

The cosine tests used vectors of 2 and 3 elements, below both kernels'
thresholds, and pass against asg017's kernels as they are.
test_vec_distance_cosine now checks lengths of 17 to 768, with and
without a remainder past the 16-element blocks, against a reference
computed in float64, since numpy sums int8 products in int8; the
zero-vector and extreme-magnitude tests repeat at 40 dimensions.
asg017's float32 kernel fails both, and with a zero-vector check added
it still fails the extreme-magnitude test. The full suite passes on
the NEON build and on the scalar one, and make test-unit, which builds
without SIMD, passes too.
@MayCXC
MayCXC force-pushed the neon-cosine-photostructure branch from a4f7976 to 6df0e50 Compare October 23, 2026 03:33
mceachen added a commit that referenced this pull request Oct 23, 2026
#5's cosine_int8_neon() kept its sums in 32 bits, so on NEON builds
vec_distance_cosine() returned NULL, 1.0, or a wrong distance such as 3.8
for int8 vectors whose squares sum past 2^31 - 1, from 131,072 elements.
Its CHANGELOG entry also overstated how closely the NEON float32 and int8
distances match the scalar code's.
@mceachen

Copy link
Copy Markdown
Member

Landed as 050bb59, thanks. Follow-up 4e7a65b: cosine_int8_neon() summed in 32 bits, so vec_distance_cosine() on int8 vectors whose squares sum past 2^31 − 1 (from 131,072 elements of −128) returned NULL, 1.0, or a wrong distance such as 3.8. Each pass's products now go into 64-bit lanes with vpadalq_s32, as in l2_sqr_int8_neon(), with tests at 2^31, 2^32, and per-lane overflow. The CHANGELOG now also says float32 distances can differ in the last digits, mostly between near-duplicate vectors, and that long int8 vectors usually, not always, get closer distances. vlasky#20 and asg017's main (0de765f) have the same 32-bit int8 kernel.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants