From 396b74c95fe13975c73c4c0d9b6651ce2bb9f627 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 18:54:04 -0500 Subject: [PATCH] Use NEON on Linux arm64 The Makefile enabled NEON only on macOS arm64, so a Linux arm64 build 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 build needs only the define, inside the guard that keeps host detection away from 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 rejected the NEON code: l1_int8_neon() passed the signed vectors vabdq_s8() returns to vpaddlq_u8(), and the unsigned sums to vaddq_s32(), conversions Clang makes implicitly and GCC refuses. They are now explicit reinterprets, as photostructure/sqlite-vec a5aec74 writes them, which changes no result: L1 distances between int8 vectors match the scalar loop's. Built with GCC 14 at -O3 on Linux arm64, an L2 KNN query on a 100,000-row float[768] table takes 32.6 ms instead of 77.1 ms for k = 10 and 39.7 ms instead of 84.3 ms for k = 100 (medians of 20 queries). L1 and Hamming distances are bit-identical to the scalar loops' on 2,000 random pairs at each of six lengths from 8 to 1,536. The float32 L1 kernel adds in another order than the scalar loop, though, both in double, so where the sum rounds the two differ in the last digits: on values spread from 1e-8 to 1e8, for 176 of 500 random pairs at 8 dimensions and 467 of 500 at 768. Float32 L2 distances differ in the last digits, since the NEON kernel splits the sum across 16 partial sums where the scalar loop keeps one, and are closer to the exact value: against a float64 reference the mean relative error at 768 dimensions is 3.0e-8 instead of 1.7e-7. 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 queries swapped such rows. The int8 L2 kernel sums exactly and the scalar loop in float, exact only below 2^24, so the two differ only past that, from about 1,536 random dimensions. The Test workflow ran nothing on Linux arm64: build-linux-aarch64-extension cross-compiles for 32-bit ARM on an x86_64 runner, where the Makefile leaves NEON off, and runs no tests. test-linux-aarch64-extension runs the x86_64 job's build and tests on ubuntu-24.04-arm. test-loadable.py passes on the NEON and the scalar GCC build, and make test-unit, which compiles without the Makefile's CFLAGS, passes too. --- .github/workflows/test.yaml | 12 ++++++++++++ CLAUDE.md | 2 +- Makefile | 10 ++++++++++ sqlite-vec.c | 27 +++++++++++++++++---------- 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 3cccd43c..330dc261 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -24,6 +24,18 @@ jobs: # after the upload, so dist/test-unit stays out of the artifact - run: make test-unit - run: cargo test --locked --manifest-path tests/Cargo.toml + test-linux-aarch64-extension: + runs-on: ubuntu-24.04-arm + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + with: + enable-cache: true + - run: ./scripts/vendor.sh + - run: make loadable static + - run: uv sync --directory tests + - run: make test-loadable python=./tests/.venv/bin/python + - run: make test-unit build-macos-x86_64-extension: runs-on: macos-15-intel steps: diff --git a/CLAUDE.md b/CLAUDE.md index 699a9253..214a764b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -36,7 +36,7 @@ Run `./scripts/vendor.sh` first to download vendored dependencies (sqlite3.c, sh - Set `AR=` to use a different archiver (default: ar) **SIMD control:** -- SIMD is auto-enabled on Darwin x86_64 (AVX) and Darwin arm64 (NEON) +- SIMD is auto-enabled on Darwin x86_64 (AVX), Darwin arm64 (NEON), and Linux arm64 (NEON) - Set `OMIT_SIMD=1` to disable SIMD optimizations ### Testing diff --git a/Makefile b/Makefile index 1845396c..a8b58874 100644 --- a/Makefile +++ b/Makefile @@ -66,6 +66,16 @@ ifndef OMIT_SIMD ifeq ($(shell uname -sm),Darwin arm64) CFLAGS += -mcpu=apple-m1 -DSQLITE_VEC_ENABLE_NEON endif + # NEON is in the arm64 baseline (GCC and Clang define __ARM_NEON for + # aarch64 with no -m flag), so the Linux 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. + ifeq ($(shell uname -s),Linux) + ifneq ($(filter aarch64-%,$(shell $(CC) -dumpmachine)),) + CFLAGS += -DSQLITE_VEC_ENABLE_NEON + endif + endif endif endif diff --git a/sqlite-vec.c b/sqlite-vec.c index ab77083c..5470b694 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -320,26 +320,32 @@ static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v, int32x4_t acc3 = vdupq_n_s32(0); int32x4_t acc4 = vdupq_n_s32(0); + // vabdq_s8 returns |a - b| (up to 255) in signed lanes, so each difference + // is reinterpreted as unsigned before the widening pairwise adds 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))); + uint8x16_t diff1 = vreinterpretq_u8_s8(vabdq_s8(v1, v2)); + acc1 = + vaddq_s32(acc1, vreinterpretq_s32_u32(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_s32(acc2, vreinterpretq_s32_u32(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_s32(acc3, vreinterpretq_s32_u32(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_s32(acc4, vreinterpretq_s32_u32(vpaddlq_u16(vpaddlq_u8(diff4)))); pVect1 += 64; pVect2 += 64; @@ -348,8 +354,9 @@ 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_s32(acc1, vreinterpretq_s32_u32(vpaddlq_u16(vpaddlq_u8(diff)))); pVect1 += 16; pVect2 += 16; }