Add a benchmark/profiler to count cycles in demo-prover (#577)

* extract data_gen into its own function

* temp push

* add prover cycles

* include multiple transactions

* reduce txns

* temp switch

* cycle tracker macro

* another temp commit

* merge main and nightly

* use run_without_prover for bench/profile

* working metrics

* add zk-cycle-utils. missed it

* temporary commit

* pass features to risc0 vm

* feature gate all the bench things

* cleanup and testing

* check in readme

* checking some cargo locks in prep for merge

* some fixes

* revert rollup config toml

* working merge

* cleanup and documentation

* simplify the macro for generating the wrapped function

* format fixes

* README changes

* simplify code as per lint

* lint fix

* add profiler to analyze trace

* fix cargo lock

* really fix Cargo.lock

* trace options

* cargo fmt

* fix log error

* fix prover bench

* add docs and address comments

* prover bench readme

* add a skip proving option to demo-prover

* add tests for macros

* demo prover format

* remove commented code

* formatting

---------

Co-authored-by: dubbelosix <dub@006.com>

Cargo.toml

@@ -11,7 +11,10 @@ members = [
     "examples/demo-nft-module",
     "full-node/db/sov-db",
     "full-node/sov-sequencer",
+    "utils/zk-cycle-macros",
+    "utils/zk-cycle-utils",
     "full-node/sov-ethereum",
+
     "full-node/sov-stf-runner",
 
     "module-system/sov-cli",

adapters/celestia/Cargo.toml

@@ -27,19 +27,25 @@ serde_json = { workspace = true }
 tokio = { workspace = true, optional = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
+zk-cycle-macros = {path = "../../utils/zk-cycle-macros", optional=true}
+risc0-zkvm = { version = "0.16", default-features = false, features = ["std"], optional=true}
+risc0-zkvm-platform = { version = "0.16", optional=true}
 
 sov-rollup-interface = { path = "../../rollup-interface" }
 nmt-rs = { git = "https://github.com/Sovereign-Labs/nmt-rs.git", rev = "dd37588444fca72825d11fe4a46838f66525c49f", features = ["serde", "borsh"] }
 
+
 [dev-dependencies]
 postcard = { version = "1", features = ["use-std"] }
 proptest = { version = "1.2" }
 wiremock = "0.5"
 
+
 [build-dependencies]
 prost-build = { version = "0.11" }
 
 [features]
 default = ["native"]
 native = ["dep:tokio", "dep:jsonrpsee"]
+bench = ["zk-cycle-macros/bench","risc0-zkvm","risc0-zkvm-platform"]
 verifier = []

adapters/celestia/src/verifier/mod.rs

@@ -13,6 +13,8 @@ pub mod address;
 pub mod proofs;
 
 use proofs::*;
+#[cfg(all(target_os = "zkvm", feature = "bench"))]
+use zk_cycle_macros::cycle_tracker;
 
 use self::address::CelestiaAddress;
 use crate::share_commit::recreate_commitment;
@@ -151,6 +153,7 @@ impl da::DaVerifier for CelestiaVerifier {
         }
     }
 
+    #[cfg_attr(all(target_os = "zkvm", feature = "bench"), cycle_tracker)]
     fn verify_relevant_tx_list<H: Digest>(
         &self,
         block_header: &<Self::Spec as DaSpec>::BlockHeader,
@@ -248,6 +251,7 @@ impl da::DaVerifier for CelestiaVerifier {
                 let mut blob_iter = blob_ref.data();
                 let mut blob_data = vec![0; blob_iter.remaining()];
                 blob_iter.copy_to_slice(blob_data.as_mut_slice());
+
                 let tx_data = tx.data().accumulator();
 
                 match tx_data {

adapters/risc0/Cargo.toml

@@ -15,13 +15,17 @@ readme = "README.md"
 anyhow = { workspace = true }
 bincode = { workspace = true }
 risc0-zkvm = { version = "0.16", default-features = false, features = ['std'] }
+risc0-zkvm-platform = { version = "0.16" }
 risc0-zkp = { version = "0.16", optional = true }
 risc0-circuit-rv32im = { version = "0.16", optional = true }
 serde = { workspace = true }
-
+bytemuck = "1.13.1"
+once_cell = { version = "1.7.2", optional = true}
+parking_lot = { version = "0.11.1", optional = true}
+zk-cycle-utils = { path = "../../utils/zk-cycle-utils"}
 sov-rollup-interface = { path = "../../rollup-interface" }
 
-
 [features]
 default = ["native"]
 native = ["risc0-zkvm/prove", "dep:risc0-zkp", "dep:risc0-circuit-rv32im"]
+bench = ["once_cell","parking_lot"]

adapters/risc0/src/host.rs

@@ -7,7 +7,11 @@ use risc0_zkvm::{
 };
 use sov_rollup_interface::zk::{Zkvm, ZkvmHost};
 use sov_rollup_interface::AddressTrait;
+#[cfg(feature = "bench")]
+use zk_cycle_utils::{cycle_count_callback, get_syscall_name, get_syscall_name_cycles};
 
+#[cfg(feature = "bench")]
+use crate::metrics::metrics_callback;
 use crate::Risc0MethodId;
 
 pub struct Risc0Host<'a> {
@@ -16,9 +20,28 @@ pub struct Risc0Host<'a> {
 }
 
 impl<'a> Risc0Host<'a> {
+    #[cfg(not(feature = "bench"))]
     pub fn new(elf: &'a [u8]) -> Self {
+        let default_env = ExecutorEnvBuilder::default();
+
+        Self {
+            env: RefCell::new(default_env),
+            elf,
+        }
+    }
+
+    #[cfg(feature = "bench")]
+    pub fn new(elf: &'a [u8]) -> Self {
+        let mut default_env = ExecutorEnvBuilder::default();
+
+        let metrics_syscall_name = get_syscall_name();
+        default_env.io_callback(metrics_syscall_name, metrics_callback);
+
+        let cycles_syscall_name = get_syscall_name_cycles();
+        default_env.io_callback(cycles_syscall_name, cycle_count_callback);
+
         Self {
-            env: RefCell::new(ExecutorEnvBuilder::default()),
+            env: RefCell::new(default_env),
             elf,
         }
     }

adapters/risc0/src/lib.rs

@@ -6,6 +6,9 @@ pub mod guest;
 #[cfg(feature = "native")]
 pub mod host;
 
+#[cfg(feature = "bench")]
+pub mod metrics;
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Risc0MethodId([u32; 8]);
 

adapters/risc0/src/metrics.rs

@@ -0,0 +1,33 @@
+use std::collections::HashMap;
+
+use once_cell::sync::Lazy;
+use parking_lot::Mutex;
+
+pub static GLOBAL_HASHMAP: Lazy<Mutex<HashMap<String, (u64, u64)>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
+pub fn add_value(metric: String, value: u64) {
+    let mut hashmap = GLOBAL_HASHMAP.lock();
+    hashmap
+        .entry(metric)
+        .and_modify(|(sum, count)| {
+            *sum += value;
+            *count += 1;
+        })
+        .or_insert((value, 1));
+}
+
+pub fn deserialize_custom(serialized: &[u8]) -> (String, u64) {
+    let null_pos = serialized.iter().position(|&b| b == 0).unwrap();
+    let (string_bytes, size_bytes_with_null) = serialized.split_at(null_pos);
+    let size_bytes = &size_bytes_with_null[1..]; // Skip the null terminator
+    let string = String::from_utf8(string_bytes.to_vec()).unwrap();
+    let size = u64::from_ne_bytes(size_bytes.try_into().unwrap()); // Convert bytes back into usize
+    (string, size)
+}
+
+pub fn metrics_callback(input: &[u8]) -> Vec<u8> {
+    let met_tuple = deserialize_custom(input);
+    add_value(met_tuple.0, met_tuple.1);
+    vec![]
+}