Add cargo risczero benchmark command

risc0/cargo-risczero/Cargo.toml

@@ -33,7 +33,8 @@ reqwest-middleware = "0.2"
 reqwest-retry = "0.3"
 risc0-build = { workspace = true }
 risc0-r0vm = { workspace = true, optional = true }
-risc0-zkvm = { workspace = true, optional = true }
+risc0-zkvm = { path = "../zkvm", default-features = false, features = ["std", "getrandom"] }
+risc0-zkvm-methods = { path = "../zkvm/methods", default-features = false }
 serde = { version = "1", features = ["derive"] }
 syn = "2.0.38"
 tar = "0.4"

risc0/cargo-risczero/src/bin/cargo-risczero.rs

@@ -30,6 +30,7 @@ fn main() -> Result<()> {
         RisczeroCmd::BuildToolchain(cmd) => cmd.run(),
         RisczeroCmd::Install(cmd) => cmd.run(),
         RisczeroCmd::New(cmd) => cmd.run(),
+        RisczeroCmd::Benchmark(cmd) => cmd.run(),
         #[cfg(feature = "experimental")]
         RisczeroCmd::BuildCrate(build) => build.run(BuildSubcommand::Build),
         #[cfg(feature = "experimental")]

risc0/cargo-risczero/src/commands/benchmark.rs

@@ -0,0 +1,111 @@
+// Copyright 2024 RISC Zero, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::{Duration, Instant};
+
+use risc0_zkvm_methods::{
+    bench::{BenchmarkSpec, SpecWithIters},
+    BENCH_ELF,
+};
+
+use risc0_zkvm::{
+    get_prover_server,
+    recursion::{join, lift},
+    ExecutorEnv, ExecutorImpl, ProverOpts, VerifierContext,
+};
+
+use anyhow::Result;
+use clap::Parser;
+
+/// `cargo risczero benchmark`
+#[derive(Parser)]
+pub struct BenchmarkCommand {
+    /// Number of iterations.
+    #[arg(short, long)]
+    pub iterations: Option<u64>,
+
+    /// Which hash function to use.
+    #[arg(short = 'f', long, default_value_t = String::from("poseidon"), value_parser = ["poseidon", "sha-256"])]
+    pub hashfn: String,
+
+    /// Specify the segment po2.
+    #[arg(short, long, default_value_t = 20)]
+    po2: u32,
+}
+
+impl BenchmarkCommand {
+    /// Execute this command.
+    pub fn run(&self) -> Result<()> {
+        // TODO: Handle the case where the user does not specify the number of iterations
+        let iterations = SpecWithIters(
+            BenchmarkSpec::SimpleLoop,
+            self.iterations.unwrap_or(4 * 1024),
+        );
+        let env = ExecutorEnv::builder()
+            .write(&iterations)?
+            .segment_limit_po2(self.po2)
+            .build()?;
+        let mut exec = ExecutorImpl::from_elf(env, BENCH_ELF)?;
+
+        // Execute
+        let (session, exec_duration) = with_duration(|| exec.run())?;
+
+        let cycles = session.get_cycles()?;
+        let segments = session.resolve()?;
+
+        let opts = ProverOpts::default();
+        let ctx = VerifierContext::default();
+        let prover = get_prover_server(&opts)?;
+
+        let mut lifts = vec![];
+        let mut prove_durations = vec![];
+        let mut lift_durations = vec![];
+
+        // Prove and Lift
+        for segment in segments.iter() {
+            let (receipt, receipt_duration) =
+                with_duration(|| prover.prove_segment(&ctx, segment))?;
+            prove_durations.push(receipt_duration);
+
+            let (lift, lift_duration) = with_duration(|| lift(&receipt))?;
+            lifts.push(lift);
+            lift_durations.push(lift_duration);
+        }
+
+        let mut join_durations = vec![];
+        // Optional Join
+        if segments.len() > 1 {
+            let (_final, duration) = with_duration(|| join(&lifts[0], &lifts[1]))?;
+            join_durations.push(duration);
+        }
+
+        println!("\nSTATS:");
+        println!("cycles:     {}", cycles.1);
+        println!("segments:   {}", segments.len());
+        println!("exec:       {exec_duration:?}");
+        println!("prove:      {prove_durations:?}");
+        println!("lift:       {lift_durations:?}");
+        println!("prove+lift: {:?}", prove_durations[0] + lift_durations[0]);
+        println!("join:       {join_durations:?}");
+
+        Ok(())
+    }
+}
+
+fn with_duration<T, F: FnOnce() -> Result<T>>(f: F) -> Result<(T, Duration)> {
+    let start = Instant::now();
+    let result = f()?;
+    let duration = start.elapsed();
+    Ok((result, duration))
+}

risc0/cargo-risczero/src/commands/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod benchmark;
 #[cfg(feature = "experimental")]
 pub mod build;
 pub mod build_guest;

risc0/cargo-risczero/src/lib.rs

@@ -27,7 +27,8 @@ use clap::{Parser, Subcommand};
 #[cfg(feature = "experimental")]
 use self::commands::build::BuildCommand;
 use self::commands::{
-    build_guest::BuildGuest, build_toolchain::BuildToolchain, install::Install, new::NewCommand,
+    benchmark::BenchmarkCommand, build_guest::BuildGuest, build_toolchain::BuildToolchain,
+    install::Install, new::NewCommand,
 };
 
 #[derive(Parser)]
@@ -58,6 +59,8 @@ pub enum RisczeroCmd {
     Install(Install),
     /// Creates a new risczero starter project.
     New(NewCommand),
+    /// Prove a zkvm program locally and report performance statistics.
+    Benchmark(BenchmarkCommand),
     /// Build a crate for RISC Zero.
     #[cfg(feature = "experimental")]
     BuildCrate(BuildCommand),