fix(sqlite-vec.c): keep the NEON int8 cosine sums in 64 bits

#5's cosine_int8_neon() kept its sums in 32 bits, so on NEON builds
vec_distance_cosine() returned NULL, 1.0, or a wrong distance such as 3.8
for int8 vectors whose squares sum past 2^31 - 1, from 131,072 elements.
Its CHANGELOG entry also overstated how closely the NEON float32 and int8
distances match the scalar code's.

Author: mceachen

CHANGELOG.md

@@ -6,7 +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`.
+- 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 can differ from the scalar code's in the last digits, as NEON adds the same double-precision products in a different order: differences are rare between random vectors and common between near-duplicate ones, and a vector's distance to itself is between -2.2e-16 and 2.2e-16 in either. `int8` distances match up to 1,024 dimensions; past that, vectors whose squared elements sum past 2^24 can get different distances, usually 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)), where asg017's kernel sums them in `float32`, and `int8` sums are kept in 64 bits, where asg017's 32-bit sums overflow from 131,072 elements.
 
 ### Fixed
 

sqlite-vec.c

@@ -540,21 +540,27 @@ static f32 cosine_float_neon(const void *pVect1v, const void *pVect2v,
 }
 
 // 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.
+// vector case. 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, as in
+// l2_sqr_int8_neon.
 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));
@@ -607,11 +613,17 @@ static f32 cosine_int8_neon(const void *pA, const void *pB, const void *pD) {
     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));
@@ -634,17 +646,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;

tests/test-loadable.py

@@ -486,6 +486,15 @@ def check(a, b, dtype=np.float32):
     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