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.

Author: MayCXC

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 - 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.
 - The linux-arm64 and linux-arm64-musl binaries use NEON, as the darwin-arm64 binary does: `make` enables it on Linux arm64, and `OMIT_SIMD=1` builds without it. An L2 KNN query on a 100,000-row `float[768]` table took 19 ms instead of 57 ms there. Hamming and `int8` L1 distances are unchanged; `float32` L1 distances can differ in the last digit when their double-precision sum rounds, and `float32` L2 distances can differ from 2.1.1's in the last digits, closer to the exact value (at 768 dimensions, a mean relative error of 3.0e-8 instead of 1.7e-7), so a KNN query can order rows with nearly equal distances differently; `int8` L2 distances whose squared differences sum past 2^24 are computed from the exact sum.
+- 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
 

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

tests/test-loadable.py

@@ -471,14 +471,22 @@ 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])
     check([-1.2, -0.1], [-0.4, 0.4])
     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():