[chore] Parallelize Poseidon trace generation (#1045)

* Parallelize Poseidon trace generation * chore: par_extend --------- Co-authored-by: Jonathan Wang <[email protected]>
openvm-org · Dec 15, 2024 · 6c0c0e2 · 6c0c0e2
1 parent a26c3de
commit 6c0c0e2
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml
@@ -38,6 +38,7 @@ derivative.workspace = true
 static_assertions.workspace = true
 async-trait.workspace = true
 getset.workspace = true
+rayon = { workspace = true, optional = true }
 
 [dev-dependencies]
 p3-dft = { workspace = true }
@@ -70,7 +71,7 @@ hex.workspace = true
 
 [features]
 default = ["parallel", "mimalloc"]
-parallel = ["openvm-stark-backend/parallel"]
+parallel = ["openvm-stark-backend/parallel", "dep:rayon"]
 test-utils = ["openvm-ecc-guest/halo2curves", "dep:openvm-stark-sdk"]
 bench-metrics = [
     "dep:metrics",

diff --git a/crates/vm/src/system/poseidon2/trace.rs b/crates/vm/src/system/poseidon2/trace.rs
@@ -6,10 +6,13 @@ use openvm_stark_backend::{
     p3_air::BaseAir,
     p3_field::PrimeField32,
     p3_matrix::dense::RowMajorMatrix,
+    p3_maybe_rayon::prelude::*,
     prover::types::AirProofInput,
     rap::{get_air_name, AnyRap},
     Chip, ChipUsageGetter,
 };
+#[cfg(feature = "parallel")]
+use rayon::iter::ParallelExtend;
 
 use super::{columns::*, Poseidon2Chip};
 
@@ -35,12 +38,21 @@ where
 
         let aux_cols_factory = memory_controller.borrow().aux_cols_factory();
         let mut flat_rows: Vec<_> = records
-            .into_iter()
+            .into_par_iter()
             .flat_map(|record| Self::record_to_cols(&aux_cols_factory, record).flatten())
             .collect();
-        for _ in 0..diff {
-            flat_rows.extend(Poseidon2VmCols::<Val<SC>>::blank_row(&air).flatten());
-        }
+        #[cfg(feature = "parallel")]
+        flat_rows.par_extend(
+            vec![Poseidon2VmCols::<Val<SC>>::blank_row(&air).flatten(); diff]
+                .into_par_iter()
+                .flatten(),
+        );
+        #[cfg(not(feature = "parallel"))]
+        flat_rows.extend(
+            vec![Poseidon2VmCols::<Val<SC>>::blank_row(&air).flatten(); diff]
+                .into_iter()
+                .flatten(),
+        );
 
         AirProofInput::simple_no_pis(
             Arc::new(air.clone()),