feat: cache proof inputs

bin/reth-benchmark/Cargo.toml

@@ -22,7 +22,9 @@ clap_builder = "4.5.34"
 bincode = { workspace = true, features = ["std"] }
 bitcode.workspace = true
 hex = "0.4.3"
+regex = "1"
 serde_json.workspace = true
+serde.workspace = true
 eyre.workspace = true
 tokio.workspace = true
 

bin/reth-benchmark/src/lib.rs

@@ -18,6 +18,8 @@ use openvm_stateless_executor::{
 };
 
 use openvm_rpc_proxy::{RpcExecutor, DEFAULT_PREIMAGE_CACHE_NIBBLES};
+#[cfg(feature = "evm-verify")]
+use openvm_sdk::keygen::Halo2ProvingKey;
 use openvm_sdk::{
     config::{AppConfig, SdkVmBuilder, SdkVmConfig},
     fs::read_object_from_file,
@@ -28,11 +30,14 @@ use openvm_sdk::{
 };
 use openvm_stark_sdk::engine::StarkFriEngine;
 use openvm_transpiler::{elf::Elf, openvm_platform::memory::MEM_SIZE};
+use regex::Regex;
 pub use reth_ethereum_primitives as reth_primitives;
 use serde_json::json;
 use std::{fs, path::PathBuf};
 use tracing::{info, info_span};
 
+#[cfg(feature = "evm-verify")]
+use serde::{Deserialize, Serialize};
 mod cli;
 
 use cli::ProviderArgs;
@@ -57,6 +62,9 @@ pub enum BenchMode {
     MakeInput,
     /// Generate fixtures file for futher benchmarking.
     GenerateFixtures,
+    /// Generate fixtures (app/root proof + app/agg/halo proving keys) for EVM benchmarking.
+    #[cfg(feature = "evm-verify")]
+    GenerateFixturesEvm,
 }
 
 impl std::fmt::Display for BenchMode {
@@ -71,6 +79,8 @@ impl std::fmt::Display for BenchMode {
             Self::ProveEvm => write!(f, "prove_evm"),
             Self::MakeInput => write!(f, "make_input"),
             Self::GenerateFixtures => write!(f, "generate_fixtures"),
+            #[cfg(feature = "evm-verify")]
+            Self::GenerateFixturesEvm => write!(f, "generate_fixtures_evm"),
         }
     }
 }
@@ -92,6 +102,12 @@ pub struct HostArgs {
     /// created from RPC data if it doesn't already exist.
     #[clap(long)]
     cache_dir: Option<PathBuf>,
+
+    /// Regex of cached part names to recompute (bypassing the cache). Matching parts are still
+    /// written back to the cache after they're recomputed.
+    #[clap(long, value_parser = Regex::new)]
+    run_parts: Option<Regex>,
+
     /// The path to the CSV file containing the execution data.
     #[clap(long, default_value = "report.csv")]
     report_path: PathBuf,
@@ -123,6 +139,10 @@ pub struct HostArgs {
     #[arg(long)]
     pub agg_pk_path: Option<PathBuf>,
 
+    /// If specified, loads the halo2 proving key from this path.
+    #[arg(long)]
+    pub halo_pk_path: Option<PathBuf>,
+
     #[arg(long, default_value_t = false)]
     pub skip_comparison: bool,
 
@@ -241,17 +261,22 @@ pub async fn run_reth_benchmark(args: HostArgs, openvm_client_eth_elf: &[u8]) ->
     println!("CUDA Backend Enabled");
 
     let vm_config = reth_vm_config(app_log_blowup);
+    println!("made vm_config");
     let app_config = args.benchmark.app_config(vm_config.clone());
+    println!("made app config");
     let sdk = Sdk::new(app_config.clone())?
         .with_agg_config(args.benchmark.agg_config())
         .with_agg_tree_config(args.benchmark.agg_tree_config);
+    println!("made sdk");
 
     if args.app_pk_path.is_some() != args.agg_pk_path.is_some() {
         eyre::bail!("app_pk_path and agg_pk_path must be provided together");
     }
     if let Some(app_pk_path) = args.app_pk_path {
+        println!("app_pk: loading from cache at {}", app_pk_path.display());
         let app_pk: AppProvingKey<SdkVmConfig> = read_object_from_file(app_pk_path)?;
         let agg_pk_path = args.agg_pk_path.unwrap();
+        println!("agg_pk: loading from cache at {}", agg_pk_path.display());
         let agg_pk: AggProvingKey = read_object_from_file(agg_pk_path)?;
         let vm_config_loaded = app_pk.app_vm_pk.vm_config.clone();
         let vm_config_json =
@@ -264,6 +289,14 @@ pub async fn run_reth_benchmark(args: HostArgs, openvm_client_eth_elf: &[u8]) ->
         );
         sdk.set_app_pk(app_pk).map_err(|_| eyre::eyre!("failed to set app pk"))?;
         sdk.set_agg_pk(agg_pk).map_err(|_| eyre::eyre!("failed to set agg pk"))?;
+    } else {
+        println!("not loading app_pk");
+    }
+    #[cfg(feature = "evm-verify")]
+    if let Some(halo_pk_path) = args.halo_pk_path.take() {
+        println!("halo_pk: loading from cache at {}", halo_pk_path.display());
+        let halo_pk: Halo2ProvingKey = read_object_from_file(halo_pk_path)?;
+        sdk.set_halo2_pk(halo_pk).map_err(|_| eyre::eyre!("failed to set halo pk"))?;
     }
 
     let elf = Elf::decode(openvm_client_eth_elf, MEM_SIZE as u32)?;
@@ -359,8 +392,28 @@ pub async fn run_reth_benchmark(args: HostArgs, openvm_client_eth_elf: &[u8]) ->
                             "halo2_wrapper_k: {}",
                             halo2_pk.wrapper.pinning.metadata.config_params.k
                         );
