From 8de0d018b033434fb632aebc58d9ae04f7a68e13 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 18:52:05 -0500 Subject: [PATCH] Use NEON for cosine distance distance_cosine_float() and distance_cosine_int8() had no NEON kernel, so cosine KNN queries in NEON builds, which include the macOS arm64 build, ran the scalar loop. asg017/sqlite-vec 0de765f added cosine_float_neon() and cosine_int8_neon(), dispatched above 16 and 15 dimensions. This ports them with three changes. asg017's float32 kernel sums in float32, the loss d75c245 fixed in the scalar loop for asg017#324: vectors with elements below about 3e-23 or above about 2e19 get wrong distances. Here each float32x4 load is widened to two float64x2 halves before it is multiplied and added, in asg017's 16-element blocks with four accumulators per sum. The int8 kernel finishes in double, as distance_cosine_int8() does, where asg017's finishes in float. And it adds each pass's 32-bit lanes into 64-bit sums, as mceachen's photostructure/sqlite-vec@4e7a65b does in that fork, where asg017's keeps the whole sum in 32 bits: squares of -128 overflow it from 131,072 elements, and vec_distance_cosine() then returns NULL or a wrong distance, 3.27 for two constant vectors of 300,000 elements, whose distance is 0. Neither kernel checks for a zero vector: vec_distance_cosine() returns NULL for one, and vec0 refuses them in cosine columns, before any distance is computed. Built with -DSQLITE_VEC_ENABLE_NEON on Linux arm64 (Clang 19, -O3), a KNN query on a 100,000-row float[768] table takes 39.2 ms instead of 102.8 ms for k = 10 and 46.2 ms instead of 126.2 ms for k = 100 (medians of 20 queries), returning the same rows and distances. The float32 distances can differ from the scalar loop's in the last digits, since the kernel adds the same double products in another order: none of 5,000 random pairs of 17 to 1,536 dimensions differed, but 68 to 924 of 1,000 near-duplicate pairs (a vector and itself plus noise of 1e-4) did at each length, and a vector's distance to itself, between -2.2e-16 and 2.2e-16 in either, changed for 175 to 575 of 1,000. The int8 distances match up to 1,024 dimensions, where every sum is exact in float; past that, vectors whose squares sum past 2^24 can get different distances, closer to the exact value: at 2,048 dimensions a mean error of 2.3e-8 instead of 2.5e-7, and at 4,096, 896 of 1,000 random pairs came closer and none farther. The cosine tests used vectors of 2 and 3 elements, below both kernels' thresholds, and pass against asg017's float32 kernel as it is. test_vec_distance_cosine now checks lengths of 17 to 768, with and without a remainder past the 16-element blocks, against a reference computed in float64, since numpy sums int8 products in int8, and mceachen's int8 vectors of 131,072 and 1,048,576 elements from 4e7a65b, whose squares overflow a 32-bit sum in total or in each lane, for which 32-bit sums return NULL. The extreme-magnitude test repeats at 40 dimensions, which asg017's float32 kernel fails. test-loadable.py passes on the NEON build and on the scalar one, and make test-unit, which builds without SIMD, passes too. --- sqlite-vec.c | 240 +++++++++++++++++++++++++++++++++++++++++ tests/test-loadable.py | 47 +++++--- 2 files changed, 274 insertions(+), 13 deletions(-) diff --git a/sqlite-vec.c b/sqlite-vec.c index ab77083c..583b4973 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -400,6 +400,236 @@ static double l1_f32_neon(const void *pVect1v, const void *pVect2v, return vaddvq_f64(acc) + sum; } + +// asg017/sqlite-vec's cosine_float_neon, with its sums kept in double as +// distance_cosine_float keeps them: each float32x4 load is widened to two +// float64x2 halves before it is multiplied and added. +static f32 cosine_float_neon(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + size_t qty16 = qty >> 4; + const f32 *pEnd1 = pVect1 + (qty16 << 4); + + float64x2_t dot0 = vdupq_n_f64(0), dot1 = vdupq_n_f64(0); + float64x2_t dot2 = vdupq_n_f64(0), dot3 = vdupq_n_f64(0); + float64x2_t amag0 = vdupq_n_f64(0), amag1 = vdupq_n_f64(0); + float64x2_t amag2 = vdupq_n_f64(0), amag3 = vdupq_n_f64(0); + float64x2_t bmag0 = vdupq_n_f64(0), bmag1 = vdupq_n_f64(0); + float64x2_t bmag2 = vdupq_n_f64(0), bmag3 = vdupq_n_f64(0); + + while (pVect1 < pEnd1) { + float32x4_t v1, v2; + float64x2_t a, b; + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + a = vcvt_f64_f32(vget_low_f32(v1)); + b = vcvt_f64_f32(vget_low_f32(v2)); + dot0 = vfmaq_f64(dot0, a, b); + amag0 = vfmaq_f64(amag0, a, a); + bmag0 = vfmaq_f64(bmag0, b, b); + a = vcvt_high_f64_f32(v1); + b = vcvt_high_f64_f32(v2); + dot0 = vfmaq_f64(dot0, a, b); + amag0 = vfmaq_f64(amag0, a, a); + bmag0 = vfmaq_f64(bmag0, b, b); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + a = vcvt_f64_f32(vget_low_f32(v1)); + b = vcvt_f64_f32(vget_low_f32(v2)); + dot1 = vfmaq_f64(dot1, a, b); + amag1 = vfmaq_f64(amag1, a, a); + bmag1 = vfmaq_f64(bmag1, b, b); + a = vcvt_high_f64_f32(v1); + b = vcvt_high_f64_f32(v2); + dot1 = vfmaq_f64(dot1, a, b); + amag1 = vfmaq_f64(amag1, a, a); + bmag1 = vfmaq_f64(bmag1, b, b); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + a = vcvt_f64_f32(vget_low_f32(v1)); + b = vcvt_f64_f32(vget_low_f32(v2)); + dot2 = vfmaq_f64(dot2, a, b); + amag2 = vfmaq_f64(amag2, a, a); + bmag2 = vfmaq_f64(bmag2, b, b); + a = vcvt_high_f64_f32(v1); + b = vcvt_high_f64_f32(v2); + dot2 = vfmaq_f64(dot2, a, b); + amag2 = vfmaq_f64(amag2, a, a); + bmag2 = vfmaq_f64(bmag2, b, b); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + a = vcvt_f64_f32(vget_low_f32(v1)); + b = vcvt_f64_f32(vget_low_f32(v2)); + dot3 = vfmaq_f64(dot3, a, b); + amag3 = vfmaq_f64(amag3, a, a); + bmag3 = vfmaq_f64(bmag3, b, b); + a = vcvt_high_f64_f32(v1); + b = vcvt_high_f64_f32(v2); + dot3 = vfmaq_f64(dot3, a, b); + amag3 = vfmaq_f64(amag3, a, a); + bmag3 = vfmaq_f64(bmag3, b, b); + } + + double dot = + vaddvq_f64(vaddq_f64(vaddq_f64(dot0, dot1), vaddq_f64(dot2, dot3))); + double aMag = + vaddvq_f64(vaddq_f64(vaddq_f64(amag0, amag1), vaddq_f64(amag2, amag3))); + double bMag = + vaddvq_f64(vaddq_f64(vaddq_f64(bmag0, bmag1), vaddq_f64(bmag2, bmag3))); + + const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4)); + while (pVect1 < pEnd2) { + dot += (double)*pVect1 * (double)*pVect2; + aMag += (double)*pVect1 * (double)*pVect1; + bMag += (double)*pVect2 * (double)*pVect2; + pVect1++; + pVect2++; + } + + return 1 - (dot / (sqrt(aMag) * sqrt(bMag))); +} + +// asg017/sqlite-vec's cosine_int8_neon, finishing as distance_cosine_int8 +// does. 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; + + 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)); + int16x8_t a1_hi = vmovl_s8(vget_high_s8(va1)); + int16x8_t b1_lo = vmovl_s8(vget_low_s8(vb1)); + int16x8_t b1_hi = vmovl_s8(vget_high_s8(vb1)); + + dot_acc1 = vmlal_s16(dot_acc1, vget_low_s16(a1_lo), vget_low_s16(b1_lo)); + dot_acc1 = vmlal_s16(dot_acc1, vget_high_s16(a1_lo), vget_high_s16(b1_lo)); + dot_acc2 = vmlal_s16(dot_acc2, vget_low_s16(a1_hi), vget_low_s16(b1_hi)); + dot_acc2 = vmlal_s16(dot_acc2, vget_high_s16(a1_hi), vget_high_s16(b1_hi)); + + aMag_acc1 = vmlal_s16(aMag_acc1, vget_low_s16(a1_lo), vget_low_s16(a1_lo)); + aMag_acc1 = + vmlal_s16(aMag_acc1, vget_high_s16(a1_lo), vget_high_s16(a1_lo)); + aMag_acc2 = vmlal_s16(aMag_acc2, vget_low_s16(a1_hi), vget_low_s16(a1_hi)); + aMag_acc2 = + vmlal_s16(aMag_acc2, vget_high_s16(a1_hi), vget_high_s16(a1_hi)); + + bMag_acc1 = vmlal_s16(bMag_acc1, vget_low_s16(b1_lo), vget_low_s16(b1_lo)); + bMag_acc1 = + vmlal_s16(bMag_acc1, vget_high_s16(b1_lo), vget_high_s16(b1_lo)); + bMag_acc2 = vmlal_s16(bMag_acc2, vget_low_s16(b1_hi), vget_low_s16(b1_hi)); + bMag_acc2 = + vmlal_s16(bMag_acc2, vget_high_s16(b1_hi), vget_high_s16(b1_hi)); + + int8x16_t va2 = vld1q_s8(a + 16); + int8x16_t vb2 = vld1q_s8(b + 16); + int16x8_t a2_lo = vmovl_s8(vget_low_s8(va2)); + int16x8_t a2_hi = vmovl_s8(vget_high_s8(va2)); + int16x8_t b2_lo = vmovl_s8(vget_low_s8(vb2)); + int16x8_t b2_hi = vmovl_s8(vget_high_s8(vb2)); + + dot_acc1 = vmlal_s16(dot_acc1, vget_low_s16(a2_lo), vget_low_s16(b2_lo)); + dot_acc1 = vmlal_s16(dot_acc1, vget_high_s16(a2_lo), vget_high_s16(b2_lo)); + dot_acc2 = vmlal_s16(dot_acc2, vget_low_s16(a2_hi), vget_low_s16(b2_hi)); + dot_acc2 = vmlal_s16(dot_acc2, vget_high_s16(a2_hi), vget_high_s16(b2_hi)); + + aMag_acc1 = vmlal_s16(aMag_acc1, vget_low_s16(a2_lo), vget_low_s16(a2_lo)); + aMag_acc1 = + vmlal_s16(aMag_acc1, vget_high_s16(a2_lo), vget_high_s16(a2_lo)); + aMag_acc2 = vmlal_s16(aMag_acc2, vget_low_s16(a2_hi), vget_low_s16(a2_hi)); + aMag_acc2 = + vmlal_s16(aMag_acc2, vget_high_s16(a2_hi), vget_high_s16(a2_hi)); + + bMag_acc1 = vmlal_s16(bMag_acc1, vget_low_s16(b2_lo), vget_low_s16(b2_lo)); + bMag_acc1 = + vmlal_s16(bMag_acc1, vget_high_s16(b2_lo), vget_high_s16(b2_lo)); + 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)); + int16x8_t a_hi = vmovl_s8(vget_high_s8(va)); + int16x8_t b_lo = vmovl_s8(vget_low_s8(vb)); + int16x8_t b_hi = vmovl_s8(vget_high_s8(vb)); + + dot_acc1 = vmlal_s16(dot_acc1, vget_low_s16(a_lo), vget_low_s16(b_lo)); + dot_acc1 = vmlal_s16(dot_acc1, vget_high_s16(a_lo), vget_high_s16(b_lo)); + dot_acc1 = vmlal_s16(dot_acc1, vget_low_s16(a_hi), vget_low_s16(b_hi)); + dot_acc1 = vmlal_s16(dot_acc1, vget_high_s16(a_hi), vget_high_s16(b_hi)); + + aMag_acc1 = vmlal_s16(aMag_acc1, vget_low_s16(a_lo), vget_low_s16(a_lo)); + aMag_acc1 = vmlal_s16(aMag_acc1, vget_high_s16(a_lo), vget_high_s16(a_lo)); + aMag_acc1 = vmlal_s16(aMag_acc1, vget_low_s16(a_hi), vget_low_s16(a_hi)); + aMag_acc1 = vmlal_s16(aMag_acc1, vget_high_s16(a_hi), vget_high_s16(a_hi)); + + bMag_acc1 = vmlal_s16(bMag_acc1, vget_low_s16(b_lo), vget_low_s16(b_lo)); + bMag_acc1 = vmlal_s16(bMag_acc1, vget_high_s16(b_lo), vget_high_s16(b_lo)); + 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; + } + + 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; + aMag += (i32)*a * (i32)*a; + bMag += (i32)*b * (i32)*b; + a++; + b++; + } + + return 1 - ((f32)dot / (sqrt((f32)aMag) * sqrt((f32)bMag))); +} #endif static f32 l2_sqr_float(const void *pVect1v, const void *pVect2v, @@ -582,6 +812,11 @@ static f32 distance_cosine_bit(const void *pA, const void *pB, static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v, const void *qty_ptr) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)qty_ptr) > 16) { + return cosine_float_neon(pVect1v, pVect2v, qty_ptr); + } +#endif f32 *pVect1 = (f32 *)pVect1v; f32 *pVect2 = (f32 *)pVect2v; size_t qty = *((size_t *)qty_ptr); @@ -604,6 +839,11 @@ static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v, } static f32 distance_cosine_int8(const void *pA, const void *pB, const void *pD) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)pD) > 15) { + return cosine_int8_neon(pA, pB, pD); + } +#endif i8 *a = (i8 *)pA; i8 *b = (i8 *)pB; size_t d = *((size_t *)pD); diff --git a/tests/test-loadable.py b/tests/test-loadable.py index 0ef935c5..7a0b3830 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -417,7 +417,8 @@ def check(a, b, dtype: type = np.float32): b = np.array(b, dtype=dtype) x = vec_distance_cosine(a, b, a=transform, b=transform) - y = npy_cosine(a, b) + # the reference in float64: numpy sums int8 products in int8 + y = npy_cosine(a.astype(np.float64), b.astype(np.float64)) assert isclose(x, y, abs_tol=1e-6) check([1.2, 0.1], [0.4, -0.4]) @@ -425,6 +426,22 @@ def check(a, b, dtype: type = np.float32): check([1, 2, 3], [-9, -8, -7], dtype=np.int8) assert vec_distance_cosine("[1.1, 1.0]", "[1.2, 1.2]") == 0.0011318627512082458 + # lengths past the SIMD kernels' thresholds, with and without a remainder + # past their 16-element blocks + rng = np.random.default_rng(324) + for n in [17, 31, 32, 40, 127, 768]: + check(rng.uniform(-1, 1, n), rng.uniform(-1, 1, n)) + check(rng.integers(-128, 128, n), rng.integers(-128, 128, n), dtype=np.int8) + # 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, dtype=np.int8) + # 2^32 in total, which wraps a 32-bit sum to 0: + check([64] * 1048576, [-64] * 1048576, dtype=np.int8) + # 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, dtype=np.int8) + vec_distance_cosine_bit = lambda *args: db.execute( "select vec_distance_cosine(vec_bit(?), vec_bit(?))", args ).fetchone()[0] @@ -477,18 +494,22 @@ def test_vec_distance_cosine_extreme_magnitudes(): "select vec_distance_cosine(?, ?)", [_f32(a), _f32(b)] ).fetchone()[0] - for scale in [1e-45, 1e-30, 1e-25, 1e-20, 1.0, 1e20, 3e38]: - # cosine distance is invariant under positive scaling - assert vec_distance_cosine([1, 0], [scale, 0]) == 0.0 - assert vec_distance_cosine([scale, 0], [1, 0]) == 0.0 - assert vec_distance_cosine([scale, 0], [scale, 0]) == 0.0 - assert vec_distance_cosine([1, 0], [-scale, 0]) == 2.0 - assert vec_distance_cosine([1, 0], [0, scale]) == 1.0 - assert isclose( - vec_distance_cosine([1, 0], [scale, scale]), - 1 - 1 / np.sqrt(2), - abs_tol=1e-6, - ) + # the same 2-element vectors padded to 40 elements, past the SIMD kernels' + # threshold, with a remainder past their 16-element blocks + for pad in [0, 38]: + z = [0] * pad + for scale in [1e-45, 1e-30, 1e-25, 1e-20, 1.0, 1e20, 3e38]: + # cosine distance is invariant under positive scaling + assert vec_distance_cosine([1, 0] + z, [scale, 0] + z) == 0.0 + assert vec_distance_cosine([scale, 0] + z, [1, 0] + z) == 0.0 + assert vec_distance_cosine([scale, 0] + z, [scale, 0] + z) == 0.0 + assert vec_distance_cosine([1, 0] + z, [-scale, 0] + z) == 2.0 + assert vec_distance_cosine([1, 0] + z, [0, scale] + z) == 1.0 + assert isclose( + vec_distance_cosine([1, 0] + z, [scale, scale] + z), + 1 - 1 / np.sqrt(2), + abs_tol=1e-6, + ) def test_ensure_vector_match_cleanup_on_second_vector_error():