From 6df0e5039f5f2df694cbee9b5c32422200c1f827 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 18:40:26 -0500 Subject: [PATCH] perf(sqlite-vec.c): 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 darwin-arm64 binary, 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 two changes. asg017's float32 kernel sums in float32, the loss asg017#324 reported and 2.1.0 fixed in the scalar loop by summing in double: 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. And both kernels return 1.0 for a zero vector, as the scalar loops do, where asg017's return NaN. On Linux arm64 (GCC 14, -O3), one 768-dimension float32 distance takes about 155 ns instead of 895 ns, and a KNN query on a 100,000-row float[768] table takes 25.9 ms instead of 94.4 ms for k = 10 and 30.1 ms instead of 98.9 ms for k = 100 (medians of 20 queries), returning the same rows. The float32 distances match the scalar loop's on 12,000 random pairs of 16 to 1,536 dimensions, except that a vector's distance to itself, 0 or about 2e-16 in either, can change between the two. The int8 kernel sums products exactly in 32-bit lanes, and the scalar loop sums them in float, which is exact while a sum stays below 2^24: their distances match up to 1,024 dimensions, and past that the kernel's are closer to the exact value for vectors whose squares sum past 2^24 (a mean error of 2.3e-8 instead of 2.5e-7 at 2,048 dimensions). The cosine tests used vectors of 2 and 3 elements, below both kernels' thresholds, and pass against asg017's kernels as they are. 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; the zero-vector and extreme-magnitude tests repeat at 40 dimensions. asg017's float32 kernel fails both, and with a zero-vector check added it still fails the extreme-magnitude test. The full suite passes on the NEON build and on the scalar one, and make test-unit, which builds without SIMD, passes too. --- CHANGELOG.md | 1 + sqlite-vec.c | 238 +++++++++++++++++++++++++++++++++++++++++ tests/test-loadable.py | 45 +++++--- 3 files changed, 271 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index da5f90d1..731bbcff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Changed - Reading a vector column without a KNN query, as in `select embedding from v`, reads each chunk's vectors through one blob handle instead of opening the chunk's blob again for every row: selecting every vector of a 10,000-row `float[768]` table took 25 ms instead of 2,080 ms, and `select vec_distance_cosine(embedding, ?) from v` 41 ms instead of 2,093 ms. +- Cosine distance between `float32` vectors of more than 16 dimensions, or `int8` vectors of more than 15, uses NEON in NEON builds, which include the darwin-arm64 binary: a KNN query on a 100,000-row `float[768]` table took 26 ms instead of 94 ms on Linux arm64. `float32` distances match the scalar code's on random vectors, though a vector's distance to an identical vector, 0 or about 2e-16 in either, can change between the two. `int8` distances match up to 1,024 dimensions; past that, vectors whose squared elements sum past 2^24 get distances closer to the exact value, as NEON sums them exactly. Ported from [asg017@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f), except that `float32` sums stay in `double` ([asg017#324](https://github.com/asg017/sqlite-vec/issues/324)); asg017's kernel sums them in `float32`. ### Fixed diff --git a/sqlite-vec.c b/sqlite-vec.c index 429b475d..218acb06 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -432,6 +432,234 @@ 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++; + } + + // Handle zero vectors: return max distance (1.0) to avoid division by zero + if (aMag == 0 || bMag == 0) { + return 1.0f; + } + return 1 - (dot / (sqrt(aMag) * sqrt(bMag))); +} + +// asg017/sqlite-vec's cosine_int8_neon, with distance_cosine_int8's zero +// vector case. Products of int8 elements are summed exactly in 32-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); + + while (a < aEnd - 31) { + 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)); + + a += 32; + b += 32; + } + + while (a < aEnd - 15) { + 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)); + + 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); + + while (a < aEnd) { + dot += (i32)*a * (i32)*b; + aMag += (i32)*a * (i32)*a; + bMag += (i32)*b * (i32)*b; + a++; + b++; + } + + // Handle zero vectors: return max distance (1.0) to avoid division by zero + if (aMag == 0 || bMag == 0) { + return 1.0f; + } + return 1.0f - ((f32)dot / (sqrtf((f32)aMag) * sqrtf((f32)bMag))); +} #endif static f32 l2_sqr_float(const void *pVect1v, const void *pVect2v, @@ -622,6 +850,11 @@ static f32 distance_cosine_bit(const void *pA, const void *pB, const void *pD) { 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); @@ -648,6 +881,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 d7776f50..67fdad31 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -471,7 +471,8 @@ def check(a, b, dtype=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]) @@ -479,6 +480,13 @@ def check(a, b, dtype=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) + vec_distance_cosine_bit = lambda *args: db.execute( "select vec_distance_cosine(vec_bit(?), vec_bit(?))", args ).fetchone()[0] @@ -524,6 +532,13 @@ def test_vec_distance_cosine_zero_vector(): assert vec_distance_cosine_int8("[0,0,0]", "[1,2,3]") == 1.0 assert vec_distance_cosine_int8("[1,2,3]", "[0,0,0]") == 1.0 + # the same past the SIMD kernels' thresholds + zero, ones = json.dumps([0] * 40), json.dumps([1] * 40) + assert vec_distance_cosine(zero, ones) == 1.0 + assert vec_distance_cosine(zero, zero) == 1.0 + assert vec_distance_cosine_int8(zero, ones) == 1.0 + assert vec_distance_cosine_int8(ones, zero) == 1.0 + # bit zero vectors vec_distance_cosine_bit = lambda a, b: db.execute( "select vec_distance_cosine(vec_bit(?), vec_bit(?))", [a, b] @@ -542,18 +557,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():