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

openvm_benchmarks_prove/
util.rs