-                        let proof = prover.prove_evm(stdin)?;
-                        let block_hash = &proof.user_public_values;
+
+                        let app_proof = load_or_run(
+                            args.fixtures_path.as_ref(),
+                            args.run_parts.as_ref(),
+                            "app_proof",
+                            || prover.stark_prover.app_prover.prove(stdin),
+                        )?;
+
+                        let root_proof = load_or_run(
+                            args.fixtures_path.as_ref(),
+                            args.run_parts.as_ref(),
+                            "root_proof",
+                            || prover.stark_prover.agg_prover.generate_root_proof(app_proof),
+                        )?;
+
+                        let evm_proof = load_or_run::<_, eyre::Report, _>(
+                            args.fixtures_path.as_ref(),
+                            args.run_parts.as_ref(),
+                            "evm_proof",
+                            || Ok(prover.halo2_prover.prove_for_evm(&root_proof)),
+                        )?;
+                        let block_hash = &evm_proof.user_public_values;
                         println!("block_hash (prove_evm): {}", ToHexExt::encode_hex(block_hash));
                     }
                     BenchMode::GenerateFixtures => {
@@ -385,6 +438,35 @@ pub async fn run_reth_benchmark(args: HostArgs, openvm_client_eth_elf: &[u8]) ->
                         agg_pk_path.push("agg_pk.bitcode");
                         fs::write(agg_pk_path, bitcode::serialize(sdk.agg_pk())?)?;
                     }
+                    #[cfg(feature = "evm-verify")]
+                    BenchMode::GenerateFixturesEvm => {
+                        let mut prover = sdk.evm_prover(elf)?.with_program_name(program_name);
+                        let app_proof = prover.stark_prover.app_prover.prove(stdin)?;
+                        let fixture_path = args.fixtures_path.unwrap();
+
+                        let mut app_proof_path = fixture_path.clone();
+                        app_proof_path.push("app_proof.bitcode");
+                        fs::write(app_proof_path, bitcode::serialize(&app_proof)?)?;
+
+                        let root_proof =
+                            prover.stark_prover.agg_prover.generate_root_proof(app_proof)?;
+
+                        let mut root_proof_path = fixture_path.clone();
+                        root_proof_path.push("root_proof.bitcode");
+                        fs::write(root_proof_path, bitcode::serialize(&root_proof)?)?;
+
+                        let mut app_pk_path = fixture_path.clone();
+                        app_pk_path.push("app_pk.bitcode");
+                        fs::write(app_pk_path, bitcode::serialize(sdk.app_pk())?)?;
+
+                        let mut agg_pk_path = fixture_path.clone();
+                        agg_pk_path.push("agg_pk.bitcode");
+                        fs::write(agg_pk_path, bitcode::serialize(sdk.agg_pk())?)?;
+
+                        let mut halo_pk_path = fixture_path.clone();
+                        halo_pk_path.push("halo_pk.bitcode");
+                        fs::write(halo_pk_path, bitcode::serialize(sdk.halo2_pk())?)?;
+                    }
                     _ => {
                         // This case is handled earlier and should not reach here
                         unreachable!();
@@ -398,6 +480,69 @@ pub async fn run_reth_benchmark(args: HostArgs, openvm_client_eth_elf: &[u8]) ->
     Ok(())
 }
 
