openvm_benchmarks_prove/
util.rs

1use std::path::PathBuf;
2
3use clap::{command, Parser};
4use eyre::Result;
5use openvm_benchmarks_utils::{build_elf, get_programs_dir};
6use openvm_circuit::{
7    arch::{
8        verify_single, Executor, MeteredExecutor, PreflightExecutor, SystemConfig, VmBuilder,
9        VmConfig, VmExecutionConfig,
10    },
11    utils::{TestRecordArena as RA, TestStarkEngine as Poseidon2Engine},
12};
13use openvm_native_circuit::{NativeBuilder as DefaultNativeBuilder, NativeConfig};
14use openvm_native_compiler::conversion::CompilerOptions;
15use openvm_sdk::{
16    config::{
17        AggregationConfig, AggregationTreeConfig, AppConfig, Halo2Config, TranspilerConfig,
18        DEFAULT_APP_LOG_BLOWUP, DEFAULT_HALO2_VERIFIER_K, DEFAULT_INTERNAL_LOG_BLOWUP,
19        DEFAULT_LEAF_LOG_BLOWUP, DEFAULT_ROOT_LOG_BLOWUP,
20    },
21    keygen::_leaf_keygen,
22    prover::{verify_app_proof, vm::new_local_prover, LeafProvingController},
23    types::ExecutableFormat,
24    GenericSdk, StdIn,
25};
26use openvm_stark_sdk::{
27    config::{baby_bear_poseidon2::BabyBearPoseidon2Config, FriParameters},
28    engine::StarkFriEngine,
29    p3_baby_bear::BabyBear,
30};
31use openvm_transpiler::elf::Elf;
32use tracing::info_span;
33
34type F = BabyBear;
35type SC = BabyBearPoseidon2Config;
36
37#[derive(Parser, Debug)]
38#[command(allow_external_subcommands = true)]
39pub struct BenchmarkCli {
40    /// Application level log blowup, default set by the benchmark
41    #[arg(short = 'p', long, alias = "app_log_blowup")]
42    pub app_log_blowup: Option<usize>,
43
44    /// Aggregation (leaf) level log blowup, default set by the benchmark
45    #[arg(short = 'g', long, alias = "leaf_log_blowup")]
46    pub leaf_log_blowup: Option<usize>,
47
48    /// Internal level log blowup, default set by the benchmark
49    #[arg(short, long, alias = "internal_log_blowup")]
50    pub internal_log_blowup: Option<usize>,
51
52    /// Root level log blowup, default set by the benchmark
53    #[arg(short, long, alias = "root_log_blowup")]
54    pub root_log_blowup: Option<usize>,
55
56    #[arg(long)]
57    pub halo2_outer_k: Option<usize>,
58
59    #[arg(long)]
60    pub halo2_wrapper_k: Option<usize>,
61
62    #[arg(long)]
63    pub kzg_params_dir: Option<PathBuf>,
64
65    /// Max trace height per chip in segment for continuations
66    #[arg(long, alias = "max_segment_length")]
67    pub max_segment_length: Option<u32>,
68
69    /// Total cells used in all chips in segment for continuations
70    #[arg(long)]
71    pub segment_max_cells: Option<usize>,
72
73    /// Controls the arity (num_children) of the aggregation tree
74    #[command(flatten)]
75    pub agg_tree_config: AggregationTreeConfig,
76
77    /// Whether to execute with additional profiling metric collection
78    #[arg(long)]
79    pub profiling: bool,
80}
81
82impl BenchmarkCli {
83    pub fn app_config<VC>(&self, mut app_vm_config: VC) -> AppConfig<VC>
84    where
85        VC: AsMut<SystemConfig>,
86    {
87        let app_log_blowup = self.app_log_blowup.unwrap_or(DEFAULT_APP_LOG_BLOWUP);
88        let leaf_log_blowup = self.leaf_log_blowup.unwrap_or(DEFAULT_LEAF_LOG_BLOWUP);
89
90        app_vm_config.as_mut().profiling = self.profiling;
91        app_vm_config.as_mut().max_constraint_degree = (1 << app_log_blowup) + 1;
92        if let Some(max_height) = self.max_segment_length {
93            app_vm_config
94                .as_mut()
95                .segmentation_limits
96                .set_max_trace_height(max_height);
97        }
98        if let Some(max_cells) = self.segment_max_cells {
99            app_vm_config.as_mut().segmentation_limits.max_cells = max_cells;
100        }
101        AppConfig {
102            app_fri_params: FriParameters::standard_with_100_bits_security(app_log_blowup).into(),
103            app_vm_config,
104            leaf_fri_params: FriParameters::standard_with_100_bits_security(leaf_log_blowup).into(),
105            compiler_options: CompilerOptions {
106                enable_cycle_tracker: self.profiling,
107                ..Default::default()
108            },
109        }
110    }
111
112    pub fn agg_config(&self) -> AggregationConfig {
113        let leaf_log_blowup = self.leaf_log_blowup.unwrap_or(DEFAULT_LEAF_LOG_BLOWUP);
114        let internal_log_blowup = self
115            .internal_log_blowup
116            .unwrap_or(DEFAULT_INTERNAL_LOG_BLOWUP);
117        let root_log_blowup = self.root_log_blowup.unwrap_or(DEFAULT_ROOT_LOG_BLOWUP);
118
119        let [leaf_fri_params, internal_fri_params, root_fri_params] =
120            [leaf_log_blowup, internal_log_blowup, root_log_blowup]
121                .map(FriParameters::standard_with_100_bits_security);
122
123        AggregationConfig {
124            leaf_fri_params,
125            internal_fri_params,
126            root_fri_params,
127            profiling: self.profiling,
128            compiler_options: CompilerOptions {
129                enable_cycle_tracker: self.profiling,
130                ..Default::default()
131            },
132            root_max_constraint_degree: root_fri_params.max_constraint_degree(),
133            ..Default::default()
134        }
135    }
136
137    pub fn halo2_config(&self) -> Halo2Config {
138        Halo2Config {
139            verifier_k: self.halo2_outer_k.unwrap_or(DEFAULT_HALO2_VERIFIER_K),
140            wrapper_k: self.halo2_wrapper_k,
141            profiling: self.profiling,
142        }
143    }
144
145    pub fn build_bench_program<VC>(
146        &self,
147        program_name: &str,
148        vm_config: &VC,
149        init_file_name: Option<&str>,
150    ) -> Result<Elf>
151    where
152        VC: VmConfig<SC>,
153    {
154        let profile = if self.profiling {
155            "profiling"
156        } else {
157            "release"
158        }
159        .to_string();
160        let manifest_dir = get_programs_dir().join(program_name);
161        vm_config.write_to_init_file(&manifest_dir, init_file_name)?;
162        build_elf(&manifest_dir, profile)
163    }
164
165    pub fn bench_from_exe<VB, VC>(
166        &self,
167        bench_name: impl ToString,
168        vm_config: VC,
169        exe: impl Into<ExecutableFormat>,
170        input_stream: StdIn,
171    ) -> Result<()>
172    where
173        VB: VmBuilder<Poseidon2Engine, VmConfig = VC, RecordArena = RA> + Clone + Default,
174        VC: VmExecutionConfig<F> + VmConfig<SC> + TranspilerConfig<F>,
175        <VC as VmExecutionConfig<F>>::Executor:
176            Executor<F> + MeteredExecutor<F> + PreflightExecutor<F, RA>,
177    {
178        let app_config = self.app_config(vm_config);
179        bench_from_exe::<Poseidon2Engine, VB, DefaultNativeBuilder>(
180            bench_name,
181            app_config,
182            exe,
183            input_stream,
184            #[cfg(not(feature = "aggregation"))]
185            None,
186            #[cfg(feature = "aggregation")]
187            Some(self.agg_config().leaf_vm_config()),
188        )
189    }
190}
191
192/// 1. Generate proving key from config.
193/// 2. Commit to the exe by generating cached trace for program.
194/// 3. Executes runtime
195/// 4. Generate trace
196/// 5. Generate STARK proofs for each segment (segmentation is determined by `config`)
197/// 6. Verify STARK proofs.
198///
199/// Returns the data necessary for proof aggregation.
200pub fn bench_from_exe<E, VB, NativeBuilder>(
201    bench_name: impl ToString,
202    app_config: AppConfig<VB::VmConfig>,
203    exe: impl Into<ExecutableFormat>,
204    input_stream: StdIn,
205    leaf_vm_config: Option<NativeConfig>,
206) -> Result<()>
207where
208    E: StarkFriEngine<SC = SC>,
209    VB: VmBuilder<E> + Clone + Default,
210    VB::VmConfig: TranspilerConfig<F>,
211    <VB::VmConfig as VmExecutionConfig<F>>::Executor:
212        Executor<F> + MeteredExecutor<F> + PreflightExecutor<F, VB::RecordArena>,
213    NativeBuilder: VmBuilder<E, VmConfig = NativeConfig> + Clone + Default,
214    <NativeConfig as VmExecutionConfig<F>>::Executor:
215        PreflightExecutor<F, <NativeBuilder as VmBuilder<E>>::RecordArena>,
216{
217    let bench_name = bench_name.to_string();
218    let sdk = GenericSdk::<E, VB, NativeBuilder>::new(app_config.clone())?;
219    // 1. Generate proving key from config.
220    let (app_pk, app_vk) = info_span!("keygen", group = &bench_name).in_scope(|| sdk.app_keygen());
221    // 3. Executes runtime
222    // 4. Generate trace
223    // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with
224    //    timer.
225    let mut prover = sdk.app_prover(exe)?.with_program_name(bench_name);
226    let app_proof = prover.prove(input_stream)?;
227    // 6. Verify STARK proofs, including boundary conditions.
228    verify_app_proof(&app_vk, &app_proof)?;
229    if let Some(leaf_vm_config) = leaf_vm_config {
230        let leaf_vm_pk = _leaf_keygen(app_config.leaf_fri_params.fri_params, leaf_vm_config)?;
231        let vk = leaf_vm_pk.vm_pk.get_vk();
232        let mut leaf_prover = new_local_prover(
233            sdk.native_builder().clone(),
234            &leaf_vm_pk,
235            app_pk.leaf_committed_exe.exe.clone(),
236        )?;
237        let leaf_controller = LeafProvingController {
238            num_children: AggregationTreeConfig::default().num_children_leaf,
239        };
240        let leaf_proofs = leaf_controller.generate_proof(&mut leaf_prover, &app_proof)?;
241        for proof in leaf_proofs {
242            verify_single(&leaf_prover.vm.engine, &vk, &proof)?;
243        }
244    }
245    Ok(())
246}
openvm_benchmarks_prove/util.rs

openvm_benchmarks_prove/
util.rs
