feat: use rayon for parallel computing (#10)

* feat: use rayon for parallel computing Signed-off-by: Keming <[email protected]> * bump version Signed-off-by: Keming <[email protected]> --------- Signed-off-by: Keming <[email protected]>
kemingy · Dec 4, 2024 · 5b58f37 · 5b58f37
1 parent d56c1db
commit 5b58f37
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 22 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "gathers"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 authors = ["Keming <[email protected]>"]
 license = "Apache-2.0"
@@ -19,6 +19,7 @@ log = "0.4.22"
 num-traits = "0.2.19"
 rand = "0.8.5"
 rand_distr = "0.4.3"
+rayon = "1.10.0"
 
 [profile.dev.package.faer]
 opt-level = 3

diff --git a/README.md b/README.md
@@ -7,8 +7,8 @@
 Clustering algorithm implementation:
 
 - [x] K-means
-- [ ] Parallel
-- [ ] Hierarchical K-means
+- [x] PyO3 binding
 - [x] RaBitQ assignment
+- [x] Parallel with Rayon
 - [ ] mini batch K-means
-- [x] PyO3 binding
+- [ ] Hierarchical K-means
diff --git a/python/Cargo.lock b/python/Cargo.lock
diff --git a/python/Cargo.toml b/python/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "gatherspy"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

diff --git a/src/kmeans.rs b/src/kmeans.rs
@@ -5,6 +5,7 @@ use std::time::Instant;
 
 use log::debug;
 use rand::Rng;
+use rayon::prelude::*;
 
 use crate::distance::{argmin, neg_dot_product, squared_euclidean, Distance};
 use crate::rabitq::RaBitQ;
@@ -15,6 +16,7 @@ const EPS: f32 = 1.0 / 1024.0;
 const MIN_POINTS_PER_CENTROID: usize = 39;
 const MAX_POINTS_PER_CENTROID: usize = 256;
 const LARGE_CLUSTER_THRESHOLD: usize = 1 << 20;
+const RAYON_BLOCK_SIZE: usize = 1024 * 32;
 
 /// Assign vectors to centroids.
 pub fn assign(vecs: &[f32], centroids: &[f32], dim: usize, distance: Distance, labels: &mut [u32]) {
@@ -36,14 +38,22 @@ pub fn assign(vecs: &[f32], centroids: &[f32], dim: usize, distance: Distance, l
             // let squared_x: Vec<f32> = vecs.chunks(dim).map(l2_norm).collect();
             // let squared_y: Vec<f32> = centroids.chunks(dim).map(l2_norm).collect();
 
-            for (i, vec) in vecs.chunks(dim).enumerate() {
-                for (j, centroid) in centroids.chunks(dim).enumerate() {
-                    distances[j] =
-                        // squared_x[i] + squared_y[j] + 2.0 * neg_dot_product(vec, centroid);
-                    squared_euclidean(vec, centroid);
-                }
-                labels[i] = argmin(&distances) as u32;
-            }
+            labels.copy_from_slice(
+                &vecs
+                    .par_chunks(dim * RAYON_BLOCK_SIZE)
+                    .flat_map(|vec| {
+                        let mut par_labels = vec![0; vec.len() / dim];
+                        let mut par_distances = vec![f32::MAX; centroids.len() / dim];
+                        for (i, v) in vec.chunks(dim).enumerate() {
+                            for (j, centroid) in centroids.chunks(dim).enumerate() {
+                                par_distances[j] = squared_euclidean(v, centroid);
+                            }
+                            par_labels[i] = argmin(&par_distances) as u32;
+                        }
+                        par_labels
+                    })
+                    .collect::<Vec<_>>(),
+            );
         }
     }
 }
@@ -53,9 +63,18 @@ pub fn assign(vecs: &[f32], centroids: &[f32], dim: usize, distance: Distance, l
 /// TODO: support dot product distance
 pub fn rabitq_assign(vecs: &[f32], centroids: &[f32], dim: usize, labels: &mut [u32]) {
     let rabitq = RaBitQ::new(centroids, dim);
-    for (i, vec) in vecs.chunks(dim).enumerate() {
-        labels[i] = rabitq.retrieve_top_one(vec) as u32;
-    }
+
+    labels.copy_from_slice(
+        &vecs
+            .par_chunks(dim * RAYON_BLOCK_SIZE)
+            .flat_map(|vec| {
+                vec.chunks(dim)
+                    .map(|v| rabitq.retrieve_top_one(v) as u32)
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>(),
+    );
+
     let (rough, precise) = rabitq.get_metrics();
     debug!(
         "RaBitQ: rough {}, precise {}, ratio: {}",
@@ -197,9 +216,7 @@ impl KMeans {
         if num > MAX_POINTS_PER_CENTROID * self.n_cluster as usize {
             let n_sample = MAX_POINTS_PER_CENTROID * self.n_cluster as usize;
             debug!("subsample to {} points", n_sample);
-            let subsampled = as_continuous_vec(&subsample(n_sample, &vecs, dim));
-            vecs.shrink_to(subsampled.len());
-            vecs.copy_from_slice(&subsampled);
+            vecs = as_continuous_vec(&subsample(n_sample, &vecs, dim));
         }
 
         let mut centroids = as_continuous_vec(&subsample(self.n_cluster as usize, &vecs, dim));