diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..89b48d40 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -119,6 +119,55 @@ static f32 l2_sqr_float_avx(const void *pVect1v, const void *pVect2v, return sqrt(TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7]); } + +#ifdef __FMA__ +#define VEC_AVX_MULADD(a, b, acc) _mm256_fmadd_ps((a), (b), (acc)) +#else +#define VEC_AVX_MULADD(a, b, acc) _mm256_add_ps(_mm256_mul_ps((a), (b)), (acc)) +#endif + +static f32 hsum_float_avx(__m256 v) { + __m128 s = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); + s = _mm_add_ps(s, _mm_movehl_ps(s, s)); + s = _mm_add_ss(s, _mm_shuffle_ps(s, s, 1)); + return _mm_cvtss_f32(s); +} + +// Cosine distance, 16 floats per step in two independent accumulator sets +// (dot, |a|^2, |b|^2) so the loop is not bound by FP-add latency. +static f32 cosine_float_avx(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + const f32 *a = (const f32 *)pVect1v; + const f32 *b = (const f32 *)pVect2v; + size_t qty = *((const size_t *)qty_ptr); + const f32 *aEnd = a + (qty & ~(size_t)15); + + __m256 dot0 = _mm256_setzero_ps(), dot1 = _mm256_setzero_ps(); + __m256 aa0 = _mm256_setzero_ps(), aa1 = _mm256_setzero_ps(); + __m256 bb0 = _mm256_setzero_ps(), bb1 = _mm256_setzero_ps(); + while (a < aEnd) { + __m256 a0 = _mm256_loadu_ps(a), a1 = _mm256_loadu_ps(a + 8); + __m256 b0 = _mm256_loadu_ps(b), b1 = _mm256_loadu_ps(b + 8); + dot0 = VEC_AVX_MULADD(a0, b0, dot0); + dot1 = VEC_AVX_MULADD(a1, b1, dot1); + aa0 = VEC_AVX_MULADD(a0, a0, aa0); + aa1 = VEC_AVX_MULADD(a1, a1, aa1); + bb0 = VEC_AVX_MULADD(b0, b0, bb0); + bb1 = VEC_AVX_MULADD(b1, b1, bb1); + a += 16; + b += 16; + } + f32 dot = hsum_float_avx(_mm256_add_ps(dot0, dot1)); + f32 aMag = hsum_float_avx(_mm256_add_ps(aa0, aa1)); + f32 bMag = hsum_float_avx(_mm256_add_ps(bb0, bb1)); + for (size_t i = qty & ~(size_t)15; i < qty; i++, a++, b++) { + dot += *a * *b; + aMag += *a * *a; + bMag += *b * *b; + } + return 1 - (dot / (sqrt(aMag) * sqrt(bMag))); +} +#undef VEC_AVX_MULADD #endif #ifdef SQLITE_VEC_ENABLE_NEON @@ -487,6 +536,11 @@ static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v, if ((*(const size_t *)qty_ptr) > 16) { return cosine_float_neon(pVect1v, pVect2v, qty_ptr); } +#endif +#ifdef SQLITE_VEC_ENABLE_AVX + if ((*(const size_t *)qty_ptr) >= 16) { + return cosine_float_avx(pVect1v, pVect2v, qty_ptr); + } #endif f32 *pVect1 = (f32 *)pVect1v; f32 *pVect2 = (f32 *)pVect2v; diff --git a/tests/test-unit.c b/tests/test-unit.c index 83cedd5f..010ef8e4 100644 --- a/tests/test-unit.c +++ b/tests/test-unit.c @@ -870,6 +870,29 @@ void test_distance_cosine_float() { assert(fabsf(d - 2.0f) < 1e-6f); } + // Lengths around the 16-float SIMD block, including odd tails, against + // a scalar reference summed in double. + { + size_t dims_list[] = {15, 16, 17, 31, 33, 47, 384, 389}; + float a[389], b[389]; + for (size_t i = 0; i < countof(a); i++) { + a[i] = sinf((float)i * 0.37f + 0.1f); + b[i] = cosf((float)i * 0.23f - 0.4f); + } + for (size_t t = 0; t < countof(dims_list); t++) { + size_t n = dims_list[t]; + double dot = 0, aa = 0, bb = 0; + for (size_t i = 0; i < n; i++) { + dot += (double)a[i] * b[i]; + aa += (double)a[i] * a[i]; + bb += (double)b[i] * b[i]; + } + float expected = (float)(1.0 - dot / (sqrt(aa) * sqrt(bb))); + d = _test_distance_cosine_float(a, b, n); + assert(fabsf(d - expected) < 1e-5f); + } + } + printf(" All distance_cosine_float tests passed.\n"); }