rewrite min/max with simd

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

Author: kemingy

src/main.rs

@@ -69,7 +69,8 @@ fn main() {
         let start_time = Instant::now();
         let res = rabitq.query_one(&query_vec, args.probe, args.topk);
         total_time += start_time.elapsed().as_secs_f64();
-        recall += calculate_recall(&truth[i], &res, args.topk);
+        let ids: Vec<i32> = res.iter().map(|(_, id)| *id as i32).collect();
+        recall += calculate_recall(&truth[i], &ids, args.topk);
     }
 
     info!(

src/rabitq.rs

@@ -104,7 +104,7 @@ fn l2_squared_distance(
 }
 
 // Get the min/max value of a vector.
-fn min_max(vec: &DVectorView<f32>) -> (f32, f32) {
+fn min_max_raw(vec: &DVectorView<f32>) -> (f32, f32) {
     let mut min = f32::MAX;
     let mut max = f32::MIN;
     for v in vec.iter() {
@@ -118,6 +118,22 @@ fn min_max(vec: &DVectorView<f32>) -> (f32, f32) {
     (min, max)
 }
 
+// Interface of `min_max`
+fn min_max(vec: &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) }
+        } else {
+            min_max_raw(vec)
+        }
+    }
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
+    {
+        min_max_raw(vec)
+    }
+}
+
 // Quantize the query residual vector.
 fn quantize_query_vector(
     quantized: &mut DVector<u8>,
@@ -279,7 +295,7 @@ impl RaBitQ {
     }
 
     /// Query the topk nearest neighbors for the given query.
-    pub fn query_one(&self, query: &DVector<f32>, probe: usize, topk: usize) -> Vec<i32> {
+    pub fn query_one(&self, query: &DVector<f32>, probe: usize, topk: usize) -> Vec<(f32, u32)> {
         let y_projected = query.tr_mul(&self.orthogonal).transpose();
         let k = self.centroids.shape().1;
         let mut lists = Vec::with_capacity(k);
@@ -337,7 +353,7 @@ impl RaBitQ {
         query: &DVector<f32>,
         rough_distances: &[(f32, u32)],
         topk: usize,
-    ) -> Vec<i32> {
+    ) -> Vec<(f32, u32)> {
         let mut threshold = f32::MAX;
         let mut recent_max_accurate = f32::MIN;
         let mut res = Vec::with_capacity(topk);
@@ -351,7 +367,7 @@ impl RaBitQ {
                     &mut residual,
                 );
                 if accurate < threshold {
-                    res.push((accurate, u as i32));
+                    res.push((accurate, u));
                     count += 1;
                     recent_max_accurate = recent_max_accurate.max(accurate);
                     if count == WINDOWS_SIZE {
@@ -366,6 +382,7 @@ impl RaBitQ {
         METRICS.add_precise_count(res.len() as u64);
         let length = topk.min(res.len());
         res.select_nth_unstable_by(length - 1, |a, b| a.0.total_cmp(&b.0));
-        res[..length].iter().map(|(_, u)| *u).collect()
+        res.truncate(length);
+        res
     }
 }

src/simd.rs

@@ -99,3 +99,54 @@ pub unsafe fn vector_binarize_query_avx2(vec: &DVectorView<u8>, binary: &mut [u6
         }
     }
 }
+
+/// Compute the min and max value of a vector.
+///
+/// # Safety
+///
+/// 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) {
+    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 f32x8 = [0.0f32; 8];
+    let mut min = f32::MAX;
+    let mut max = f32::MIN;
+    let length = vec.len();
+    let rest = length & 0b111;
+
+    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);
+    }
+    _mm256_storeu_ps(f32x8.as_mut_ptr(), min_32x8);
+    for &x in f32x8.iter() {
+        if x < min {
+            min = x;
+        }
+    }
+    _mm256_storeu_ps(f32x8.as_mut_ptr(), max_32x8);
+    for &x in f32x8.iter() {
+        if x > max {
+            max = x;
+        }
+    }
+
+    for _ in 0..rest {
+        if *ptr < min {
+            min = *ptr;
+        }
+        if *ptr > max {
+            max = *ptr;
+        }
+        ptr = ptr.add(1);
+    }
+
+    (min, max)
+}