feat: hand-writing avx2

Signed-off-by: usamoi <[email protected]>

Author: usamoi

crates/c/.gitignore

@@ -0,0 +1,2 @@
+*.s
+*.o

crates/c/src/c.c

@@ -25,8 +25,8 @@ vectors_f16_cosine_axv512(_Float16 const *restrict a,
     xx = _mm512_fmadd_ph(x, x, xx);
     yy = _mm512_fmadd_ph(y, y, yy);
   }
-  return (float)(_mm512_reduce_add_ps(xy) /
-                 sqrt(_mm512_reduce_add_ps(xx) * _mm512_reduce_add_ps(yy)));
+  return (float)(_mm512_reduce_add_ph(xy) /
+                 sqrt(_mm512_reduce_add_ph(xx) * _mm512_reduce_add_ph(yy)));
 }
 
 __attribute__((target("avx512fp16,bmi2"))) extern float
@@ -71,3 +71,47 @@ vectors_f16_distance_squared_l2_axv512(_Float16 const *restrict a,
 
   return (float)_mm512_reduce_add_ph(dd);
 }
+
+__attribute__((target("avx2"))) extern float
+vectors_f16_cosine_axv2(_Float16 const *restrict a, _Float16 const *restrict b,
+                        size_t n) {
+  float xy = 0;
+  float xx = 0;
+  float yy = 0;
+#pragma clang loop vectorize_width(8)
+  for (size_t i = 0; i < n; i++) {
+    float x = a[i];
+    float y = b[i];
+    xy += x * y;
+    xx += x * x;
+    yy += y * y;
+  }
+  return xy / sqrt(xx * yy);
+}
+
+__attribute__((target("avx2"))) extern float
+vectors_f16_dot_axv2(_Float16 const *restrict a, _Float16 const *restrict b,
+                     size_t n) {
+  float xy = 0;
+#pragma clang loop vectorize_width(8)
+  for (size_t i = 0; i < n; i++) {
+    float x = a[i];
+    float y = b[i];
+    xy += x * y;
+  }
+  return xy;
+}
+
+__attribute__((target("avx2"))) extern float
+vectors_f16_distance_squared_l2_axv2(_Float16 const *restrict a,
+                                     _Float16 const *restrict b, size_t n) {
+  float dd = 0;
+#pragma clang loop vectorize_width(8)
+  for (size_t i = 0; i < n; i++) {
+    float x = a[i];
+    float y = b[i];
+    float d = x - y;
+    dd += d * d;
+  }
+  return dd;
+}

crates/c/src/c.h

@@ -7,3 +7,9 @@ extern float vectors_f16_dot_axv512(_Float16 const *, _Float16 const *,
                                     size_t n);
 extern float vectors_f16_distance_squared_l2_axv512(_Float16 const *,
                                                     _Float16 const *, size_t n);
+
+extern float vectors_f16_cosine_axv2(_Float16 const *, _Float16 const *,
+                                     size_t n);
+extern float vectors_f16_dot_axv2(_Float16 const *, _Float16 const *, size_t n);
+extern float vectors_f16_distance_squared_l2_axv2(_Float16 const *,
+                                                  _Float16 const *, size_t n);

crates/c/src/c.rs

@@ -3,4 +3,7 @@ extern "C" {
     pub fn vectors_f16_cosine_axv512(a: *const u16, b: *const u16, n: usize) -> f32;
     pub fn vectors_f16_dot_axv512(a: *const u16, b: *const u16, n: usize) -> f32;
     pub fn vectors_f16_distance_squared_l2_axv512(a: *const u16, b: *const u16, n: usize) -> f32;
+    pub fn vectors_f16_cosine_axv2(a: *const u16, b: *const u16, n: usize) -> f32;
+    pub fn vectors_f16_dot_axv2(a: *const u16, b: *const u16, n: usize) -> f32;
+    pub fn vectors_f16_distance_squared_l2_axv2(a: *const u16, b: *const u16, n: usize) -> f32;
 }

crates/service/src/prelude/global/avx2.rs

@@ -0,0 +1,9 @@
+#[cfg(not(target_arch = "x86_64"))]
+pub fn detect() -> bool {
+    false
+}
+
+#[cfg(target_arch = "x86_64")]
+pub fn detect() -> bool {
+    std_detect::is_x86_feature_detected!("avx2")
+}

crates/service/src/prelude/global/f16.rs

@@ -20,7 +20,15 @@ pub fn cosine(lhs: &[F16], rhs: &[F16]) -> F32 {
         unsafe {
             assert!(lhs.len() == rhs.len());
             let n = lhs.len();
-            c::vectors_f16_cosine_axv512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n);
+            return c::vectors_f16_cosine_axv512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n)
+                .into();
+        }
+    }
+    if super::avx2::detect() {
+        unsafe {
+            assert!(lhs.len() == rhs.len());
+            let n = lhs.len();
+            return c::vectors_f16_cosine_axv2(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n).into();
         }
     }
     cosine(lhs, rhs)
@@ -42,7 +50,14 @@ pub fn dot(lhs: &[F16], rhs: &[F16]) -> F32 {
         unsafe {
             assert!(lhs.len() == rhs.len());
             let n = lhs.len();
-            c::vectors_f16_dot_axv512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n);
+            return c::vectors_f16_dot_axv512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n).into();
+        }
+    }
+    if super::avx2::detect() {
+        unsafe {
+            assert!(lhs.len() == rhs.len());
+            let n = lhs.len();
+            return c::vectors_f16_dot_axv2(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n).into();
         }
     }
     cosine(lhs, rhs)
@@ -65,7 +80,24 @@ pub fn distance_squared_l2(lhs: &[F16], rhs: &[F16]) -> F32 {
         unsafe {
             assert!(lhs.len() == rhs.len());
             let n = lhs.len();
-            c::vectors_f16_distance_squared_l2_axv512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), n);
+            return c::vectors_f16_distance_squared_l2_axv512(
+                lhs.as_ptr().cast(),
+                rhs.as_ptr().cast(),
+                n,
+            )
+            .into();
+        }
+    }
+    if super::avx2::detect() {
+        unsafe {
+            assert!(lhs.len() == rhs.len());
+            let n = lhs.len();
+            return c::vectors_f16_distance_squared_l2_axv2(
+                lhs.as_ptr().cast(),
+                rhs.as_ptr().cast(),
+                n,
+            )
+            .into();
         }
     }
     distance_squared_l2(lhs, rhs)

crates/service/src/prelude/global/mod.rs

@@ -1,3 +1,4 @@
+mod avx2;
 mod avx512fp16;
 mod f16;
 mod f16_cos;