Merge pull request #177 from promised-ai/feature/parallel-slice

Feature/parallel slice

CHANGELOG.md

@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `DataParseError::DataFrameMissingColumn` variant for when a column is in the codebook but not in the initial dataframe.
 - Python's `Engine.update` uses `tqdm.auto` for progress bar reporting.
 
+### Changed
+- Added parallelism to `Slice` row reassignment kernel. Run time is ~6x faster.
+
 ### Fixed
 - Initializing an engine with a codebook that has a different number of rows than the data will result in an error instead of printing a bunch on nonsense.
 - Pylace default transition sets didn't hit all required transitions

cli/Cargo.toml

@@ -26,6 +26,3 @@ serde_yaml = "0.9.4"
 approx = "0.5.1"
 tempfile = "3.4"
 indoc = "2.0.3"
-
-[profile.release]
-lto = "fat"

lace/lace_cc/src/view.rs

@@ -966,8 +966,9 @@ impl View {
         row_alg: RowAssignAlg,
         rng: &mut impl Rng,
     ) {
-        // TODO: parallelize over rows_mut somehow?
-        logps.rows_mut().enumerate().for_each(|(k, logp)| {
+        use rayon::prelude::*;
+
+        logps.par_rows_mut().enumerate().for_each(|(k, logp)| {
             self.ftrs.values().for_each(|ftr| {
                 ftr.accum_score(logp, k);
             })

lace/lace_utils/Cargo.toml

@@ -10,6 +10,7 @@ description = "Miscellaneous utilities for Lace and shared libraries"
 
 [dependencies]
 rand = {version="0.8", features=["serde1"]}
+rayon = "1.5"
 
 [dev-dependencies]
 approx = "0.5.1"

lace/lace_utils/src/matrix.rs

@@ -1,3 +1,4 @@
+use rayon::prelude::{ParallelSlice, ParallelSliceMut};
 use std::ops::Index;
 
 /// A lightweight Matrix abstraction that does almost nothing.
@@ -71,8 +72,14 @@ impl<T: Send + Sync> Matrix<T> {
     /// assert_eq!(mat.raw_values(), &vec![1, 2, 3, 4, 5, 6])
     /// ```
     #[inline]
-    pub fn rows_mut(&mut self) -> std::slice::ChunksMut<T> {
-        self.values.chunks_mut(self.n_cols)
+    pub fn rows_mut(&mut self) -> std::slice::ChunksExactMut<T> {
+        self.values.chunks_exact_mut(self.n_cols)
+    }
+
+    /// Parallel version of `rows_mut`
+    #[inline]
+    pub fn par_rows_mut(&mut self) -> rayon::slice::ChunksExactMut<T> {
+        self.values.par_chunks_exact_mut(self.n_cols)
     }
 
     /// Create an iterator through rows
@@ -96,8 +103,14 @@ impl<T: Send + Sync> Matrix<T> {
     /// assert_eq!(rowsum, vec![3_u8, 12_u8])
     /// ```
     #[inline]
-    pub fn rows(&self) -> std::slice::Chunks<T> {
-        self.values.chunks(self.n_cols)
+    pub fn rows(&self) -> std::slice::ChunksExact<T> {
+        self.values.chunks_exact(self.n_cols)
+    }
+
+    /// Parallel version of `rows`
+    #[inline]
+    pub fn par_rows(&self) -> rayon::slice::ChunksExact<T> {
+        self.values.par_chunks_exact(self.n_cols)
     }
 
     /// Does an implicit transpose by inverting coordinates.