Add an AVX cosine distance for float vectors

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;

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");
 }