From 92258b0f1c1d7c5cd8708d643133251074420ece Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 19:09:30 -0500 Subject: [PATCH] Sum cosine distances and vec_normalize() in double distance_cosine_float(), cosine_float_neon() and vec_normalize() summed squared elements in float32. The square of an element below about 3e-23 underflows to zero and the square of one above about 2e19 overflows to Inf, so a nonzero finite vector got a cosine distance of +-Inf, NaN or 1.0, ranked ahead of or behind its true neighbours in a cosine vec0 table, and vec_normalize() returned [Inf, NaN] or a vector that was not of unit length (#324). The scalar loops now sum in double, where every product of two float32 values is representable, as vlasky/sqlite-vec@d75c245 does in that fork; its tests come with it. cosine_float_neon(), which that fork lacks and which fails the same way above 16 dimensions, keeps its 16-element blocks and four accumulators per sum but widens each float32x4 load to two float64x2 halves before it is multiplied and added, and finishes as the scalar loop does. A zero vector gets the same results as before: NULL from vec_distance_cosine() and NaN elements from vec_normalize(). On Linux arm64 (Clang 19, -O3; medians of 20 queries on a 100,000-row float[768] cosine table, each build in its own process), a KNN query takes 77.5 ms instead of 72.0 ms for k = 10 and 84.5 ms instead of 78.1 ms for k = 100 on a scalar build, and 25.3 ms instead of 18.3 ms and 32.0 ms instead of 26.8 ms on a NEON build, whose float64x2 lanes hold half as many elements as float32x4. The queries return the same rows at every rank. Their distances change in the last digits, and the NEON build's now equal the scalar build's on every row, where on main they differed; test_vec_distance_cosine expects the new value of one. The tests check the cases from #324 on the scalar path and, padded to 40 dimensions, on the NEON path: vec_distance_cosine() under positive scaling from 1e-45 to 3e38, vec_normalize() of tiny and huge vectors, and the KNN order of a cosine table holding them, along with lengths of 17 to 768 against a float64 reference. All four tests fail on main with a scalar and a NEON build. The suite passes on GCC and Clang scalar builds and a Clang NEON build (482 passed, 132 skipped); GCC builds the NEON code once #322 is in. --- sqlite-vec.c | 109 ++++++++++++++++++++++++++++------------- tests/test-loadable.py | 89 ++++++++++++++++++++++++++++++++- 2 files changed, 161 insertions(+), 37 deletions(-) diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..1a9ab762 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -185,6 +185,9 @@ static f32 l2_sqr_float_neon(const void *pVect1v, const void *pVect2v, return sqrt(sum_scalar); } +// Sums are 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; @@ -193,53 +196,83 @@ static f32 cosine_float_neon(const void *pVect1v, const void *pVect2v, size_t qty16 = qty >> 4; const f32 *pEnd1 = pVect1 + (qty16 << 4); - float32x4_t dot0 = vdupq_n_f32(0), dot1 = vdupq_n_f32(0); - float32x4_t dot2 = vdupq_n_f32(0), dot3 = vdupq_n_f32(0); - float32x4_t amag0 = vdupq_n_f32(0), amag1 = vdupq_n_f32(0); - float32x4_t amag2 = vdupq_n_f32(0), amag3 = vdupq_n_f32(0); - float32x4_t bmag0 = vdupq_n_f32(0), bmag1 = vdupq_n_f32(0); - float32x4_t bmag2 = vdupq_n_f32(0), bmag3 = vdupq_n_f32(0); + 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; - dot0 = vfmaq_f32(dot0, v1, v2); - amag0 = vfmaq_f32(amag0, v1, v1); - bmag0 = vfmaq_f32(bmag0, v2, v2); + 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; - dot1 = vfmaq_f32(dot1, v1, v2); - amag1 = vfmaq_f32(amag1, v1, v1); - bmag1 = vfmaq_f32(bmag1, v2, v2); + 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; - dot2 = vfmaq_f32(dot2, v1, v2); - amag2 = vfmaq_f32(amag2, v1, v1); - bmag2 = vfmaq_f32(bmag2, v2, v2); + 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; - dot3 = vfmaq_f32(dot3, v1, v2); - amag3 = vfmaq_f32(amag3, v1, v1); - bmag3 = vfmaq_f32(bmag3, v2, v2); - } - - f32 dot_s = vaddvq_f32(vaddq_f32(vaddq_f32(dot0, dot1), vaddq_f32(dot2, dot3))); - f32 amag_s = vaddvq_f32(vaddq_f32(vaddq_f32(amag0, amag1), vaddq_f32(amag2, amag3))); - f32 bmag_s = vaddvq_f32(vaddq_f32(vaddq_f32(bmag0, bmag1), vaddq_f32(bmag2, bmag3))); + 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_s = vaddvq_f64(vaddq_f64(vaddq_f64(dot0, dot1), vaddq_f64(dot2, dot3))); + double amag_s = vaddvq_f64(vaddq_f64(vaddq_f64(amag0, amag1), vaddq_f64(amag2, amag3))); + double bmag_s = vaddvq_f64(vaddq_f64(vaddq_f64(bmag0, bmag1), vaddq_f64(bmag2, bmag3))); const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4)); while (pVect1 < pEnd2) { - dot_s += *pVect1 * *pVect2; - amag_s += *pVect1 * *pVect1; - bmag_s += *pVect2 * *pVect2; + dot_s += (double)*pVect1 * (double)*pVect2; + amag_s += (double)*pVect1 * (double)*pVect1; + bmag_s += (double)*pVect2 * (double)*pVect2; pVect1++; pVect2++; } - return 1.0f - (dot_s / (sqrtf(amag_s) * sqrtf(bmag_s))); + return 1 - (dot_s / (sqrt(amag_s) * sqrt(bmag_s))); } static f32 l2_sqr_int8_neon(const void *pVect1v, const void *pVect2v, @@ -492,13 +525,17 @@ static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v, f32 *pVect2 = (f32 *)pVect2v; size_t qty = *((size_t *)qty_ptr); - f32 dot = 0; - f32 aMag = 0; - f32 bMag = 0; + // Accumulate in double: the square of a small (or large) f32 element + // underflows (or overflows) f32, which made the magnitude 0 (or Inf) and + // the distance +-Inf, NaN or plain wrong for nonzero finite vectors. Every + // f32 product is representable in double. + double dot = 0; + double aMag = 0; + double bMag = 0; for (size_t i = 0; i < qty; i++) { - dot += *pVect1 * *pVect2; - aMag += *pVect1 * *pVect1; - bMag += *pVect2 * *pVect2; + dot += (double)*pVect1 * (double)*pVect2; + aMag += (double)*pVect1 * (double)*pVect1; + bMag += (double)*pVect2 * (double)*pVect2; pVect1++; pVect2++; } @@ -2047,9 +2084,11 @@ static void vec_normalize(sqlite3_context *context, int argc, f32 *v = (f32 *)vector; - f32 norm = 0; + // Accumulate in double so that squaring small or large f32 elements does + // not underflow or overflow, see distance_cosine_float(). + double norm = 0; for (size_t i = 0; i < dimensions; i++) { - norm += v[i] * v[i]; + norm += (double)v[i] * (double)v[i]; } norm = sqrt(norm); for (size_t i = 0; i < dimensions; i++) { diff --git a/tests/test-loadable.py b/tests/test-loadable.py index 0044144b..28ce010a 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -273,13 +273,48 @@ 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.001131898257881403 + 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) + + +def test_vec_distance_cosine_extreme_magnitudes(): + # squaring very small or very large f32 elements used to underflow or + # overflow the f32 magnitude accumulators, returning +-Inf, NaN or a wrong + # distance for nonzero finite vectors + # https://github.com/asg017/sqlite-vec/issues/324 + vec_distance_cosine = lambda a, b: db.execute( + "select vec_distance_cosine(?, ?)", [_f32(a), _f32(b)] + ).fetchone()[0] + + # 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_vec_distance_hamming(): @@ -482,6 +517,18 @@ def test_vec_normalize(): -0.6324555277824402, ] + # small and large nonzero vectors are normalized rather than treated as + # zero-magnitude or returned as a zero vector, the squared magnitude used + # to underflow/overflow f32 + # https://github.com/asg017/sqlite-vec/issues/324 + for scale in [1e-45, 1e-30, 1e-25, 1e-22, 1e-20, 1e20, 3e38]: + assert vec_normalize(_f32([scale, 0])) == _f32([1, 0]) + assert vec_normalize(_f32([0, -scale])) == _f32([0, -1]) + assert np.allclose( + np.frombuffer(vec_normalize(_f32([scale, scale])), dtype=np.float32), + [1 / np.sqrt(2)] * 2, + ) + def test_vec_slice(): vec_slice = lambda *args, f="?": db.execute( @@ -2183,6 +2230,44 @@ def test_vec0_distance_metric(): ] +def test_vec0_cosine_extreme_magnitudes(): + # tiny or huge nonzero vectors used to get +-Inf, NaN or wrong cosine + # distances, which ranked them ahead of (or behind) the true nearest + # neighbors + # https://github.com/asg017/sqlite-vec/issues/324 + db = connect(EXT_PATH) + db.execute("create virtual table v using vec0(a float[2] distance_metric=cosine)") + db.executemany( + "insert into v(rowid, a) values (?, ?)", + [ + (1, _f32([0.9, 0.1])), + (2, _f32([1e-25, 0])), + (3, _f32([1e-25, 1e-25])), + (4, _f32([-1e-25, 0])), + (5, _f32([1e-30, 0])), + (6, _f32([3e38, 3e38])), + ], + ) + + expected = [ + {"rowid": 2, "distance": 0.0}, + {"rowid": 5, "distance": 0.0}, + {"rowid": 1, "distance": 0.006116265896707773}, + {"rowid": 3, "distance": 0.2928932309150696}, + {"rowid": 6, "distance": 0.2928932309150696}, + {"rowid": 4, "distance": 2.0}, + ] + for query in ([1, 0], [1e-25, 0], [3e38, 0]): + rows = execute_all( + db, + "select rowid, distance from v where a match ? and k = 6", + [_f32(query)], + ) + # the order of equidistant rows is unspecified + assert sorted(rows, key=lambda r: (r["distance"], r["rowid"])) == expected + assert [r["distance"] for r in rows] == [r["distance"] for r in expected] + + def test_vec0_vacuum(): db = connect(EXT_PATH) db.execute("create virtual table vec_t using vec0(a float[1]);")