Add option to output zk_os_runner flamegraph

evm_tester/Cargo.toml

@@ -63,6 +63,8 @@ zk_os_evm_interpreter = { package = "evm_interpreter",  git="ssh://[email protected]
 zk_os_basic_system = { package = "basic_system",  git="ssh://[email protected]/matter-labs/zk_ee.git", branch = "main"  }
 zk_os_basic_bootloader = { package = "basic_bootloader",  git="ssh://[email protected]/matter-labs/zk_ee.git", branch = "main" , features = ["code_in_kernel_space"] }
 zk_os_system_hooks = { package = "system_hooks",  git="ssh://[email protected]/matter-labs/zk_ee.git", branch = "main"  }
+zk_os_oracle_provider = {package = "oracle_provider", git="ssh://[email protected]/matter-labs/zk_ee.git", branch = "main"  }
+zk_os_runner = {package = "zk_os_runner", git="ssh://[email protected]/matter-labs/zk_ee.git", branch = "main"  }
 
 [dependencies.web3]
 version = "=0.19.0"

evm_tester/src/evm_tester/arguments.rs

@@ -40,7 +40,10 @@ pub struct Arguments {
     #[structopt(long = "environment")]
     pub environment: Option<evm_tester::Environment>,
 
-    /// Choose between `build` to compile tests only without running, and `run` to compile and run.
+    /// Choose between `build` to compile tests only without running, `run` to compile and run
+    /// or `bench` to also produce flamegraphs.
+    /// Note that you might want to set the ZKOS_DIR env var to point to the directory
+    /// containing the app.elf and app.bin from ZK OS to run benchmarks.
     #[structopt(long = "workflow", default_value = "run")]
     pub workflow: evm_tester::Workflow,
 }

evm_tester/src/lib.rs

@@ -101,7 +101,11 @@ impl EvmTester {
         let _: Vec<()> = tests
             .into_par_iter()
             .map(|test| {
-                test.run_zk_os(self.summary.clone(), vm.clone());
+                test.run_zk_os(
+                    self.summary.clone(),
+                    vm.clone(),
+                    matches!(self.workflow, Workflow::Bench),
+                );
             })
             .collect();
 

evm_tester/src/test/case/mod.rs

@@ -479,11 +479,12 @@ impl Case {
         vm: ZkOS,
         test_name: String,
         test_group: Option<String>,
+        bench: bool,
     ) {
         let calldata = self.transaction.data.0.clone();
         let name = self.label.clone();
         let result = std::panic::catch_unwind(|| {
-            self.run_zk_os_inner(summary.clone(), vm, test_name.clone(), test_group)
+            self.run_zk_os_inner(summary.clone(), vm, test_name.clone(), test_group, bench)
         });
         if let Err(e) = result {
             Summary::panicked(
@@ -501,6 +502,7 @@ impl Case {
         mut vm: ZkOS,
         test_name: String,
         test_group: Option<String>,
+        bench: bool,
     ) {
         let name = self.label;
 
@@ -552,6 +554,7 @@ impl Case {
         if let Some(random) = self.env.current_random {
             system_context.block_difficulty = utils::u256_to_h256(&random);
         }
+        let test_id = format!("{}-{}", test_name, name);
         let run_result = vm.execute_transaction(
             self.transaction.secret_key,
             self.transaction.to.0,
@@ -560,6 +563,8 @@ impl Case {
             self.transaction.gas_limit,
             self.transaction.nonce.try_into().expect("Nonce overflow"),
             system_context,
+            bench,
+            test_id,
         );
 
         let mut check_successful = true;

evm_tester/src/test/mod.rs

@@ -149,7 +149,7 @@ impl Test {
     ///
     /// Runs the test on ZK OS.
     ///
-    pub fn run_zk_os(self, summary: Arc<Mutex<Summary>>, vm: Arc<ZkOS>) {
+    pub fn run_zk_os(self, summary: Arc<Mutex<Summary>>, vm: Arc<ZkOS>, bench: bool) {
         for case in self.cases {
             if let Some(filter_calldata) = self.skipped_calldatas.as_ref() {
                 if filter_calldata.contains(&case.transaction.data) {
@@ -166,7 +166,13 @@ impl Test {
             }
 
             let vm = ZkOS::clone(vm.clone());
-            case.run_zk_os(summary.clone(), vm, self.name.clone(), self.group.clone());
+            case.run_zk_os(
+                summary.clone(),
+                vm,
+                self.name.clone(),
+                self.group.clone(),
+                bench,
+            );
         }
     }
 }

evm_tester/src/vm/zk_ee/mod.rs

@@ -93,6 +93,8 @@ impl ZkOS {
         gas_limit: U256,
         nonce: u32,
         system_context: ZkOsEVMContext,
+        bench: bool,
+        test_id: String,
     ) -> anyhow::Result<ZkOsExecutionResult, String> {
         let fee = Fee {
             gas_limit,
@@ -146,6 +148,36 @@ impl ZkOS {
         let tree = self.tree.clone();
         let preimage_source = self.preimage_source.clone();
 
+        // Output flamegraphs if on benchmarking mode
+        if bench {
+            use zk_os_forward_system::run::io_implementer_init_data;
+            use zk_os_forward_system::run::ForwardRunningOracle;
+            use zk_os_oracle_provider::BasicZkEEOracleWrapper;
+            use zk_os_oracle_provider::ReadWitnessSource;
+            use zk_os_oracle_provider::ZkEENonDeterminismSource;
+
+            let oracle: ForwardRunningOracle<InMemoryTree, InMemoryPreimageSource, TxListSource> =
+                ForwardRunningOracle {
+                    io_implementer_init_data: Some(io_implementer_init_data(Some(
+                        storage_commitment,
+                    ))),
+                    block_metadata: context,
+                    tree: tree.clone(),
+                    preimage_source: preimage_source.clone(),
+                    tx_source: tx_source.clone(),
+                    next_tx: None,
+                };
+            let oracle_wrapper = BasicZkEEOracleWrapper::new(oracle.clone());
+            let mut non_determinism_source = ZkEENonDeterminismSource::default();
+            non_determinism_source.add_external_processor(oracle_wrapper);
+            let copy_source = ReadWitnessSource::new(non_determinism_source);
+            let path = std::env::current_dir()
+                .unwrap()
+                .join(format!("{}.svg", test_id));
+            let _output =
+                zk_os_runner::run_default_with_flamegraph_path(1 << 25, copy_source, Some(path));
+        }
+
         let result = run_batch(
             context,
             storage_commitment,

evm_tester/src/workflow.rs

@@ -13,6 +13,8 @@ pub enum Workflow {
     BuildOnly,
     /// Build and execute tests.
     BuildAndRun,
+    /// Build, execute and benchmark tests.
+    Bench,
 }
 
 impl FromStr for Workflow {
@@ -22,6 +24,7 @@ impl FromStr for Workflow {
         match day {
             "build" => Ok(Workflow::BuildOnly),
             "run" => Ok(Workflow::BuildAndRun),
+            "bench" => Ok(Workflow::Bench),
             string => anyhow::bail!(
                 "Unknown workflow `{}`. Supported workflows: {}",
                 string,
@@ -40,6 +43,7 @@ impl std::fmt::Display for Workflow {
         match self {
             Workflow::BuildOnly => write!(f, "build"),
             Workflow::BuildAndRun => write!(f, "run"),
+            Workflow::Bench => write!(f, "bench"),
         }
     }
 }