+/// Load a `bitcode`-serialized object named `{name}.bitcode` from `fixtures_path`, falling back
+/// to `run_fn` if loading is skipped or fails. On a successful run the result is written back to
+/// the fixture so subsequent runs can skip the step. Loading is skipped (with a logged reason) if
+/// `fixtures_path` is `None`, the file is missing, the `run_parts` regex matches `name`, or
+/// deserialization fails.
+#[cfg(feature = "evm-verify")]
+fn load_or_run<T, E, F>(
+    fixtures_path: Option<&PathBuf>,
+    run_parts: Option<&Regex>,
+    name: &str,
+    run_fn: F,
+) -> eyre::Result<T>
+where
+    T: Serialize + for<'de> Deserialize<'de>,
+    E: Into<eyre::Report>,
+    F: FnOnce() -> Result<T, E>,
+{
+    let f = || {
+
+    let file_path = fixtures_path.map(|p| p.join(format!("{name}.bitcode")));
+
+    let skip_reason: Option<String> = if fixtures_path.is_none() {
+        Some("fixtures_path is not set".to_string())
+    } else if run_parts.is_some_and(|re| re.is_match(name)) {
+        Some(format!("run_parts regex matches '{name}'"))
+    } else {
+        let p = file_path.as_ref().unwrap();
+        (!p.exists()).then(|| format!("{} does not exist", p.display()))
+    };
+
+    if skip_reason.is_none() {
+        let p = file_path.as_ref().unwrap();
+        let load_result = std::fs::read(p)
+            .map_err(eyre::Report::from)
+            .and_then(|bytes| bitcode::deserialize::<T>(&bytes).map_err(eyre::Report::from));
+        match load_result {
+            Ok(loaded) => {
+                println!("{name}: loading from cache at {}", p.display());
+                return Ok(loaded);
+            }
+            Err(err) => {
+                println!("{name}: failed to load from {} ({err}), running instead", p.display());
+            }
+        }
+    } else if let Some(reason) = skip_reason {
+        println!("{name}: running, reason: {reason}");
+    }
+
+    let result = run_fn().map_err(Into::into)?;
+    if let Some(p) = file_path.as_ref() {
+        std::fs::write(p, bitcode::serialize(&result)?)?;
+        println!("{name}: wrote to {}", p.display());
+    }
+    Ok(result)
+    };
+
+    let t0 = std::time::Instant::now();
+    let res = f();
+    println!("{name} took {:.3}", t0.elapsed().as_secs_f64());
+
+    res
+}
+
 fn try_load_input_from_cache(
     cache_dir: Option<&PathBuf>,
     chain_id: u64,

run.sh

@@ -101,6 +101,12 @@ PROFILE_OVERRIDE=""
 BLOCK_NUMBER_OVERRIDE=""
 USE_CUDA=false
 CUDA_REASON=""
+SKIP_COMPARISON=true
+FIXTURES_PATH=""
+APP_PK_PATH=""
+AGG_PK_PATH=""
+HALO_PK_PATH=""
+RUN_PARTS=""
 
 while [[ $# -gt 0 ]]; do
     case $1 in
@@ -125,6 +131,34 @@ while [[ $# -gt 0 ]]; do
             PREIMAGE_CACHE_NIBBLES="$2"
             shift 2
             ;;
+        --skip-comparison)
+            SKIP_COMPARISON=true
+            shift
+            ;;
+        --no-skip-comparison)
+            SKIP_COMPARISON=false
+            shift
+            ;;
+        --fixtures-path)
+            FIXTURES_PATH="$2"
+            shift 2
+            ;;
+        --app-pk-path)
+            APP_PK_PATH="$2"
+            shift 2
+            ;;
+        --agg-pk-path)
+            AGG_PK_PATH="$2"
+            shift 2
+            ;;
+        --halo-pk-path)
+            HALO_PK_PATH="$2"
+            shift 2
+            ;;
+        --run-parts)
+            RUN_PARTS="$2"
+            shift 2
+            ;;
         *)
             echo "Unknown argument: $1"
             exit 1
