rewrite sub_to with avx2

Signed-off-by: Keming <kemingy94@gmail.com>

Author: kemingy

src/rabitq.rs

@@ -118,19 +118,25 @@ fn min_max_raw(vec: &DVectorView<f32>) -> (f32, f32) {
     (min, max)
 }
 
-// Interface of `min_max`
-fn min_max(vec: &DVectorView<f32>) -> (f32, f32) {
+// Interface of `min_max_residual`
+fn min_max_residual(
+    res: &mut DVector<f32>,
+    x: &DVectorView<f32>,
+    y: &DVectorView<f32>,
+) -> (f32, f32) {
     #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
     {
         if is_x86_feature_detected!("avx") {
-            unsafe { crate::simd::min_max_avx(vec) }
+            unsafe { crate::simd::min_max_residual_avx(res, x, y) }
         } else {
-            min_max_raw(vec)
+            x.sub_to(y, res);
+            min_max_raw(&res.as_view())
         }
     }
     #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
     {
-        min_max_raw(vec)
+        x.sub_to(y, &mut res);
+        min_max_raw(&res.as_view())
     }
 }
 
@@ -333,8 +339,11 @@ impl RaBitQ {
         let mut quantized = DVector::<u8>::zeros(self.dim as usize);
         let mut binary_vec = vec![0u64; query.len() * THETA_LOG_DIM as usize / 64];
         for &(dist, i) in lists[..length].iter() {
-            y_projected.sub_to(&self.centroids.column(i), &mut residual);
-            let (lower_bound, upper_bound) = min_max(&residual.as_view());
+            let (lower_bound, upper_bound) = min_max_residual(
+                &mut residual,
+                &y_projected.as_view(),
+                &self.centroids.column(i),
+            );
             let delta = (upper_bound - lower_bound) / ((1 << THETA_LOG_DIM) as f32 - 1.0);
             let one_over_delta = 1.0 / delta;
             let scalar_sum = scalar_quantize(

src/simd.rs

@@ -107,23 +107,35 @@ pub unsafe fn vector_binarize_query_avx2(vec: &DVectorView<u8>, binary: &mut [u6
 /// This function is marked unsafe because it requires the AVX intrinsics.
 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 #[target_feature(enable = "avx")]
-pub unsafe fn min_max_avx(vec: &DVectorView<f32>) -> (f32, f32) {
+pub unsafe fn min_max_residual_avx(
+    res: &mut DVector<f32>,
+    x: &DVectorView<f32>,
+    y: &DVectorView<f32>,
+) -> (f32, f32) {
     use std::arch::x86_64::*;
 
     let mut min_32x8 = _mm256_set1_ps(f32::MAX);
     let mut max_32x8 = _mm256_set1_ps(f32::MIN);
-    let mut ptr = vec.as_ptr();
+    let mut x_ptr = x.as_ptr();
+    let mut y_ptr = y.as_ptr();
+    let mut res_ptr = res.as_mut_ptr();
     let mut f32x8 = [0.0f32; 8];
     let mut min = f32::MAX;
     let mut max = f32::MIN;
-    let length = vec.len();
+    let length = res.len();
     let rest = length & 0b111;
+    let (mut x256, mut y256, mut res256);
 
     for _ in 0..(length / 8) {
-        let v = _mm256_loadu_ps(ptr);
-        ptr = ptr.add(8);
-        min_32x8 = _mm256_min_ps(min_32x8, v);
-        max_32x8 = _mm256_max_ps(max_32x8, v);
+        x256 = _mm256_loadu_ps(x_ptr);
+        y256 = _mm256_loadu_ps(y_ptr);
+        res256 = _mm256_sub_ps(x256, y256);
+        _mm256_storeu_ps(res_ptr, res256);
+        x_ptr = x_ptr.add(8);
+        y_ptr = y_ptr.add(8);
+        res_ptr = res_ptr.add(8);
+        min_32x8 = _mm256_min_ps(min_32x8, res256);
+        max_32x8 = _mm256_max_ps(max_32x8, res256);
     }
     _mm256_storeu_ps(f32x8.as_mut_ptr(), min_32x8);
     for &x in f32x8.iter() {
@@ -139,23 +151,30 @@ pub unsafe fn min_max_avx(vec: &DVectorView<f32>) -> (f32, f32) {
     }
 
     for _ in 0..rest {
-        if *ptr < min {
-            min = *ptr;
+        *res_ptr = *x_ptr - *y_ptr;
+        if *res_ptr < min {
+            min = *res_ptr;
         }
-        if *ptr > max {
-            max = *ptr;
+        if *res_ptr > max {
+            max = *res_ptr;
         }
-        ptr = ptr.add(1);
+        res_ptr = res_ptr.add(1);
+        x_ptr = x_ptr.add(1);
+        y_ptr = y_ptr.add(1);
     }
 
     (min, max)
 }
 
 /// Compute the u8 scalar quantization of a f32 vector.
 ///
+/// This function doesn't need `bias` because it *round* the f32 to u32 instead of *floor*.
+///
 /// # Safety
 ///
 /// This function is marked unsafe because it requires the AVX intrinsics.
+#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+#[target_feature(enable = "avx2")]
 pub unsafe fn scalar_quantize_avx2(
     quantized: &mut DVector<u8>,
     vec: &DVectorView<f32>,