Skip to content

Keep the NEON int8 cosine sums in 64 bits #356

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

Open
MayCXC wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Keep the NEON int8 cosine sums in 64 bits #356

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
37 changes: 24 additions & 13 deletions sqlite-vec.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -523,20 +523,26 @@ static f32 cosine_int8(const void *pA, const void *pB, const void *pD) {
}

#ifdef SQLITE_VEC_ENABLE_NEON
// 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.
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));
Expand Down Expand Up @@ -581,11 +587,17 @@ static f32 cosine_int8_neon(const void *pA, const void *pB, const void *pD) {
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));

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));
Expand All @@ -608,17 +620,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;
Expand Down
23 changes: 23 additions & 0 deletions tests/test-loadable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,29 @@ def npy_l2(a, b):
return np.linalg.norm(a - b)


def test_vec_distance_cosine_int8_overflow():
vec_distance_cosine = lambda a, b: db.execute(
"select vec_distance_cosine(vec_int8(?), vec_int8(?))", [a, b]
).fetchone()[0]

def check(a, b):
a = np.array(a, dtype=np.int8)
b = np.array(b, dtype=np.int8)
# the reference in float64: numpy sums int8 products in int8
y = npy_cosine(a.astype(np.float64), b.astype(np.float64))
assert isclose(vec_distance_cosine(a, b), y, abs_tol=1e-6)

# 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)
# 2^32 in total, which wraps a 32-bit sum to 0:
check([64] * 1048576, [-64] * 1048576)
# 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)


def test_vec_distance_cosine():
vec_distance_cosine = lambda *args, a="?", b="?": db.execute(
f"select vec_distance_cosine({a}, {b})", args
Expand Down