From bffe7fb034673bd53d2341101e01b521dfdfefe7 Mon Sep 17 00:00:00 2001 From: eureka-cpu Date: Thu, 14 Nov 2024 00:44:23 -0800 Subject: [PATCH] add command line options, add test case --- Cargo.lock | 1 + cli/Cargo.toml | 1 + cli/src/command.rs | 14 ++++- cli/src/estimate.rs | 127 +++++++++++++++++++++++++++++++++++++++----- cli/src/main.rs | 7 ++- 5 files changed, 134 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5c903d5..5c02f1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1195,6 +1195,7 @@ dependencies = [ "bonsol-prover", "bonsol-sdk", "byte-unit", + "bytemuck", "bytes", "cargo_toml 0.20.5", "clap 4.5.20", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 6fecfbd..761007e 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -18,6 +18,7 @@ atty = "0.2.14" bincode = "1.3.3" bonsol-prover = { path = "../prover" } bonsol-sdk = { path = "../sdk" } +bytemuck = "1.15.0" hex = "0.4.3" byte-unit = "4.0.19" bytes = "1.4.0" diff --git a/cli/src/command.rs b/cli/src/command.rs index 882e294..23db329 100644 --- a/cli/src/command.rs +++ b/cli/src/command.rs @@ -251,6 +251,16 @@ pub enum Command { } )] elf: String, + + #[arg( + help = "Define the maximum number of cycles a single segment can take as a power of two, must be between 13 and 24 [default: 20usize]", + short = 's', + long, + )] + segment_limit_po2: Option, + + #[arg(help = "Set the maximum number of cycles [default: 16777216u64]", short = 'm', long)] + max_cycles: Option, }, Execute { #[arg(short = 'f', long)] @@ -315,6 +325,8 @@ pub enum ParsedCommand { }, Estimate { elf: PathBuf, + segment_limit_po2: Option, + max_cycles: Option, }, Execute { execution_request_file: Option, @@ -371,7 +383,7 @@ impl TryFrom for ParsedCommand { ), }), Command::Build { zk_program_path } => Ok(ParsedCommand::Build { zk_program_path }), - Command::Estimate { elf } => Ok(ParsedCommand::Estimate { elf: PathBuf::from(elf) }), + Command::Estimate { elf, segment_limit_po2, max_cycles } => Ok(ParsedCommand::Estimate { elf: PathBuf::from(elf), segment_limit_po2, max_cycles }), Command::Execute { execution_request_file, program_id, diff --git a/cli/src/estimate.rs b/cli/src/estimate.rs index ed6a61f..a55b8a6 100644 --- a/cli/src/estimate.rs +++ b/cli/src/estimate.rs @@ -6,32 +6,131 @@ //! specific functions as sov-cycle-tracer does and the //! defaults made into command line args. -use std::{fs, path::PathBuf}; - use anyhow::Result; +use indicatif::{ProgressIterator, ProgressStyle}; +use risc0_binfmt::{MemoryImage, Program}; use risc0_circuit_rv32im::prove::emu::{ - exec::{execute_elf, DEFAULT_SEGMENT_LIMIT_PO2}, - testutil::{NullSyscall, DEFAULT_SESSION_LIMIT}, + exec::{execute, DEFAULT_SEGMENT_LIMIT_PO2}, + testutil::DEFAULT_SESSION_LIMIT, }; +use risc0_zkvm::GUEST_MAX_MEM; +use risc0_zkvm_platform::PAGE_SIZE; + +use self::emu_syscall::BasicSyscall; + +pub fn estimate( + elf: E, + segment_limit_po2: Option, + max_cycles: Option, +) -> Result<()> { + let cycles: usize = get_cycle_count(elf, segment_limit_po2, max_cycles)?; + println!("number of cycles: {cycles}"); + + Ok(()) +} -pub fn estimate(elf: PathBuf) -> Result<()> { - // TODO: We probably want to be able to let the user decide some of this - let cycles: usize = execute_elf( - &fs::read(elf)?, - DEFAULT_SEGMENT_LIMIT_PO2, - DEFAULT_SESSION_LIMIT, - &NullSyscall::default(), +/// Get the total number of cycles by stepping through the ELF using emulation +/// tools from the risc0_circuit_rv32im module. +pub fn get_cycle_count( + elf: E, + segment_limit_po2: Option, + max_cycles: Option, +) -> Result { + execute( + elf.mk_image()?, + segment_limit_po2.unwrap_or(DEFAULT_SEGMENT_LIMIT_PO2), + max_cycles.or(DEFAULT_SESSION_LIMIT), + &BasicSyscall::default(), None, )? .segments .iter() + .progress() .try_fold(0, |acc, s| -> Result { let trace = s.preflight()?; let segment_cycles = trace.pre.cycles.len() + trace.body.cycles.len(); Ok(acc + segment_cycles) - })?; + }) +} - println!("number of cycles: {cycles}"); +/// Helper trait for loading an image from an elf. +pub trait MkImage { + fn mk_image(self) -> Result; +} +impl<'a> MkImage for &'a [u8] { + fn mk_image(self) -> Result { + let program = Program::load_elf(self, GUEST_MAX_MEM as u32)?; + MemoryImage::new(&program, PAGE_SIZE as u32) + } +} - Ok(()) +pub mod emu_syscall { + //! The following is copied from risc0 emu test utils, likely this is okay for our use case since we only want the cycle count. + //! https://github.com/anagrambuild/risc0/blob/eb331d7ee30bc9ccf944bb1ea4835e60e21c25a2/risc0/circuit/rv32im/src/prove/emu/exec/tests.rs#L41 + + use std::cell::RefCell; + + use anyhow::Result; + use risc0_circuit_rv32im::prove::emu::{ + addr::ByteAddr, + exec::{Syscall, SyscallContext}, + }; + use risc0_zkvm_platform::syscall::reg_abi::{REG_A4, REG_A5}; + + #[derive(Default, Clone)] + pub struct BasicSyscallState { + syscall: String, + from_guest: Vec, + into_guest: Vec, + } + + #[derive(Default)] + pub struct BasicSyscall { + state: RefCell, + } + + impl Syscall for BasicSyscall { + fn syscall( + &self, + syscall: &str, + ctx: &mut dyn SyscallContext, + guest_buf: &mut [u32], + ) -> Result<(u32, u32)> { + self.state.borrow_mut().syscall = syscall.to_string(); + let buf_ptr = ByteAddr(ctx.peek_register(REG_A4)?); + let buf_len = ctx.peek_register(REG_A5)?; + self.state.borrow_mut().from_guest = ctx.peek_region(buf_ptr, buf_len)?; + let guest_buf_bytes: &mut [u8] = bytemuck::cast_slice_mut(guest_buf); + let into_guest = &self.state.borrow().into_guest; + guest_buf_bytes[..into_guest.len()].clone_from_slice(into_guest); + Ok((0, 0)) + } + } +} + +#[cfg(test)] +mod estimate_tests { + use anyhow::Result; + use risc0_binfmt::MemoryImage; + use risc0_circuit_rv32im::prove::emu::testutil::basic as basic_test_program; + use risc0_zkvm::PAGE_SIZE; + + use super::MkImage; + use crate::estimate; + + impl MkImage for MemoryImage { + fn mk_image(self) -> Result { + Ok(self) + } + } + + #[test] + fn test_estimate() { + let program = basic_test_program(); + let image = MemoryImage::new(&program, PAGE_SIZE as u32) + .expect("failed to create image from basic program"); + let res = estimate::get_cycle_count(image, None, None); + + assert_eq!(res.ok(), Some(15790)); + } } diff --git a/cli/src/main.rs b/cli/src/main.rs index 017a3a2..f6fd2c3 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,3 +1,4 @@ +use std::fs; use std::io::{self, Read}; use std::path::Path; @@ -60,7 +61,11 @@ async fn main() -> anyhow::Result<()> { } deploy::deploy(rpc, keypair, deploy_args).await } - ParsedCommand::Estimate { elf } => estimate::estimate(elf), + ParsedCommand::Estimate { + elf, + segment_limit_po2, + max_cycles, + } => estimate::estimate(fs::read(elf)?.as_slice(), segment_limit_po2, max_cycles), ParsedCommand::Execute { execution_request_file, program_id,