Parallelize test_only segment simulation (#498)

* Parallelize test_only segment simulation * Address comments
0xPolygonZero · Aug 15, 2024 · 9923297 · 9923297
1 parent 74d499d
commit 9923297
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/zero_bin/ops/src/lib.rs b/zero_bin/ops/src/lib.rs
@@ -4,12 +4,14 @@ use std::time::Instant;
 #[cfg(not(feature = "test_only"))]
 use evm_arithmetization::generation::TrimmedGenerationInputs;
 use evm_arithmetization::proof::PublicValues;
-#[cfg(not(feature = "test_only"))]
-use paladin::operation::FatalStrategy;
+#[cfg(feature = "test_only")]
+use evm_arithmetization::{prover::testing::simulate_execution_all_segments, GenerationInputs};
 use paladin::{
-    operation::{FatalError, Monoid, Operation, Result},
+    operation::{FatalError, FatalStrategy, Monoid, Operation, Result},
     registry, RemoteExecute,
 };
+#[cfg(feature = "test_only")]
+use proof_gen::types::Field;
 use proof_gen::{
     proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof},
     proof_types::{
@@ -25,6 +27,25 @@ use zero_bin_common::{debug_utils::save_inputs_to_disk, prover_state::p_state};
 
 registry!();
 
+#[cfg(feature = "test_only")]
+#[derive(Deserialize, Serialize, RemoteExecute)]
+pub struct BatchTestOnly {
+    pub save_inputs_on_error: bool,
+}
+
+#[cfg(feature = "test_only")]
+impl Operation for BatchTestOnly {
+    type Input = (GenerationInputs, usize);
+    type Output = ();
+
+    fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
+        simulate_execution_all_segments::<Field>(inputs.0, inputs.1)
+            .map_err(|err| FatalError::from_anyhow(err, FatalStrategy::Terminate))?;
+
+        Ok(())
+    }
+}
+
 #[derive(Deserialize, Serialize, RemoteExecute)]
 pub struct SegmentProof {
     pub save_inputs_on_error: bool,
@@ -73,8 +94,8 @@ impl Operation for SegmentProof {
     type Input = AllData;
     type Output = ();
 
-    fn execute(&self, _input: Self::Input) -> Result<Self::Output> {
-        todo!() // currently unused, change or remove
+    fn execute(&self, _all_data: Self::Input) -> Result<Self::Output> {
+        Ok(())
     }
 }
 

diff --git a/zero_bin/prover/src/lib.rs b/zero_bin/prover/src/lib.rs
@@ -135,12 +135,20 @@ impl BlockProverInput {
     #[cfg(feature = "test_only")]
     pub async fn prove(
         self,
-        _runtime: &Runtime,
+        runtime: &Runtime,
         previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
         prover_config: ProverConfig,
     ) -> Result<GeneratedBlockProof> {
-        use evm_arithmetization::prover::testing::simulate_execution_all_segments;
-        use plonky2::field::goldilocks_field::GoldilocksField;
+        use std::iter::repeat;
+
+        use futures::StreamExt;
+        use paladin::directive::{Directive, IndexedStream};
+
+        let ProverConfig {
+            max_cpu_len_log,
+            batch_size,
+            save_inputs_on_error,
+        } = prover_config;
 
         let block_number = self.get_block_number();
         info!("Testing witness generation for block {block_number}.");
@@ -149,13 +157,21 @@ impl BlockProverInput {
         let txs = self.block_trace.into_txn_proof_gen_ir(
             &ProcessingMeta::new(resolve_code_hash_fn),
             other_data.clone(),
-            prover_config.batch_size,
+            batch_size,
         )?;
 
-        type F = GoldilocksField;
-        for txn in txs.into_iter() {
-            simulate_execution_all_segments::<F>(txn, prover_config.max_cpu_len_log)?;
-        }
+        let batch_ops = ops::BatchTestOnly {
+            save_inputs_on_error,
+        };
+
+        let simulation = Directive::map(
+            IndexedStream::from(txs.into_iter().zip(repeat(max_cpu_len_log))),
+            &batch_ops,
+        );
+
+        let result = simulation.run(runtime).await?;
+
+        result.collect::<Vec<_>>().await;
 
         // Wait for previous block proof
         let _prev = match previous {