openvm_sdk/keygen/
mod.rs

1use std::sync::Arc;
2
3use derivative::Derivative;
4use openvm_circuit::{
5    arch::{AirInventoryError, SystemConfig, VirtualMachine, VirtualMachineError, VmCircuitConfig},
6    system::memory::dimensions::MemoryDimensions,
7};
8use openvm_continuations::verifier::{
9    internal::InternalVmVerifierConfig, leaf::LeafVmVerifierConfig, root::RootVmVerifierConfig,
10};
11use openvm_native_circuit::{NativeConfig, NativeCpuBuilder, NATIVE_MAX_TRACE_HEIGHTS};
12use openvm_native_compiler::ir::DIGEST_SIZE;
13use openvm_stark_backend::{
14    config::Val,
15    engine::StarkEngine,
16    p3_field::{BasedVectorSpace, PrimeField32, TwoAdicField},
17};
18use openvm_stark_sdk::{
19    config::{
20        baby_bear_poseidon2::BabyBearPoseidon2Engine,
21        baby_bear_poseidon2_root::BabyBearPoseidon2RootEngine, FriParameters,
22    },
23    engine::StarkFriEngine,
24    openvm_stark_backend::{
25        config::{Com, StarkGenericConfig},
26        keygen::types::MultiStarkVerifyingKey,
27        proof::Proof,
28    },
29};
30use serde::{Deserialize, Serialize};
31use tracing::{info_span, instrument};
32#[cfg(feature = "evm-prove")]
33use {
34    openvm_continuations::static_verifier::StaticVerifierPvHandler,
35    openvm_native_recursion::halo2::{
36        utils::Halo2ParamsReader, verifier::Halo2VerifierProvingKey,
37        wrapper::Halo2WrapperProvingKey,
38    },
39};
40
41#[cfg(feature = "evm-prove")]
42use crate::config::Halo2Config;
43use crate::{
44    commit::VmCommittedExe,
45    config::{AggregationConfig, AppConfig},
46    keygen::{
47        dummy::{compute_root_proof_heights, dummy_internal_proof_riscv_app_vm},
48        perm::AirIdPermutation,
49    },
50    prover::vm::types::VmProvingKey,
51    util::check_max_constraint_degrees,
52    RootSC, SC,
53};
54
55pub mod asm;
56pub(crate) mod dummy;
57pub mod perm;
58#[cfg(feature = "evm-prove")]
59pub mod static_verifier;
60
61/// This is lightweight to clone as it contains smart pointers to the proving keys.
62#[derive(Clone, Serialize, Deserialize)]
63pub struct AppProvingKey<VC> {
64    /// The committed executable of the leaf verifier program that verifies proofs of the App VM
65    /// circuit. The App VM circuit constraints are statically compiled into this executable.
66    pub leaf_committed_exe: Arc<VmCommittedExe<SC>>,
67    pub leaf_fri_params: FriParameters,
68    pub app_vm_pk: Arc<VmProvingKey<SC, VC>>,
69}
70
71#[derive(Clone, Serialize, Deserialize)]
72pub struct AppVerifyingKey {
73    /// We store the FRI parameters used to generate the proof separately.
74    pub fri_params: FriParameters,
75    /// STARK backend verifying key
76    pub vk: MultiStarkVerifyingKey<SC>,
77    pub memory_dimensions: MemoryDimensions,
78}
79
80/// The STARK proving keys necessary for aggregation of app proofs into a single aggregate STARK
81/// proof.
82///
83/// This is lightweight to clone as it contains smart pointers to the proving keys.
84#[derive(Clone, Serialize, Deserialize)]
85pub struct AggProvingKey {
86    pub leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
87    pub internal_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
88    pub internal_committed_exe: Arc<VmCommittedExe<SC>>,
89    pub root_verifier_pk: RootVerifierProvingKey,
90}
91
92#[derive(Clone, Serialize, Deserialize)]
93pub struct AggVerifyingKey {
94    pub(super) leaf_fri_params: FriParameters,
95    pub(super) leaf_vk: MultiStarkVerifyingKey<SC>,
96    /// FRI parameters used to generate the last internal proof.
97    pub(super) internal_fri_params: FriParameters,
98    pub(super) internal_vk: MultiStarkVerifyingKey<SC>,
99    pub(super) internal_verifier_program_commit: Com<SC>,
100}
101
102/// Attention: the serialized size of this struct is VERY large, usually >10GB.
103///
104/// This is lightweight to clone as it contains smart pointers to the proving keys.
105#[cfg(feature = "evm-prove")]
106#[derive(Clone, Serialize, Deserialize)]
107pub struct Halo2ProvingKey {
108    /// Static verifier to verify a stark proof of the root verifier.
109    pub verifier: Arc<Halo2VerifierProvingKey>,
110    /// Wrapper circuit to verify static verifier and reduce the verification costs in the final
111    /// proof.
112    pub wrapper: Arc<Halo2WrapperProvingKey>,
113    /// Whether to collect detailed profiling metrics
114    pub profiling: bool,
115}
116
117impl<VC> AppProvingKey<VC>
118where
119    VC: Clone + VmCircuitConfig<SC> + AsRef<SystemConfig>,
120{
121    pub fn keygen(config: AppConfig<VC>) -> Result<Self, AirInventoryError> {
122        let app_engine = BabyBearPoseidon2Engine::new(config.app_fri_params.fri_params);
123        let app_vm_pk = {
124            let vm_pk = config.app_vm_config.create_airs()?.keygen(&app_engine);
125            assert!(
126                vm_pk.max_constraint_degree
127                    <= config.app_fri_params.fri_params.max_constraint_degree()
128            );
129            check_max_constraint_degrees(
130                config.app_vm_config.as_ref(),
131                &config.app_fri_params.fri_params,
132            );
133            VmProvingKey {
134                fri_params: config.app_fri_params.fri_params,
135                vm_config: config.app_vm_config.clone(),
136                vm_pk,
137            }
138        };
139        check_recursive_verifier_size(
140            &app_vm_pk.vm_pk.get_vk(),
141            config.app_fri_params.fri_params,
142            config.leaf_fri_params.fri_params.log_blowup,
143        );
144        let leaf_committed_exe = {
145            let leaf_engine = BabyBearPoseidon2Engine::new(config.leaf_fri_params.fri_params);
146            let leaf_program = LeafVmVerifierConfig {
147                app_fri_params: config.app_fri_params.fri_params,
148                app_system_config: config.app_vm_config.as_ref().clone(),
149                compiler_options: config.compiler_options,
150            }
151            .build_program(&app_vm_pk.vm_pk.get_vk());
152            Arc::new(VmCommittedExe::commit(
153                leaf_program.into(),
154                leaf_engine.config().pcs(),
155            ))
156        };
157        Ok(Self {
158            leaf_committed_exe,
159            leaf_fri_params: config.leaf_fri_params.fri_params,
160            app_vm_pk: Arc::new(app_vm_pk),
161        })
162    }
163
164    pub fn num_public_values(&self) -> usize {
165        self.app_vm_pk.vm_config.as_ref().num_public_values
166    }
167
168    pub fn get_app_vk(&self) -> AppVerifyingKey {
169        AppVerifyingKey {
170            fri_params: self.app_vm_pk.fri_params,
171            vk: self.app_vm_pk.vm_pk.get_vk(),
172            memory_dimensions: self
173                .app_vm_pk
174                .vm_config
175                .as_ref()
176                .memory_config
177                .memory_dimensions(),
178        }
179    }
180
181    pub fn leaf_verifier_program_commit(&self) -> Com<SC> {
182        self.leaf_committed_exe.get_program_commit()
183    }
184
185    pub fn app_fri_params(&self) -> FriParameters {
186        self.app_vm_pk.fri_params
187    }
188
189    pub fn vm_config(&self) -> &VC {
190        &self.app_vm_pk.vm_config
191    }
192
193    pub fn app_config(&self) -> AppConfig<VC> {
194        AppConfig {
195            app_fri_params: self.app_fri_params().into(),
196            app_vm_config: self.vm_config().clone(),
197            leaf_fri_params: self.leaf_fri_params.into(),
198            compiler_options: Default::default(),
199        }
200    }
201}
202
203/// Try to determine statically if there will be an issue with the recursive verifier size and log
204/// a warning if so.
205///
206/// `next_log_blowup` refers to the `log_blowup` of the next verifier in the chain; this determines
207/// a maximum trace height.
208fn check_recursive_verifier_size<SC: StarkGenericConfig>(
209    vk: &MultiStarkVerifyingKey<SC>,
210    fri_params: FriParameters,
211    next_log_blowup: usize,
212) where
213    Val<SC>: PrimeField32 + TwoAdicField,
214{
215    let vk = &vk.inner;
216
217    // for each round we will compute the pair (total_width, num_airs, num_pts)
218    let mut rounds = vec![];
219
220    // Preprocessed rounds.
221    rounds.extend(
222        vk.per_air
223            .iter()
224            .filter_map(|vk| vk.params.width.preprocessed)
225            .map(|width| (width, 1, 2)),
226    );
227
228    let common_main_total_width = vk
229        .per_air
230        .iter()
231        .map(|vk| vk.params.width.common_main)
232        .sum();
233    rounds.push((common_main_total_width, vk.per_air.len(), 2));
234
235    for vk in vk.per_air.iter() {
236        for &cached_main_width in &vk.params.width.cached_mains {
237            rounds.push((cached_main_width, 1, 2));
238        }
239    }
240
241    let mut after_challenge_rounds = vec![];
242    for vk in vk.per_air.iter() {
243        let widths = &vk.params.width.after_challenge;
244        if widths.len() > after_challenge_rounds.len() {
245            after_challenge_rounds.resize(widths.len(), (0, 0, 2));
246        }
247        for (i, &width) in widths.iter().enumerate() {
248            after_challenge_rounds[i].0 += SC::Challenge::DIMENSION * width;
249            after_challenge_rounds[i].1 += 1;
250        }
251    }
252    rounds.extend(after_challenge_rounds);
253
254    let quotient_round = (
255        vk.per_air
256            .iter()
257            .map(|vk| SC::Challenge::DIMENSION * vk.quotient_degree as usize)
258            .sum(),
259        vk.per_air.len(),
260        1,
261    );
262    rounds.push(quotient_round);
263
264    // This computes the number of rows in the `FRI_REDUCED_OPENING` chip, which is the expected
265    // bottleneck of the recursive verifier.
266    let fri_reduced_opening_trace_height = fri_params.num_queries
267        * rounds
268            .iter()
269            .map(|(total_width, num_airs, total_pts)| total_pts * (total_width + 2 * num_airs))
270            .sum::<usize>();
271    // First check: is FriReducedOpening trace height too large?
272    if fri_reduced_opening_trace_height > (1 << (Val::<SC>::TWO_ADICITY - next_log_blowup)) {
273        tracing::warn!("recursive verifier size may be too large; FriReducedOpening height ({fri_reduced_opening_trace_height}) > {}", 1 << (Val::<SC>::TWO_ADICITY - next_log_blowup));
274    }
275    // Second check: static check for log up soundness constraints using FriReducedOpening trace
276    // height as proxy.
277    let native_max_height: u32 = *NATIVE_MAX_TRACE_HEIGHTS.iter().max().unwrap();
278    if fri_reduced_opening_trace_height as u32 > native_max_height {
279        tracing::warn!(
280            "recursive verifier size may violate log up soundness constraints; {} > {}",
281            fri_reduced_opening_trace_height,
282            native_max_height
283        );
284    }
285}
286
287impl AggProvingKey {
288    #[instrument(
289        name = "agg_stark_keygen",
290        fields(group = "agg_stark_keygen"),
291        skip_all
292    )]
293    pub fn keygen(config: AggregationConfig) -> Result<Self, VirtualMachineError> {
294        let (pk, _) = Self::dummy_proof_and_keygen(config)?;
295        Ok(pk)
296    }
297
298    #[tracing::instrument(level = "info", fields(group = "agg_keygen"), skip_all)]
299    pub(crate) fn dummy_proof_and_keygen(
300        config: AggregationConfig,
301    ) -> Result<(Self, Proof<SC>), VirtualMachineError> {
302        let leaf_vm_config = config.leaf_vm_config();
303        let internal_vm_config = config.internal_vm_config();
304        let root_vm_config = config.root_verifier_vm_config();
305
306        let leaf_engine = BabyBearPoseidon2Engine::new(config.leaf_fri_params);
307        let leaf_vm_pk = {
308            let (_, vm_pk) = VirtualMachine::new_with_keygen(
309                leaf_engine,
310                NativeCpuBuilder,
311                leaf_vm_config.clone(),
312            )?;
313            assert!(vm_pk.max_constraint_degree <= config.leaf_fri_params.max_constraint_degree());
314            check_max_constraint_degrees(&leaf_vm_config.system, &config.leaf_fri_params);
315            Arc::new(VmProvingKey {
316                fri_params: config.leaf_fri_params,
317                vm_config: leaf_vm_config,
318                vm_pk,
319            })
320        };
321        let leaf_vm_vk = leaf_vm_pk.vm_pk.get_vk();
322        check_recursive_verifier_size(
323            &leaf_vm_vk,
324            config.leaf_fri_params,
325            config.internal_fri_params.log_blowup,
326        );
327
328        let internal_engine = BabyBearPoseidon2Engine::new(config.internal_fri_params);
329        let (internal_vm, vm_pk) = VirtualMachine::new_with_keygen(
330            internal_engine,
331            NativeCpuBuilder,
332            internal_vm_config.clone(),
333        )?;
334        check_max_constraint_degrees(&internal_vm_config.system, &config.internal_fri_params);
335        assert!(vm_pk.max_constraint_degree <= config.internal_fri_params.max_constraint_degree());
336        let internal_vm_pk = Arc::new(VmProvingKey {
337            fri_params: config.internal_fri_params,
338            vm_config: internal_vm_config,
339            vm_pk,
340        });
341        let internal_vm_vk = internal_vm_pk.vm_pk.get_vk();
342        check_recursive_verifier_size(
343            &internal_vm_vk,
344            config.internal_fri_params,
345            config.internal_fri_params.log_blowup,
346        );
347
348        let internal_program = InternalVmVerifierConfig {
349            leaf_fri_params: config.leaf_fri_params,
350            internal_fri_params: config.internal_fri_params,
351            compiler_options: config.compiler_options,
352        }
353        .build_program(&leaf_vm_vk, &internal_vm_vk);
354        let internal_committed_exe = Arc::new(VmCommittedExe::commit(
355            internal_program.into(),
356            internal_vm.engine.config().pcs(),
357        ));
358
359        let internal_proof = dummy_internal_proof_riscv_app_vm(
360            leaf_vm_pk.clone(),
361            internal_vm_pk.clone(),
362            internal_committed_exe.clone(),
363            config.max_num_user_public_values,
364        )?;
365
366        let root_verifier_pk = {
367            let mut root_engine = BabyBearPoseidon2RootEngine::new(config.root_fri_params);
368            root_engine.max_constraint_degree = config.root_max_constraint_degree;
369            let root_program = RootVmVerifierConfig {
370                leaf_fri_params: config.leaf_fri_params,
371                internal_fri_params: config.internal_fri_params,
372                num_user_public_values: config.max_num_user_public_values,
373                internal_vm_verifier_commit: internal_committed_exe.get_program_commit().into(),
374                compiler_options: config.compiler_options,
375            }
376            .build_program(&leaf_vm_vk, &internal_vm_vk);
377            let (mut vm, mut vm_pk) = VirtualMachine::new_with_keygen(
378                root_engine,
379                NativeCpuBuilder,
380                root_vm_config.clone(),
381            )?;
382            let root_committed_exe = Arc::new(VmCommittedExe::commit(
383                root_program.into(),
384                vm.engine.config().pcs(),
385            ));
386
387            assert!(vm_pk.max_constraint_degree <= config.root_fri_params.max_constraint_degree());
388
389            let air_heights =
390                compute_root_proof_heights(&mut vm, &root_committed_exe, &internal_proof)?;
391            let root_air_perm = AirIdPermutation::compute(&air_heights);
392            // ATTENTION: make sure to permute everything in vm_pk that references the original AIR
393            // ID ordering:
394            root_air_perm.permute(&mut vm_pk.per_air);
395            #[cfg(not(feature = "legacy-v1-3-evm-verifier"))]
396            for thc in &mut vm_pk.trace_height_constraints {
397                root_air_perm.permute(&mut thc.coefficients);
398            }
399
400            RootVerifierProvingKey {
401                vm_pk: Arc::new(VmProvingKey {
402                    fri_params: config.root_fri_params,
403                    vm_config: root_vm_config,
404                    vm_pk,
405                }),
406                root_committed_exe,
407                air_heights,
408            }
409        };
410        Ok((
411            Self {
412                leaf_vm_pk,
413                internal_vm_pk,
414                internal_committed_exe,
415                root_verifier_pk,
416            },
417            internal_proof,
418        ))
419    }
420
421    pub fn get_agg_vk(&self) -> AggVerifyingKey {
422        let leaf_fri_params = self.leaf_vm_pk.fri_params;
423        let leaf_vk = self.leaf_vm_pk.vm_pk.get_vk();
424        let internal_fri_params = self.internal_vm_pk.fri_params;
425        let internal_vk = self.internal_vm_pk.vm_pk.get_vk();
426        let internal_verifier_program_commit = self.internal_committed_exe.get_program_commit();
427        AggVerifyingKey {
428            leaf_fri_params,
429            leaf_vk,
430            internal_fri_params,
431            internal_vk,
432            internal_verifier_program_commit,
433        }
434    }
435
436    pub fn num_user_public_values(&self) -> usize {
437        self.root_verifier_pk
438            .vm_pk
439            .vm_config
440            .system
441            .num_public_values
442            - (2 * DIGEST_SIZE)
443    }
444}
445
446/// Proving key for the root verifier.
447/// Properties:
448/// - Traces heights of each AIR is constant. This is required by the static verifier.
449/// - Instead of the AIR order specified by VmConfig. AIRs are ordered by trace heights.
450#[derive(Serialize, Deserialize, Derivative)]
451#[derivative(Clone(bound = "Com<SC>: Clone"))]
452pub struct RootVerifierProvingKey {
453    /// VM Proving key for the root verifier.
454    /// - AIR proving key in `MultiStarkProvingKey` is ordered by trace height.
455    /// - `VmConfig.overridden_executor_heights` is specified and is in the original AIR order.
456    /// - `VmConfig.memory_config.boundary_air_height` is specified.
457    pub vm_pk: Arc<VmProvingKey<RootSC, NativeConfig>>,
458    /// Committed executable for the root VM.
459    pub root_committed_exe: Arc<VmCommittedExe<RootSC>>,
460    /// The constant trace heights, ordered by AIR ID (the original ordering from VmConfig).
461    pub air_heights: Vec<u32>,
462}
463
464#[cfg(feature = "evm-prove")]
465impl RootVerifierProvingKey {
466    pub(crate) fn air_id_permutation(&self) -> AirIdPermutation {
467        AirIdPermutation::compute(&self.air_heights)
468    }
469}
470
471#[cfg(feature = "evm-prove")]
472impl Halo2ProvingKey {
473    /// Attention:
474    /// - This function is very expensive. Usually it requires >64GB memory and takes >10 minutes.
475    ///   /// - Please make sure SRS(KZG parameters) is already downloaded.
476    #[tracing::instrument(level = "info", fields(group = "halo2_keygen"), skip_all)]
477    pub fn keygen(
478        halo2_config: Halo2Config,
479        reader: &impl Halo2ParamsReader,
480        pv_handler: &impl StaticVerifierPvHandler,
481        agg_pk: &AggProvingKey,
482        dummy_internal_proof: Proof<SC>,
483    ) -> Result<Self, VirtualMachineError> {
484        let dummy_root_proof = agg_pk
485            .root_verifier_pk
486            .generate_dummy_root_proof(dummy_internal_proof)?;
487        let verifier = agg_pk.root_verifier_pk.keygen_static_verifier(
488            &reader.read_params(halo2_config.verifier_k),
489            dummy_root_proof,
490            pv_handler,
491        );
492        let dummy_snark = verifier.generate_dummy_snark(reader);
493        let wrapper = if let Some(wrapper_k) = halo2_config.wrapper_k {
494            Halo2WrapperProvingKey::keygen(&reader.read_params(wrapper_k), dummy_snark)
495        } else {
496            Halo2WrapperProvingKey::keygen_auto_tune(reader, dummy_snark)
497        };
498        Ok(Halo2ProvingKey {
499            verifier: Arc::new(verifier),
500            wrapper: Arc::new(wrapper),
501            profiling: halo2_config.profiling,
502        })
503    }
504}
505
506/// For internal use only.
507pub fn _leaf_keygen(
508    fri_params: FriParameters,
509    leaf_vm_config: NativeConfig,
510) -> Result<Arc<VmProvingKey<SC, NativeConfig>>, AirInventoryError> {
511    let leaf_engine = BabyBearPoseidon2Engine::new(fri_params);
512    let leaf_vm_pk = info_span!("keygen", group = "leaf").in_scope(|| {
513        leaf_vm_config
514            .create_airs()
515            .map(|airs| airs.keygen(&leaf_engine))
516    })?;
517    Ok(Arc::new(VmProvingKey {
518        fri_params,
519        vm_config: leaf_vm_config,
520        vm_pk: leaf_vm_pk,
521    }))
522}
openvm_sdk/keygen/mod.rs

openvm_sdk/keygen/
mod.rs
