openvm_benchmarks_prove/
util.rs

Help
1use std::{path::PathBuf, sync::Arc};
2
3use clap::{command, Parser};
4use eyre::Result;
5use openvm_benchmarks_utils::{build_elf, get_programs_dir};
6use openvm_circuit::arch::{instructions::exe::VmExe, DefaultSegmentationStrategy, VmConfig};
7use openvm_native_circuit::NativeConfig;
8use openvm_native_compiler::conversion::CompilerOptions;
9use openvm_sdk::{
10    commit::commit_app_exe,
11    config::{
12        AggConfig, AggStarkConfig, AppConfig, Halo2Config, DEFAULT_APP_LOG_BLOWUP,
13        DEFAULT_INTERNAL_LOG_BLOWUP, DEFAULT_LEAF_LOG_BLOWUP, DEFAULT_ROOT_LOG_BLOWUP,
14    },
15    keygen::{leaf_keygen, AppProvingKey},
16    prover::{
17        vm::local::VmLocalProver, AppProver, LeafProvingController, DEFAULT_NUM_CHILDREN_LEAF,
18    },
19    Sdk, StdIn,
20};
21use openvm_stark_backend::utils::metrics_span;
22use openvm_stark_sdk::{
23    config::{
24        baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
25        FriParameters,
26    },
27    engine::StarkFriEngine,
28    openvm_stark_backend::Chip,
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 segment length for continuations
66    #[arg(short, long, alias = "max_segment_length")]
67    pub max_segment_length: Option<usize>,
68
69    /// Whether to execute with additional profiling metric collection
70    #[arg(long)]
71    pub profiling: bool,
72}
73
74impl BenchmarkCli {
75    pub fn app_config<VC: VmConfig<BabyBear>>(&self, mut app_vm_config: VC) -> AppConfig<VC> {
76        let app_log_blowup = self.app_log_blowup.unwrap_or(DEFAULT_APP_LOG_BLOWUP);
77        let leaf_log_blowup = self.leaf_log_blowup.unwrap_or(DEFAULT_LEAF_LOG_BLOWUP);
78
79        app_vm_config.system_mut().profiling = self.profiling;
80        if let Some(max_segment_length) = self.max_segment_length {
81            app_vm_config
82                .system_mut()
83                .set_segmentation_strategy(Arc::new(
84                    DefaultSegmentationStrategy::new_with_max_segment_len(max_segment_length),
85                ));
86        }
87        AppConfig {
88            app_fri_params: FriParameters::standard_with_100_bits_conjectured_security(
89                app_log_blowup,
90            )
91            .into(),
92            app_vm_config,
93            leaf_fri_params: FriParameters::standard_with_100_bits_conjectured_security(
94                leaf_log_blowup,
95            )
96            .into(),
97            compiler_options: CompilerOptions {
98                enable_cycle_tracker: self.profiling,
99                ..Default::default()
100            },
101        }
102    }
103
104    pub fn agg_config(&self) -> AggConfig {
105        let leaf_log_blowup = self.leaf_log_blowup.unwrap_or(DEFAULT_LEAF_LOG_BLOWUP);
106        let internal_log_blowup = self
107            .internal_log_blowup
108            .unwrap_or(DEFAULT_INTERNAL_LOG_BLOWUP);
109        let root_log_blowup = self.root_log_blowup.unwrap_or(DEFAULT_ROOT_LOG_BLOWUP);
110
111        let [leaf_fri_params, internal_fri_params, root_fri_params] =
112            [leaf_log_blowup, internal_log_blowup, root_log_blowup]
113                .map(FriParameters::standard_with_100_bits_conjectured_security);
114
115        AggConfig {
116            agg_stark_config: AggStarkConfig {
117                leaf_fri_params,
118                internal_fri_params,
119                root_fri_params,
120                profiling: self.profiling,
121                compiler_options: CompilerOptions {
122                    enable_cycle_tracker: self.profiling,
123                    ..Default::default()
124                },
125                root_max_constraint_degree: root_fri_params.max_constraint_degree(),
126                ..Default::default()
127            },
128            halo2_config: Halo2Config {
129                verifier_k: self.halo2_outer_k.unwrap_or(23),
130                wrapper_k: self.halo2_wrapper_k,
131                profiling: self.profiling,
132            },
133        }
134    }
135
136    pub fn build_bench_program(&self, program_name: &str) -> Result<Elf> {
137        let profile = if self.profiling {
138            "profiling"
139        } else {
140            "release"
141        }
142        .to_string();
143        let manifest_dir = get_programs_dir().join(program_name);
144        build_elf(&manifest_dir, profile)
145    }
146
147    pub fn bench_from_exe<VC>(
148        &self,
149        bench_name: impl ToString,
150        vm_config: VC,
151        exe: impl Into<VmExe<F>>,
152        input_stream: StdIn,
153    ) -> Result<()>
154    where
155        VC: VmConfig<F>,
156        VC::Executor: Chip<SC>,
157        VC::Periphery: Chip<SC>,
158    {
159        let app_config = self.app_config(vm_config);
160        bench_from_exe::<VC, BabyBearPoseidon2Engine>(
161            bench_name,
162            app_config,
163            exe,
164            input_stream,
165            #[cfg(not(feature = "aggregation"))]
166            None,
167            #[cfg(feature = "aggregation")]
168            Some(self.agg_config().agg_stark_config.leaf_vm_config()),
169        )
170    }
171}
172
173/// 1. Generate proving key from config.
174/// 2. Commit to the exe by generating cached trace for program.
175/// 3. Executes runtime
176/// 4. Generate trace
177/// 5. Generate STARK proofs for each segment (segmentation is determined by `config`)
178/// 6. Verify STARK proofs.
179///
180/// Returns the data necessary for proof aggregation.
181pub fn bench_from_exe<VC, E: StarkFriEngine<SC>>(
182    bench_name: impl ToString,
183    app_config: AppConfig<VC>,
184    exe: impl Into<VmExe<F>>,
185    input_stream: StdIn,
186    leaf_vm_config: Option<NativeConfig>,
187) -> Result<()>
188where
189    VC: VmConfig<F>,
190    VC::Executor: Chip<SC>,
191    VC::Periphery: Chip<SC>,
192{
193    let bench_name = bench_name.to_string();
194    // 1. Generate proving key from config.
195    let app_pk = info_span!("keygen", group = &bench_name).in_scope(|| {
196        metrics_span("keygen_time_ms", || {
197            AppProvingKey::keygen(app_config.clone())
198        })
199    });
200    // 2. Commit to the exe by generating cached trace for program.
201    let committed_exe = info_span!("commit_exe", group = &bench_name).in_scope(|| {
202        metrics_span("commit_exe_time_ms", || {
203            commit_app_exe(app_config.app_fri_params.fri_params, exe)
204        })
205    });
206    // 3. Executes runtime
207    // 4. Generate trace
208    // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with timer.
209    let app_vk = app_pk.get_app_vk();
210    let prover =
211        AppProver::<VC, E>::new(app_pk.app_vm_pk, committed_exe).with_program_name(bench_name);
212    let app_proof = prover.generate_app_proof(input_stream);
213    // 6. Verify STARK proofs, including boundary conditions.
214    let sdk = Sdk::new();
215    sdk.verify_app_proof(&app_vk, &app_proof)
216        .expect("Verification failed");
217    if let Some(leaf_vm_config) = leaf_vm_config {
218        let leaf_vm_pk = leaf_keygen(app_config.leaf_fri_params.fri_params, leaf_vm_config);
219        let leaf_prover =
220            VmLocalProver::<SC, NativeConfig, E>::new(leaf_vm_pk, app_pk.leaf_committed_exe);
221        let leaf_controller = LeafProvingController {
222            num_children: DEFAULT_NUM_CHILDREN_LEAF,
223        };
224        leaf_controller.generate_proof(&leaf_prover, &app_proof);
225    }
226    Ok(())
227}
openvm_benchmarks_prove/util.rs

openvm_benchmarks_prove/
util.rs
