From dcaa5cde25333ea0d85cc99b808e5931f0768629 Mon Sep 17 00:00:00 2001 From: Arianna Method Date: Wed, 23 Sep 2026 05:55:36 +0300 Subject: [PATCH 1/2] Fix signedness in NEON int8 L1 distance --- sqlite-vec.c | 33 +++++++++++++++++---------------- tests/test-loadable.py | 4 ++++ 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..d6c43a92 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -292,31 +292,32 @@ static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v, const int8_t *pEnd1 = pVect1 + qty; - int32x4_t acc1 = vdupq_n_s32(0); - int32x4_t acc2 = vdupq_n_s32(0); - int32x4_t acc3 = vdupq_n_s32(0); - int32x4_t acc4 = vdupq_n_s32(0); + uint32x4_t acc1 = vdupq_n_u32(0); + uint32x4_t acc2 = vdupq_n_u32(0); + uint32x4_t acc3 = vdupq_n_u32(0); + uint32x4_t acc4 = vdupq_n_u32(0); while (pVect1 < pEnd1 - 63) { int8x16_t v1 = vld1q_s8(pVect1); int8x16_t v2 = vld1q_s8(pVect2); - int8x16_t diff1 = vabdq_s8(v1, v2); - acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff1))); + // Absolute differences can reach 255, so widen them as unsigned bytes. + uint8x16_t diff1 = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc1 = vaddq_u32(acc1, vpaddlq_u16(vpaddlq_u8(diff1))); v1 = vld1q_s8(pVect1 + 16); v2 = vld1q_s8(pVect2 + 16); - int8x16_t diff2 = vabdq_s8(v1, v2); - acc2 = vaddq_s32(acc2, vpaddlq_u16(vpaddlq_u8(diff2))); + uint8x16_t diff2 = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc2 = vaddq_u32(acc2, vpaddlq_u16(vpaddlq_u8(diff2))); v1 = vld1q_s8(pVect1 + 32); v2 = vld1q_s8(pVect2 + 32); - int8x16_t diff3 = vabdq_s8(v1, v2); - acc3 = vaddq_s32(acc3, vpaddlq_u16(vpaddlq_u8(diff3))); + uint8x16_t diff3 = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc3 = vaddq_u32(acc3, vpaddlq_u16(vpaddlq_u8(diff3))); v1 = vld1q_s8(pVect1 + 48); v2 = vld1q_s8(pVect2 + 48); - int8x16_t diff4 = vabdq_s8(v1, v2); - acc4 = vaddq_s32(acc4, vpaddlq_u16(vpaddlq_u8(diff4))); + uint8x16_t diff4 = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc4 = vaddq_u32(acc4, vpaddlq_u16(vpaddlq_u8(diff4))); pVect1 += 64; pVect2 += 64; @@ -325,13 +326,13 @@ static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v, while (pVect1 < pEnd1 - 15) { int8x16_t v1 = vld1q_s8(pVect1); int8x16_t v2 = vld1q_s8(pVect2); - int8x16_t diff = vabdq_s8(v1, v2); - acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff))); + uint8x16_t diff = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc1 = vaddq_u32(acc1, vpaddlq_u16(vpaddlq_u8(diff))); pVect1 += 16; pVect2 += 16; } - int32x4_t acc = vaddq_s32(vaddq_s32(acc1, acc2), vaddq_s32(acc3, acc4)); + uint32x4_t acc = vaddq_u32(vaddq_u32(acc1, acc2), vaddq_u32(acc3, acc4)); int32_t sum = 0; while (pVect1 < pEnd1) { @@ -341,7 +342,7 @@ static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v, pVect2++; } - return vaddvq_s32(acc) + sum; + return vaddvq_u32(acc) + sum; } static double l1_f32_neon(const void *pVect1v, const void *pVect2v, diff --git a/tests/test-loadable.py b/tests/test-loadable.py index 0044144b..64bf0d45 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -326,6 +326,10 @@ def check(a, b, dtype=np.float32): # check overflow check([127] * 20, [-128] * 20, dtype=np.int8) check([-128, 127], [127, -128], dtype=np.int8) + # Exercise NEON's 16- and 64-element loops, including scalar leftovers. + for n in (15, 16, 17, 63, 64, 65, 80, 81): + check([127] * n, [-128] * n, dtype=np.int8) + check([-128] * n, [127] * n, dtype=np.int8) check( [1, 2, 3, 4, 5, 6, 7, 8, 1, 1, 2, 3, 4, 5, 6, 7, 8, 1], [1, 20, 38, 23, 29, 4, 10, 9, 3, 1, 20, 38, 23, 29, 4, 10, 9, 3], From 8d133f2f2e51f3a229a8a3a8693c62b9018a4a03 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 19:05:16 -0500 Subject: [PATCH 2/2] Use NEON on Linux arm64 The Makefile enabled NEON only on macOS arm64, so a Linux arm64 build, the linux-aarch64 release binary among them, ran every distance's scalar loop. NEON is part of the arm64 baseline: GCC and Clang define __ARM_NEON for aarch64 with no -m flag, so the Linux block needs only the define, beside its AVX check and inside its guard against Android cross-compiles. OMIT_SIMD=1 still builds without it. The define follows the compiler's target, `$(CC) -dumpmachine`, not uname, as mceachen's photostructure/sqlite-vec@350c67a has it: 32-bit Raspberry Pi OS runs a 64-bit kernel, so uname says aarch64 while gcc targets 32-bit ARM, which lacks the AArch64-only intrinsics the NEON kernels use. Under an aarch64 kernel, `make loadable CC=arm-linux-gnueabihf-gcc` fails when uname decides (implicit declaration of vaddvq_f32, unknown type name float64x2_t) and builds a 32-bit ARM extension when the compiler does; for gcc and clang targeting aarch64 the build commands are the same either way. GCC cannot build the NEON code on main, because l1_int8_neon() relies on implicit signed/unsigned vector conversions; this builds on #322, which makes them explicit. Built with GCC 14 at -O3 on Linux arm64, a KNN query on a 100,000-row float[768] table takes 19.6 ms instead of 59.2 ms for k = 10 and 24.2 ms instead of 65.8 ms for k = 100 with the L2 metric, and 17.5 ms instead of 85.3 ms and 23.3 ms instead of 90.6 ms with cosine (medians of 20 queries, each build in its own process: the extension exports its internal functions, so a second build loaded into the same process runs the first one's KNN code). L1 and Hamming distances were bit-identical to the scalar loops' on 2,000 random pairs at each of six lengths from 8 to 1,536, though the float32 L1 kernel adds in another order than the scalar loop, both in double, so the two can differ in the last digit on other inputs. Float32 L2 and cosine distances differ in the last digits, since the NEON kernels split their sums across 16 partial sums where the scalar loops keep one; for L2 the NEON result is closer to the exact value, a mean relative error of 3.0e-8 instead of 1.7e-7 at 768 dimensions. So a KNN query can order rows with nearly equal distances differently, as the macOS arm64 build already does; 4 of the 2,000 ranks of the k = 100 L2 queries swapped such rows. Int8 cosine distances can differ in the last digit as well, since cosine_int8_neon() finishes in float and the scalar loop in double. The Test workflow's build-linux-aarch64-extension job already runs the suite on an arm64 runner, so it now tests the NEON build. Locally the suite passes on it (478 passed, 132 skipped, 204 snapshots). --- Makefile | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/Makefile b/Makefile index 175ab169..1f6da873 100644 --- a/Makefile +++ b/Makefile @@ -47,6 +47,14 @@ ifndef OMIT_SIMD ifneq ($(filter avx,$(shell grep -o 'avx[^ ]*' /proc/cpuinfo 2>/dev/null | head -1)),) CFLAGS += -mavx -mavx2 -DSQLITE_VEC_ENABLE_AVX endif + # NEON is in the arm64 baseline (GCC and Clang define __ARM_NEON for + # aarch64 with no -m flag), so the build needs only the define. It follows + # the compiler's target, not uname: 32-bit Raspberry Pi OS runs a 64-bit + # kernel, so uname says aarch64 while gcc targets 32-bit ARM, which lacks + # the AArch64-only intrinsics the NEON kernels use. + ifneq ($(filter aarch64-%,$(shell $(CC) -dumpmachine)),) + CFLAGS += -DSQLITE_VEC_ENABLE_NEON + endif endif endif endif