chore: modify based on the comments

zkMIPS · Dec 28, 2024 · ab39a6c · ab39a6c
1 parent eded172
commit ab39a6c
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 34 deletions.
diff --git a/prover/Cargo.toml b/prover/Cargo.toml
@@ -24,7 +24,6 @@ log = { version = "0.4.14", default-features = false }
 anyhow = "1.0.75"
 num = "0.4.0"
 num-bigint = "0.4.3"
-rayon = "1.10"
 serde = { version = "1.0.144", features = ["derive"] }
 serde_json = "1.0"
 tiny-keccak = "2.0.2"

diff --git a/prover/src/witness/traces.rs b/prover/src/witness/traces.rs
@@ -1,7 +1,9 @@
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::hash::hash_types::RichField;
+use plonky2::timed;
 use plonky2::util::timing::TimingTree;
+use plonky2_maybe_rayon::rayon;
 use std::cmp::max;
 
 use crate::all_stark::{AllStark, MIN_TRACE_LEN, NUM_TABLES};
@@ -134,7 +136,7 @@ impl<T: Copy> Traces<T> {
         self,
         all_stark: &AllStark<T, D>,
         config: &StarkConfig,
-        _timing: &mut TimingTree,
+        timing: &mut TimingTree,
     ) -> [Vec<PolynomialValues<T>>; NUM_TABLES]
     where
         T: RichField + Extendable<D>,
@@ -157,38 +159,42 @@ impl<T: Copy> Traces<T> {
         let mut poseidon_sponge_trace= vec![];
         let mut logic_trace = vec![];
 
-        rayon::join(
-            ||
-                rayon::join (
-                    ||  memory_trace = all_stark.memory_stark.generate_trace(
-                            memory_ops,
-                            &mut TimingTree::new("memory", log::Level::Info),
-                        ),
-                    ||  arithmetic_trace = all_stark.arithmetic_stark.generate_trace(arithmetic_ops),
-                ),
-            ||  {
-                rayon::join (
-                    ||  cpu_trace = trace_rows_to_poly_values(
-                            cpu.into_iter().map(|x| x.into()).collect(),
-                        ),
-                    ||  poseidon_trace =  all_stark.poseidon_stark.generate_trace(
-                            poseidon_inputs, min_rows,
-                            &mut TimingTree::new("poseidon", log::Level::Info),
-                        ),
-                );
-                rayon::join (
-                    ||  poseidon_sponge_trace =  all_stark.poseidon_sponge_stark.generate_trace(
-                            poseidon_sponge_ops,
-                            min_rows,
-                            &mut TimingTree::new("poseidon_sponge", log::Level::Info),
-                        ),
-                    ||  logic_trace =  all_stark.logic_stark.generate_trace(
-                            logic_ops,
-                            min_rows,
-                            &mut TimingTree::new("logic", log::Level::Info),
-                        ),
-                );
-            },
+        timed!(
+            timing,
+            "convert trace to table parallelly",
+            rayon::join(
+                ||
+                    rayon::join (
+                        ||  memory_trace = all_stark.memory_stark.generate_trace(
+                                memory_ops,
+                                &mut TimingTree::new("memory", log::Level::Info),
+                            ),
+                        ||  arithmetic_trace = all_stark.arithmetic_stark.generate_trace(arithmetic_ops),
+                    ),
+                ||  {
+                    rayon::join (
+                        ||  cpu_trace = trace_rows_to_poly_values(
+                                cpu.into_iter().map(|x| x.into()).collect(),
+                            ),
+                        ||  poseidon_trace =  all_stark.poseidon_stark.generate_trace(
+                                poseidon_inputs, min_rows,
+                                &mut TimingTree::new("poseidon", log::Level::Info),
+                            ),
+                    );
+                    rayon::join (
+                        ||  poseidon_sponge_trace =  all_stark.poseidon_sponge_stark.generate_trace(
+                                poseidon_sponge_ops,
+                                min_rows,
+                                &mut TimingTree::new("poseidon_sponge", log::Level::Info),
+                            ),
+                        ||  logic_trace =  all_stark.logic_stark.generate_trace(
+                                logic_ops,
+                                min_rows,
+                                &mut TimingTree::new("logic", log::Level::Info),
+                            ),
+                    );
+                },
+            )
         );
 
         [