Add a benchmark/profiler to count cycles in demo-prover

adapters/risc0/src/host.rs

@@ -20,19 +20,25 @@ pub struct Risc0Host<'a> {
 }
 
 impl<'a> Risc0Host<'a> {
+    #[cfg(not(feature = "bench"))]
     pub fn new(elf: &'a [u8]) -> Self {
-        #[cfg(not(feature = "bench"))]
         let default_env = ExecutorEnvBuilder::default();
-        #[cfg(feature = "bench")]
-        let mut default_env = ExecutorEnvBuilder::default();
-        #[cfg(feature = "bench")]
-        {
-            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(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(default_env),

examples/demo-prover/host/benches/prover_bench.rs

@@ -9,7 +9,7 @@ use anyhow::Context;
 use const_rollup_config::{ROLLUP_NAMESPACE_RAW, SEQUENCER_DA_ADDRESS};
 use demo_stf::app::{App, DefaultPrivateKey};
 use demo_stf::genesis_config::create_demo_genesis_config;
-use jupiter::da_service::{CelestiaService, DaServiceConfig};
+use jupiter::da_service::CelestiaService;
 use jupiter::types::{FilteredCelestiaBlock, NamespaceId};
 use jupiter::verifier::address::CelestiaAddress;
 use jupiter::verifier::{ChainValidityCondition, RollupParams};
@@ -18,13 +18,12 @@ use log4rs::config::{Appender, Config, Root};
 use methods::ROLLUP_ELF;
 use regex::Regex;
 use risc0_adapter::host::Risc0Host;
-use serde::Deserialize;
 use sov_modules_api::PrivateKey;
 use sov_rollup_interface::services::da::DaService;
 use sov_rollup_interface::stf::StateTransitionFunction;
 use sov_rollup_interface::zk::ZkvmHost;
 use sov_state::storage::Storage;
-use sov_stf_runner::{from_toml_path, Config as RunnerConfig};
+use sov_stf_runner::{from_toml_path, RollupConfig};
 use tempfile::TempDir;
 
 #[derive(Debug)]
@@ -33,12 +32,6 @@ struct RegexAppender {
     file: Arc<Mutex<File>>,
 }
 
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
-struct HexKey {
-    hex_priv_key: String,
-    address: String,
-}
-
 impl RegexAppender {
     fn new(pattern: &str, file_path: &str) -> Self {
         if Path::new(file_path).exists() {
@@ -96,12 +89,12 @@ use std::collections::HashMap;
 #[cfg(feature = "bench")]
 use risc0_adapter::metrics::GLOBAL_HASHMAP;
 
-#[derive(Debug, Clone, PartialEq, Deserialize)]
-pub struct RollupConfig {
-    pub start_height: u64,
-    pub da: DaServiceConfig,
-    pub runner: RunnerConfig,
-}
+// #[derive(Debug, Clone, PartialEq, Deserialize)]
+// pub struct RollupConfig {
+//     pub start_height: u64,
+//     pub da: DaServiceConfig,
+//     pub runner: RunnerConfig,
+// }
 
 // The rollup stores its data in the namespace b"sov-test" on Celestia
 const ROLLUP_NAMESPACE: NamespaceId = NamespaceId(ROLLUP_NAMESPACE_RAW);

examples/demo-prover/methods/build.rs

@@ -1,17 +1,24 @@
 use std::collections::HashMap;
 
 fn main() {
-    #[cfg(not(feature = "bench"))]
-    let guest_pkg_to_options = HashMap::new();
-    #[cfg(feature = "bench")]
+    let guest_pkg_to_options = get_guest_options();
+    risc0_build::embed_methods_with_options(guest_pkg_to_options);
+}
+
+#[cfg(not(feature = "bench"))]
+fn get_guest_options() -> HashMap<&'static str, risc0_build::GuestOptions> {
+    HashMap::new()
+}
+
+#[cfg(feature = "bench")]
+fn get_guest_options() -> HashMap<&'static str, risc0_build::GuestOptions> {
     let mut guest_pkg_to_options = HashMap::new();
-    #[cfg(feature = "bench")]
     guest_pkg_to_options.insert(
         "sov-demo-prover-guest",
         risc0_build::GuestOptions {
             features: vec!["bench".to_string()],
             std: true,
         },
     );
-    risc0_build::embed_methods_with_options(guest_pkg_to_options);
+    guest_pkg_to_options
 }

utils/zk-cycle-macros/src/lib.rs

@@ -1,9 +1,23 @@
+#![deny(missing_docs)]
+#![doc = include_str!("../README.md")]
 extern crate proc_macro;
 
 use proc_macro::TokenStream;
 use quote::quote;
 use syn::{parse_macro_input, ItemFn};
 
+/// This macro is used to annotate functions that we want to track the number of riscV cycles being
+/// generated inside the VM. The purpose of the this macro is to measure how many cycles a rust
+/// function takes because prover time is directly proportional to the number of riscv cycles
+/// generated. It does this by making use of a risc0 provided function
+/// ```
+/// risc0_zkvm::guest::env::get_cycle_count
+/// ```
+/// The macro essentially generates new function with the same name by wrapping the body with a get_cycle_count
+/// at the beginning and end of the function, subtracting it and then emitting it out using the
+/// a custom syscall that is generated when the prover is run with the `bench` feature.
+/// `send_recv_slice` is used to communicate and pass a slice to the syscall that we defined.
+/// The handler for the syscall can be seen in adapters/risc0/src/host.rs and adapters/risc0/src/metrics.rs
 #[proc_macro_attribute]
 pub fn cycle_tracker(_attr: TokenStream, item: TokenStream) -> TokenStream {
     let input = parse_macro_input!(item as ItemFn);