add query count to metrics, update dependencies

Signed-off-by: Keming <kemingy94@gmail.com>
kemingy · Aug 28, 2024 · 4c5e135 · 4c5e135
1 parent a81f937
commit 4c5e135
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 28 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/main.rs b/src/main.rs
@@ -32,7 +32,7 @@ struct Args {
 }
 
 fn main() {
-    let env = Env::default().filter_or("RABITQ_LOG", "info");
+    let env = Env::default().filter_or("RABITQ_LOG", "debug");
     env_logger::init_from_env(env);
 
     let args: Args = argh::from_env();

diff --git a/src/metrics.rs b/src/metrics.rs
@@ -9,6 +9,8 @@ pub struct Metrics {
     pub rough: AtomicU64,
     /// precise count
     pub precise: AtomicU64,
+    /// query count
+    pub query: AtomicU64,
 }
 
 impl Metrics {
@@ -17,6 +19,7 @@ impl Metrics {
         Self {
             rough: AtomicU64::new(0),
             precise: AtomicU64::new(0),
+            query: AtomicU64::new(0),
         }
     }
 
@@ -25,7 +28,8 @@ impl Metrics {
         let rough = self.rough.load(Ordering::Relaxed);
         let precise = self.precise.load(Ordering::Relaxed);
         format!(
-            "rough: {}, precise: {}, ratio: {:.2}",
+            "query: {}, rough: {}, precise: {}, ratio: {:.2}",
+            self.query.load(Ordering::Relaxed),
             rough,
             precise,
             rough as f64 / precise as f64,
@@ -41,6 +45,11 @@ impl Metrics {
     pub fn add_precise_count(&self, count: u64) {
         self.precise.fetch_add(count, Ordering::Relaxed);
     }
+
+    /// add query count
+    pub fn add_query_count(&self, count: u64) {
+        self.query.fetch_add(count, Ordering::Relaxed);
+    }
 }
 
 /// Metrics instance.

diff --git a/src/rabitq.rs b/src/rabitq.rs
@@ -431,6 +431,7 @@ impl RaBitQ {
         }
 
         METRICS.add_precise_count(res.len() as u64);
+        METRICS.add_query_count(1);
         let length = topk.min(res.len());
         res.select_nth_unstable_by(length - 1, |a, b| a.0.total_cmp(&b.0));
         res.truncate(length);

diff --git a/src/simd.rs b/src/simd.rs
@@ -23,8 +23,7 @@ pub unsafe fn l2_squared_distance(lhs: &DVectorView<f32>, rhs: &DVectorView<f32>
     let mut rhs_ptr = rhs.as_ptr();
     let block_16_num = lhs.len() >> 4;
     let rest_num = lhs.len() & 0b1111;
-    let mut temp_block = [0.0f32; 8];
-    let temp_block_ptr = temp_block.as_mut_ptr();
+    let mut f32x8 = [0.0f32; 8];
     let (mut diff, mut vx, mut vy): (__m256, __m256, __m256);
     let mut sum = _mm256_setzero_ps();
 
@@ -52,16 +51,9 @@ pub unsafe fn l2_squared_distance(lhs: &DVectorView<f32>, rhs: &DVectorView<f32>
         diff = _mm256_sub_ps(vx, vy);
         sum = _mm256_fmadd_ps(diff, diff, sum);
     }
-    _mm256_store_ps(temp_block_ptr, sum);
-
-    let mut res = temp_block[0]
-        + temp_block[1]
-        + temp_block[2]
-        + temp_block[3]
-        + temp_block[4]
-        + temp_block[5]
-        + temp_block[6]
-        + temp_block[7];
+    _mm256_store_ps(f32x8.as_mut_ptr(), sum);
+    let mut res =
+        f32x8[0] + f32x8[1] + f32x8[2] + f32x8[3] + f32x8[4] + f32x8[5] + f32x8[6] + f32x8[7];
 
     for _ in 0..rest_num {
         let residual = *lhs_ptr - *rhs_ptr;