From 4866c62f1e651e7ebb5ee72e8cc9c68694d668f3 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Fri, 23 Oct 2026 11:40:01 -0500 Subject: [PATCH] Keep the NEON int8 cosine sums in 64 bits cosine_int8_neon() summed the products of int8 elements exactly, but in 32 bits, which the squares of -128 overflow from 131,072 elements. vec_distance_cosine() then returned NULL where a magnitude wrapped negative or to zero, or a wrong distance: 3.27 for a vector of 300,000 elements of 127 and one of 300,000 elements of 64, whose distance is 0. Each pass now sums its products in 32-bit lanes, which hold at most 4 products each, and adds them into 64-bit lanes, as mceachen's photostructure/sqlite-vec@4e7a65b does in that fork for the same kernel. Widening only the final sum would not be enough: over 1,048,576 elements, the lane that holds every 8th square of -128 reaches 2^31 on its own. test_vec_distance_cosine_int8_overflow checks mceachen's three cases from that commit against a float64 reference: 131,072 squares of 128, whose sum is 2^31; 2^32 in total, which wraps a 32-bit sum to 0; and 1,048,576 squares of 128, which overflow each 32-bit lane. Their squares keep every partial sum exact in float, so the scalar loop matches too. On main the NEON build returns NULL for all three. --- sqlite-vec.c | 37 ++++++++++++++++++++++++------------- tests/test-loadable.py | 23 +++++++++++++++++++++++ 2 files changed, 47 insertions(+), 13 deletions(-) diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..62a2cfc7 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -523,20 +523,26 @@ static f32 cosine_int8(const void *pA, const void *pB, const void *pD) { } #ifdef SQLITE_VEC_ENABLE_NEON +// Products of int8 elements are summed exactly: a 32-bit lane holds one pass's +// products, but 131,072 squares of -128 overflow a 32-bit sum, so each pass's +// lanes are added into 64-bit lanes. static f32 cosine_int8_neon(const void *pA, const void *pB, const void *pD) { const i8 *a = (const i8 *)pA; const i8 *b = (const i8 *)pB; size_t d = *((const size_t *)pD); const i8 *aEnd = a + d; - int32x4_t dot_acc1 = vdupq_n_s32(0); - int32x4_t dot_acc2 = vdupq_n_s32(0); - int32x4_t aMag_acc1 = vdupq_n_s32(0); - int32x4_t aMag_acc2 = vdupq_n_s32(0); - int32x4_t bMag_acc1 = vdupq_n_s32(0); - int32x4_t bMag_acc2 = vdupq_n_s32(0); + int64x2_t dot_acc = vdupq_n_s64(0); + int64x2_t aMag_acc = vdupq_n_s64(0); + int64x2_t bMag_acc = vdupq_n_s64(0); while (a < aEnd - 31) { + int32x4_t dot_acc1 = vdupq_n_s32(0); + int32x4_t dot_acc2 = vdupq_n_s32(0); + int32x4_t aMag_acc1 = vdupq_n_s32(0); + int32x4_t aMag_acc2 = vdupq_n_s32(0); + int32x4_t bMag_acc1 = vdupq_n_s32(0); + int32x4_t bMag_acc2 = vdupq_n_s32(0); int8x16_t va1 = vld1q_s8(a); int8x16_t vb1 = vld1q_s8(b); int16x8_t a1_lo = vmovl_s8(vget_low_s8(va1)); @@ -581,11 +587,17 @@ static f32 cosine_int8_neon(const void *pA, const void *pB, const void *pD) { bMag_acc2 = vmlal_s16(bMag_acc2, vget_low_s16(b2_hi), vget_low_s16(b2_hi)); bMag_acc2 = vmlal_s16(bMag_acc2, vget_high_s16(b2_hi), vget_high_s16(b2_hi)); + dot_acc = vpadalq_s32(dot_acc, vaddq_s32(dot_acc1, dot_acc2)); + aMag_acc = vpadalq_s32(aMag_acc, vaddq_s32(aMag_acc1, aMag_acc2)); + bMag_acc = vpadalq_s32(bMag_acc, vaddq_s32(bMag_acc1, bMag_acc2)); a += 32; b += 32; } while (a < aEnd - 15) { + int32x4_t dot_acc1 = vdupq_n_s32(0); + int32x4_t aMag_acc1 = vdupq_n_s32(0); + int32x4_t bMag_acc1 = vdupq_n_s32(0); int8x16_t va = vld1q_s8(a); int8x16_t vb = vld1q_s8(b); int16x8_t a_lo = vmovl_s8(vget_low_s8(va)); @@ -608,17 +620,16 @@ static f32 cosine_int8_neon(const void *pA, const void *pB, const void *pD) { bMag_acc1 = vmlal_s16(bMag_acc1, vget_low_s16(b_hi), vget_low_s16(b_hi)); bMag_acc1 = vmlal_s16(bMag_acc1, vget_high_s16(b_hi), vget_high_s16(b_hi)); + dot_acc = vpadalq_s32(dot_acc, dot_acc1); + aMag_acc = vpadalq_s32(aMag_acc, aMag_acc1); + bMag_acc = vpadalq_s32(bMag_acc, bMag_acc1); a += 16; b += 16; } - int32x4_t dot_sum = vaddq_s32(dot_acc1, dot_acc2); - int32x4_t aMag_sum = vaddq_s32(aMag_acc1, aMag_acc2); - int32x4_t bMag_sum = vaddq_s32(bMag_acc1, bMag_acc2); - - i32 dot = vaddvq_s32(dot_sum); - i32 aMag = vaddvq_s32(aMag_sum); - i32 bMag = vaddvq_s32(bMag_sum); + i64 dot = vaddvq_s64(dot_acc); + i64 aMag = vaddvq_s64(aMag_acc); + i64 bMag = vaddvq_s64(bMag_acc); while (a < aEnd) { dot += (i32)*a * (i32)*b; diff --git a/tests/test-loadable.py b/tests/test-loadable.py index 0044144b..f45765f0 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -259,6 +259,29 @@ def npy_l2(a, b): return np.linalg.norm(a - b) +def test_vec_distance_cosine_int8_overflow(): + vec_distance_cosine = lambda a, b: db.execute( + "select vec_distance_cosine(vec_int8(?), vec_int8(?))", [a, b] + ).fetchone()[0] + + def check(a, b): + a = np.array(a, dtype=np.int8) + b = np.array(b, dtype=np.int8) + # the reference in float64: numpy sums int8 products in int8 + y = npy_cosine(a.astype(np.float64), b.astype(np.float64)) + assert isclose(vec_distance_cosine(a, b), y, abs_tol=1e-6) + + # check overflow of 32-bit sums. Squares of 128 (2^14) and 64 (2^12) keep + # every partial sum exact in f32, so the scalar build matches too. + # 131,072 squares of 128 sum to 2^31, past INT32_MAX: + check([-128] * 131072, [-128] * 131072) + # 2^32 in total, which wraps a 32-bit sum to 0: + check([64] * 1048576, [-64] * 1048576) + # 32-bit partial sums of every 8th square reach 2^31, so widening only the + # final sum to 64 bits is not enough: + check([-128] * 1048576, [-128] * 1048576) + + def test_vec_distance_cosine(): vec_distance_cosine = lambda *args, a="?", b="?": db.execute( f"select vec_distance_cosine({a}, {b})", args