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; }