use qr to get random orthogonal, add simd query binarize func

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

Author: kemingy

src/consts.rs

@@ -0,0 +1,6 @@
+//! Constants used in the program.
+
+pub(crate) const DEFAULT_X_DOT_PRODUCT: f32 = 0.8;
+pub(crate) const EPSILON: f32 = 1.9;
+pub(crate) const THETA_LOG_DIM: u32 = 4;
+pub(crate) const WINDOWS_SIZE: usize = 12;

src/lib.rs

@@ -1,9 +1,10 @@
 //! RaBitQ implementation in Rust.
 
 #![forbid(missing_docs)]
-pub mod distance;
+mod consts;
 pub mod metrics;
 pub mod rabitq;
+pub mod simd;
 pub mod utils;
 
 pub use rabitq::RaBitQ;

src/metrics.rs

@@ -22,10 +22,13 @@ impl Metrics {
 
     /// get the instance
     pub fn to_str(&self) -> String {
+        let rough = self.rough.load(Ordering::Relaxed);
+        let precise = self.precise.load(Ordering::Relaxed);
         format!(
-            "rough: {}, precise: {}",
-            self.rough.load(Ordering::Relaxed),
-            self.precise.load(Ordering::Relaxed)
+            "rough: {}, precise: {}, ratio: {:.2}",
+            rough,
+            precise,
+            rough as f64 / precise as f64,
         )
     }
 

src/rabitq.rs

@@ -7,13 +7,9 @@ use nalgebra::{DMatrix, DVector, DVectorView};
 use num_traits::ToPrimitive;
 use serde::{Deserialize, Serialize};
 
+use crate::consts::{DEFAULT_X_DOT_PRODUCT, EPSILON, THETA_LOG_DIM, WINDOWS_SIZE};
 use crate::metrics::METRICS;
-use crate::utils::{gen_random_bias, gen_random_orthogonal, matrix_from_fvecs};
-
-const DEFAULT_X_DOT_PRODUCT: f32 = 0.8;
-const EPSILON: f32 = 1.9;
-const THETA_LOG_DIM: u32 = 4;
-const WINDOWS_SIZE: usize = 16;
+use crate::utils::{gen_random_bias, gen_random_qr_orthogonal, matrix_from_fvecs};
 
 /// Convert the vector to binary format and store in a u64 vector.
 fn vector_binarize_u64(vec: &DVector<f32>) -> Vec<u64> {
@@ -32,8 +28,25 @@ fn vector_binarize_one(vec: &DVector<f32>) -> DVector<f32> {
     DVector::from_fn(vec.len(), |i, _| if vec[i] > 0.0 { 1.0 } else { -1.0 })
 }
 
+/// Interface of `vector_binarize_query`
+fn vector_binarize_query(vec: &DVector<u8>) -> Vec<u64> {
+    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+    {
+        if is_x86_feature_detected!("avx2") {
+            unsafe { crate::simd::vector_binarize_query_avx2(&vec.as_view()) }
+        } else {
+            vector_binarize_query_raw(vec)
+        }
+    }
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
+    {
+        vector_binarize_query_raw(vec)
+    }
+}
+
 /// Convert the vector to binary format (one value to multiple bits) and store in a u64 vector.
-fn query_vector_binarize(vec: &DVector<u8>) -> Vec<u64> {
+#[inline]
+fn vector_binarize_query_raw(vec: &DVector<u8>) -> Vec<u64> {
     let length = vec.len();
     let mut binary = vec![0u64; length * THETA_LOG_DIM as usize / 64];
     for j in 0..THETA_LOG_DIM as usize {
@@ -78,7 +91,7 @@ fn kmeans_nearest_cluster(centroids: &DMatrix<f32>, vec: &DVectorView<f32>) -> u
             #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
             {
                 if is_x86_feature_detected!("avx2") {
-                    unsafe { crate::distance::l2_squared_distance_avx2(&centroid, vec) }
+                    unsafe { crate::simd::l2_squared_distance_avx2(&centroid, vec) }
                 } else {
                     vec.sub_to(&centroid, &mut residual);
                     residual.norm_squared()
@@ -123,7 +136,7 @@ impl RaBitQ {
         let centroids = matrix_from_fvecs(centroid_path);
         let k = centroids.shape().0;
         debug!("n: {}, dim: {}, k: {}", n, dim, k);
-        let orthogonal = gen_random_orthogonal(dim);
+        let orthogonal = gen_random_qr_orthogonal(dim);
         let rand_bias = gen_random_bias(dim);
 
         // projection
@@ -226,10 +239,7 @@ impl RaBitQ {
                 {
                     if is_x86_feature_detected!("avx2") {
                         unsafe {
-                            crate::distance::l2_squared_distance_avx2(
-                                &centroid,
-                                &y_projected.as_view(),
-                            )
+                            crate::simd::l2_squared_distance_avx2(&centroid, &y_projected.as_view())
                         }
                     } else {
                         y_projected.sub_to(&centroid, &mut residual);
@@ -257,7 +267,7 @@ impl RaBitQ {
             let y_scaled = residual.add_scalar(-lower_bound) * one_over_delta + &self.rand_bias;
             let y_quantized = y_scaled.map(|v| v.to_u8().expect("convert to u8 error"));
             let scalar_sum = y_quantized.iter().fold(0u32, |acc, &v| acc + v as u32);
-            let y_binary_vec = query_vector_binarize(&y_quantized);
+            let y_binary_vec = vector_binarize_query(&y_quantized);
             let dist_sqrt = dist.sqrt();
             for j in self.offsets[i]..self.offsets[i + 1] {
                 let ju = j as usize;
@@ -300,7 +310,7 @@ impl RaBitQ {
                     {
                         if is_x86_feature_detected!("avx2") {
                             unsafe {
-                                crate::distance::l2_squared_distance_avx2(
+                                crate::simd::l2_squared_distance_avx2(
                                     &self.base.column(u as usize),
                                     &query.as_view(),
                                 )

src/distance.rs renamed to src/simd.rs

@@ -1,7 +1,9 @@
-//! Compute the distance between two vectors.
+//! Accelerate with SIMD.
 
 use nalgebra::DVectorView;
 
+use crate::consts::THETA_LOG_DIM;
+
 /// Compute the squared Euclidean distance between two vectors.
 /// Code refer to https://github.com/nmslib/hnswlib/blob/master/hnswlib/space_l2.h
 ///
@@ -69,3 +71,34 @@ pub unsafe fn l2_squared_distance_avx2(lhs: &DVectorView<f32>, rhs: &DVectorView
     }
     res
 }
+
+/// Convert an [u8] to 4x binary vector stored as u64.
+///
+/// # 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 vector_binarize_query_avx2(vec: &DVectorView<u8>) -> Vec<u64> {
+    use std::arch::x86_64::*;
+
+    let length = vec.len();
+    let mut ptr = vec.as_ptr() as *const __m256i;
+    let mut binary = vec![0u64; length * THETA_LOG_DIM as usize / 64];
+
+    for i in (0..length).step_by(32) {
+        // since it's not guaranteed that the vec is fully-aligned
+        let mut v = _mm256_loadu_si256(ptr);
+        ptr = ptr.add(1);
+        v = _mm256_slli_epi32(v, 4);
+        for j in 0..THETA_LOG_DIM as usize {
+            let mask = (_mm256_movemask_epi8(v) as u32) as u64;
+            // let shift = if (i / 32) % 2 == 0 { 32 } else { 0 };
+            let shift = ((i >> 5) & 1) << 5;
+            binary[(3 - j) * (length >> 6) + (i >> 6)] |= mask << shift;
+            v = _mm256_slli_epi32(v, 1);
+        }
+    }
+
+    binary
+}

src/utils.rs

@@ -18,6 +18,13 @@ pub fn gen_random_orthogonal(dim: usize) -> DMatrix<f32> {
     random.unwrap()
 }
 
+/// Generate a random orthogonal matrix from QR decomposition.
+pub fn gen_random_qr_orthogonal(dim: usize) -> DMatrix<f32> {
+    let mut rng = thread_rng();
+    let random = DMatrix::from_fn(dim, dim, |_, _| rng.gen::<f32>());
+    random.qr().q()
+}
+
 /// Generate an identity matrix as a special orthogonal matrix.
 ///
 /// Use this function to debug the logic.