Skip to content

Use NEON on Linux arm64 #339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MayCXC wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Use NEON on Linux arm64 #339

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dcaa5cd
Fix signedness in NEON int8 L1 distance
ariannamethod Sep 23, 2026
8d133f2
Use NEON on Linux arm64
MayCXC Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
33 changes: 17 additions & 16 deletions sqlite-vec.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) {
Expand All @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions tests/test-loadable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down