Skip to content

Use NEON for cosine distance - #20

Merged
vlasky merged 1 commit into
vlasky:mainfrom
MayCXC:neon-cosine-vlasky
Oct 28, 2026
Merged

vlasky merged 1 commit into
vlasky:mainfrom
MayCXC:neon-cosine-vlasky

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 macOS arm64 build, run the scalar loop. Built with -DSQLITE_VEC_ENABLE_NEON on Linux arm64 (Clang 19, -O3; medians of 20 queries on a 100,000-row float[768] table):

k scalar this change
10 102.8 ms 39.2 ms
100 126.2 ms 46.2 ms

asg017/sqlite-vec 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 three changes:

Neither kernel checks for a zero vector: vec_distance_cosine() returns NULL for one, and vec0 refuses them in cosine columns (#8), before any distance is computed.

The KNN queries above return the same rows and distances. The float32 kernel's distances can differ from the scalar loop's in the last digits, since it adds the same double products in another order: none of 5,000 random pairs of 17 to 1,536 dimensions differed, but 68 to 924 of 1,000 near-duplicate pairs (a vector and itself plus noise of 1e-4) did at each length, and a vector's distance to itself, between -2.2e-16 and 2.2e-16 in either, changed for 175 to 575 of 1,000. The int8 distances match up to 1,024 dimensions, where every sum is exact in float; past that, vectors whose squares sum past 2^24 can get different distances, closer to the exact value: at 2,048 dimensions a mean error of 2.3e-8 instead of 2.5e-7, and at 4,096, 896 of 1,000 random pairs came closer and none farther. mceachen's 4e7a65b corrected the same two claims in photostructure's CHANGELOG, and this description had overstated them too.

The cosine tests used vectors of 2 and 3 elements, below both kernels' thresholds, and pass against asg017's float32 kernel 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 mceachen's three int8 cases from 4e7a65b, vectors of 131,072 and 1,048,576 elements whose squares overflow a 32-bit sum in total or in each lane; with this PR's earlier head, which kept 32-bit sums, the kernel returned NULL for all three. The extreme-magnitude test repeats at 40 dimensions, which asg017's float32 kernel fails. make test-loadable passes on the NEON build and on the scalar one (143 passed, 4 skipped, 259 snapshots), and make test-unit, which builds without SIMD, passes too.

The numbers come from Clang because GCC cannot build the NEON code on main: l1_int8_neon() relies on implicit signed/unsigned vector conversions that GCC refuses. #21 fixes that and turns NEON on for Linux arm64. The two merge cleanly with each other and with #9 through #14, #17 through #19 and #22 through #24, and together they build with GCC and Clang without warnings and pass the suite. The same change is on photostructure/sqlite-vec's main, from photostructure#5 with mceachen's 4e7a65b on top. On asg017/sqlite-vec, asg017#338 sums asg017's own float32 kernel and scalar loop in double, which is the part of asg017#324 your comment there left for its NEON path, and asg017#356 keeps its int8 kernel's sums in 64 bits.

distance_cosine_float() and distance_cosine_int8() had no NEON kernel,
so cosine KNN queries in NEON builds, which include the macOS arm64
build, 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 three changes. asg017's float32 kernel sums in float32, the
loss d75c245 fixed in the scalar loop for asg017#324: 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. The int8 kernel finishes in double, as
distance_cosine_int8() does, where asg017's finishes in float. And it
adds each pass's 32-bit lanes into 64-bit sums, as mceachen's
photostructure/sqlite-vec@4e7a65b does in that fork, where asg017's
keeps the whole sum in 32 bits: squares of -128 overflow it from
131,072 elements, and vec_distance_cosine() then returns NULL or a
wrong distance, 3.27 for two constant vectors of 300,000 elements,
whose distance is 0. Neither kernel checks for a zero vector:
vec_distance_cosine() returns NULL for one, and vec0 refuses them in
cosine columns, before any distance is computed.

Built with -DSQLITE_VEC_ENABLE_NEON on Linux arm64 (Clang 19, -O3), a
KNN query on a 100,000-row float[768] table takes 39.2 ms instead of
102.8 ms for k = 10 and 46.2 ms instead of 126.2 ms for k = 100
(medians of 20 queries), returning the same rows and distances. The
float32 distances can differ from the scalar loop's in the last
digits, since the kernel adds the same double products in another
order: none of 5,000 random pairs of 17 to 1,536 dimensions differed,
but 68 to 924 of 1,000 near-duplicate pairs (a vector and itself plus
noise of 1e-4) did at each length, and a vector's distance to itself,
between -2.2e-16 and 2.2e-16 in either, changed for 175 to 575 of
1,000. The int8 distances match up to 1,024
dimensions, where every sum is exact in float; past that, vectors
whose squares sum past 2^24 can get different distances, closer to the
exact value: at 2,048 dimensions a mean error of 2.3e-8 instead of
2.5e-7, and at 4,096, 896 of 1,000 random pairs came closer and none
farther.

The cosine tests used vectors of 2 and 3 elements, below both kernels'
thresholds, and pass against asg017's float32 kernel as it is.
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, and
mceachen's int8 vectors of 131,072 and 1,048,576 elements from
4e7a65b, whose squares overflow a 32-bit sum in total or in each lane,
for which 32-bit sums return NULL. The extreme-magnitude test repeats
at 40 dimensions, which asg017's float32 kernel fails.
test-loadable.py passes on the NEON build and on the scalar one, and
make test-unit, which builds without SIMD, passes too.
@vlasky

vlasky commented Oct 28, 2026

Copy link
Copy Markdown
Owner

Thanks for the update. The updated int8 kernel checks out: folding each pass's lanes into int64x2_t with vpadalq_s32() keeps every partial sum exact, and the three new cases (131,072 squares of -128, the 2^32 wrap, and the per-lane 2^31 at 2^20 elements) are exactly the ones I had found on the previous head, where the first two returned NULL and a 2^20-element random pair returned a wrong value with no error. The float kernel's double widening checks out end to end, and the speedups reproduce here (about 3x for float, 5x for int8 on an M-series Mac). Full suite and make test-unit pass on the NEON and OMIT_SIMD=1 builds, and GCC 16 is clean apart from the l1_int8_neon conversions that #21 fixes.

Two nits, not blocking: the loop guards form aEnd - 31 and aEnd - 15 before comparing, which for a 16-element input is a pointer before the object (undefined in C, harmless on every flat-address target; aEnd - a >= 32 avoids it), and the final (f32)dot and (f32)aMag casts drop bits above 2^24 before the division, where (double) costs nothing. Merged as is.

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