@@ -154,12 +188,14 @@ mkdir -p rpc-cache
 source .env
 
 cd "$WORKDIR/bin/stateless-guest"
-OPENVM_RUST_TOOLCHAIN=nightly-2026-01-01 cargo openvm build
-mkdir -p ../reth-benchmark/elf
 SRC="target/riscv32im-risc0-zkvm-elf/release/openvm-stateless-guest"
 DEST="../reth-benchmark/elf/openvm-stateless-guest"
 
-if [ ! -f "$DEST" ] || ! cmp -s "$SRC" "$DEST"; then
+if [ -f "$DEST" ]; then
+    echo "Stateless guest ELF already exists at $DEST, skipping build."
+else
+    OPENVM_RUST_TOOLCHAIN=nightly-2026-01-01 cargo openvm build
+    mkdir -p ../reth-benchmark/elf
     cp "$SRC" "$DEST"
 fi
 cd "$WORKDIR"
@@ -187,7 +223,7 @@ VPMM_PAGES=$((12 * $MAX_SEGMENT_LENGTH/ $VPMM_PAGE_SIZE))
 if [ "$USE_CUDA" = "true" ]; then
     FEATURES="$FEATURES,cuda"
 fi
-if [ "$MODE" = "prove-evm" ]; then
+if [ "$MODE" = "prove-evm" ] || [ "$MODE" = "generate-fixtures-evm" ]; then
     FEATURES="$FEATURES,evm-verify"
 fi
 
@@ -217,6 +253,27 @@ else
     TARGET_DIR="$PROFILE"
 fi
 
+EXTRA_ARGS=()
+if [ "$SKIP_COMPARISON" = "true" ]; then
+    EXTRA_ARGS+=(--skip-comparison)
+fi
+if [ -n "$FIXTURES_PATH" ]; then
+    mkdir -p "$FIXTURES_PATH"
+    EXTRA_ARGS+=(--fixtures-path "$FIXTURES_PATH")
+fi
+if [ -n "$APP_PK_PATH" ]; then
+    EXTRA_ARGS+=(--app-pk-path "$APP_PK_PATH")
+fi
+if [ -n "$AGG_PK_PATH" ]; then
+    EXTRA_ARGS+=(--agg-pk-path "$AGG_PK_PATH")
+fi
+if [ -n "$HALO_PK_PATH" ]; then
+    EXTRA_ARGS+=(--halo-pk-path "$HALO_PK_PATH")
+fi
+if [ -n "$RUN_PARTS" ]; then
+    EXTRA_ARGS+=(--run-parts "$RUN_PARTS")
+fi
+
 RUST_LOG="info,p3_=warn" OUTPUT_PATH="metrics.json" VPMM_PAGES=$VPMM_PAGES VPMM_PAGE_SIZE=$VPMM_PAGE_SIZE ./target/$TARGET_DIR/$BIN_NAME \
 --kzg-params-dir $PARAMS_DIR \
 --mode $MODE \
@@ -231,4 +288,6 @@ RUST_LOG="info,p3_=warn" OUTPUT_PATH="metrics.json" VPMM_PAGES=$VPMM_PAGES VPMM_
 --max-segment-length $MAX_SEGMENT_LENGTH \
 --segment-max-cells $SEGMENT_MAX_CELLS \
 --num-children-leaf 1 \
---num-children-internal 3
+--num-children-internal 3 \
+"${EXTRA_ARGS[@]}"
+