Skip to content

Use NEON for cosine distance #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
MayCXC wants to merge 1 commit into from
Closed

Use NEON for cosine distance #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
238 changes: 238 additions & 0 deletions sqlite-vec.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down
45 changes: 32 additions & 13 deletions tests/test-loadable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand All @@ -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():
Expand Down