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_conjectured_security(
103                app_log_blowup,
104            )
105            .into(),
106            app_vm_config,
107            leaf_fri_params: FriParameters::standard_with_100_bits_conjectured_security(
108                leaf_log_blowup,
109            )
110            .into(),
111            compiler_options: CompilerOptions {
112                enable_cycle_tracker: self.profiling,
113                ..Default::default()
114            },
115        }
116    }
117
118    pub fn agg_config(&self) -> AggregationConfig {
119        let leaf_log_blowup = self.leaf_log_blowup.unwrap_or(DEFAULT_LEAF_LOG_BLOWUP);
120        let internal_log_blowup = self
121            .internal_log_blowup
122            .unwrap_or(DEFAULT_INTERNAL_LOG_BLOWUP);
123        let root_log_blowup = self.root_log_blowup.unwrap_or(DEFAULT_ROOT_LOG_BLOWUP);
124
125        let [leaf_fri_params, internal_fri_params, root_fri_params] =
126            [leaf_log_blowup, internal_log_blowup, root_log_blowup]
127                .map(FriParameters::standard_with_100_bits_conjectured_security);
128
129        AggregationConfig {
130            leaf_fri_params,
131            internal_fri_params,
132            root_fri_params,
133            profiling: self.profiling,
134            compiler_options: CompilerOptions {
135                enable_cycle_tracker: self.profiling,
136                ..Default::default()
137            },
138            root_max_constraint_degree: root_fri_params.max_constraint_degree(),
139            ..Default::default()
140        }
141    }
142
143    pub fn halo2_config(&self) -> Halo2Config {
144        Halo2Config {
145            verifier_k: self.halo2_outer_k.unwrap_or(DEFAULT_HALO2_VERIFIER_K),
146            wrapper_k: self.halo2_wrapper_k,
147            profiling: self.profiling,
148        }
149    }
150
151    pub fn build_bench_program<VC>(
152        &self,
153        program_name: &str,
154        vm_config: &VC,
155        init_file_name: Option<&str>,
156    ) -> Result<Elf>
157    where
158        VC: VmConfig<SC>,
159    {
160        let profile = if self.profiling {
161            "profiling"
162        } else {
163            "release"
164        }
165        .to_string();
166        let manifest_dir = get_programs_dir().join(program_name);
167        vm_config.write_to_init_file(&manifest_dir, init_file_name)?;
168        build_elf(&manifest_dir, profile)
169    }
170
171    pub fn bench_from_exe<VB, VC>(
172        &self,
173        bench_name: impl ToString,
174        vm_config: VC,
175        exe: impl Into<ExecutableFormat>,
176        input_stream: StdIn,
177    ) -> Result<()>
178    where
179        VB: VmBuilder<Poseidon2Engine, VmConfig = VC, RecordArena = RA> + Clone + Default,
180        VC: VmExecutionConfig<F> + VmConfig<SC> + TranspilerConfig<F>,
181        <VC as VmExecutionConfig<F>>::Executor:
182            Executor<F> + MeteredExecutor<F> + PreflightExecutor<F, RA>,
183    {
184        let app_config = self.app_config(vm_config);
185        bench_from_exe::<Poseidon2Engine, VB, DefaultNativeBuilder>(
186            bench_name,
187            app_config,
188            exe,
189            input_stream,
190            #[cfg(not(feature = "aggregation"))]
191            None,
192            #[cfg(feature = "aggregation")]
193            Some(self.agg_config().leaf_vm_config()),
194        )
195    }
196}
197
198/// 1. Generate proving key from config.
199/// 2. Commit to the exe by generating cached trace for program.
200/// 3. Executes runtime
201/// 4. Generate trace
202/// 5. Generate STARK proofs for each segment (segmentation is determined by `config`)
203/// 6. Verify STARK proofs.
204///
205/// Returns the data necessary for proof aggregation.
206pub fn bench_from_exe<E, VB, NativeBuilder>(
207    bench_name: impl ToString,
208    app_config: AppConfig<VB::VmConfig>,
209    exe: impl Into<ExecutableFormat>,
210    input_stream: StdIn,
211    leaf_vm_config: Option<NativeConfig>,
212) -> Result<()>
213where
214    E: StarkFriEngine<SC = SC>,
215    VB: VmBuilder<E> + Clone + Default,
216    VB::VmConfig: TranspilerConfig<F>,
217    <VB::VmConfig as VmExecutionConfig<F>>::Executor:
218        Executor<F> + MeteredExecutor<F> + PreflightExecutor<F, VB::RecordArena>,
219    NativeBuilder: VmBuilder<E, VmConfig = NativeConfig> + Clone + Default,
220    <NativeConfig as VmExecutionConfig<F>>::Executor:
221        PreflightExecutor<F, <NativeBuilder as VmBuilder<E>>::RecordArena>,
222{
223    let bench_name = bench_name.to_string();
224    let sdk = GenericSdk::<E, VB, NativeBuilder>::new(app_config.clone())?;
225    // 1. Generate proving key from config.
226    let (app_pk, app_vk) = info_span!("keygen", group = &bench_name).in_scope(|| sdk.app_keygen());
227    // 3. Executes runtime
228    // 4. Generate trace
229    // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with
230    //    timer.
231    let mut prover = sdk.app_prover(exe)?.with_program_name(bench_name);
232    let app_proof = prover.prove(input_stream)?;
233    // 6. Verify STARK proofs, including boundary conditions.
234    verify_app_proof(&app_vk, &app_proof)?;
235    if let Some(leaf_vm_config) = leaf_vm_config {
236        let leaf_vm_pk = _leaf_keygen(app_config.leaf_fri_params.fri_params, leaf_vm_config)?;
237        let vk = leaf_vm_pk.vm_pk.get_vk();
238        let mut leaf_prover = new_local_prover(
239            sdk.native_builder().clone(),
240            &leaf_vm_pk,
241            app_pk.leaf_committed_exe.exe.clone(),
242        )?;
243        let leaf_controller = LeafProvingController {
244            num_children: AggregationTreeConfig::default().num_children_leaf,
245        };
246        let leaf_proofs = leaf_controller.generate_proof(&mut leaf_prover, &app_proof)?;
247        for proof in leaf_proofs {
248            verify_single(&leaf_prover.vm.engine, &vk, &proof)?;
249        }
250    }
251    Ok(())
252}
openvm_benchmarks_prove/util.rs

openvm_benchmarks_prove/
util.rs
