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