openvm_sdk/keygen/
mod.rs

1use std::sync::Arc;
2
3use derivative::Derivative;
4// use dummy::{compute_root_proof_heights, dummy_internal_proof_riscv_app_vm};
5use openvm_circuit::{
6    arch::{AirInventoryError, SystemConfig, VirtualMachine, VirtualMachineError, VmCircuitConfig},
7    system::memory::dimensions::MemoryDimensions,
8};
9use openvm_continuations::verifier::{
10    internal::InternalVmVerifierConfig, leaf::LeafVmVerifierConfig, root::RootVmVerifierConfig,
11};
12use openvm_native_circuit::{NativeConfig, NativeCpuBuilder};
13use openvm_native_compiler::ir::DIGEST_SIZE;
14use openvm_stark_backend::{
15    config::Val,
16    engine::StarkEngine,
17    p3_field::{FieldExtensionAlgebra, PrimeField32, TwoAdicField},
18};
19use openvm_stark_sdk::{
20    config::{
21        baby_bear_poseidon2::BabyBearPoseidon2Engine,
22        baby_bear_poseidon2_root::BabyBearPoseidon2RootEngine, FriParameters,
23    },
24    engine::StarkFriEngine,
25    openvm_stark_backend::{
26        config::{Com, StarkGenericConfig},
27        keygen::types::MultiStarkVerifyingKey,
28        proof::Proof,
29    },
30};
31use serde::{Deserialize, Serialize};
32use tracing::{info_span, instrument};
33#[cfg(feature = "evm-prove")]
34use {
35    openvm_continuations::static_verifier::StaticVerifierPvHandler,
36    openvm_native_recursion::halo2::{
37        utils::Halo2ParamsReader, verifier::Halo2VerifierProvingKey,
38        wrapper::Halo2WrapperProvingKey,
39    },
40};
41
42#[cfg(feature = "evm-prove")]
43use crate::config::Halo2Config;
44use crate::{
45    commit::VmCommittedExe,
46    config::{AggregationConfig, AppConfig},
47    keygen::{
48        dummy::{compute_root_proof_heights, dummy_internal_proof_riscv_app_vm},
49        perm::AirIdPermutation,
50    },
51    prover::vm::types::VmProvingKey,
52    util::check_max_constraint_degrees,
53    RootSC, SC,
54};
55
56pub mod asm;
57pub(crate) mod dummy;
58pub mod perm;
59#[cfg(feature = "evm-prove")]
60pub mod static_verifier;
61
62/// This is lightweight to clone as it contains smart pointers to the proving keys.
63#[derive(Clone, Serialize, Deserialize)]
64pub struct AppProvingKey<VC> {
65    /// The committed executable of the leaf verifier program that verifies proofs of the App VM
66    /// circuit. The App VM circuit constraints are statically compiled into this executable.
67    pub leaf_committed_exe: Arc<VmCommittedExe<SC>>,
68    pub leaf_fri_params: FriParameters,
69    pub app_vm_pk: Arc<VmProvingKey<SC, VC>>,
70}
71
72#[derive(Clone, Serialize, Deserialize)]
73pub struct AppVerifyingKey {
74    /// We store the FRI parameters used to generate the proof separately.
75    pub fri_params: FriParameters,
76    /// STARK backend verifying key
77    pub vk: MultiStarkVerifyingKey<SC>,
78    pub memory_dimensions: MemoryDimensions,
79}
80
81/// The STARK proving keys necessary for aggregation of app proofs into a single aggregate STARK
82/// proof.
83///
84/// This is lightweight to clone as it contains smart pointers to the proving keys.
85#[derive(Clone, Serialize, Deserialize)]
86pub struct AggProvingKey {
87    pub leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
88    pub internal_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
89    pub internal_committed_exe: Arc<VmCommittedExe<SC>>,
90    pub root_verifier_pk: RootVerifierProvingKey,
91}
92
93#[derive(Clone, Serialize, Deserialize)]
94pub struct AggVerifyingKey {
95    pub(super) leaf_fri_params: FriParameters,
96    pub(super) leaf_vk: MultiStarkVerifyingKey<SC>,
97    /// FRI parameters used to generate the last internal proof.
98    pub(super) internal_fri_params: FriParameters,
99    pub(super) internal_vk: MultiStarkVerifyingKey<SC>,
100    pub(super) internal_verifier_program_commit: Com<SC>,
101}
102
103/// Attention: the serialized size of this struct is VERY large, usually >10GB.
104///
105/// This is lightweight to clone as it contains smart pointers to the proving keys.
106#[cfg(feature = "evm-prove")]
107#[derive(Clone, Serialize, Deserialize)]
108pub struct Halo2ProvingKey {
109    /// Static verifier to verify a stark proof of the root verifier.
110    pub verifier: Arc<Halo2VerifierProvingKey>,
111    /// Wrapper circuit to verify static verifier and reduce the verification costs in the final
112    /// proof.
113    pub wrapper: Arc<Halo2WrapperProvingKey>,
114    /// Whether to collect detailed profiling metrics
115    pub profiling: bool,
116}
117
118impl<VC> AppProvingKey<VC>
119where
120    VC: Clone + VmCircuitConfig<SC> + AsRef<SystemConfig>,
121{
122    pub fn keygen(config: AppConfig<VC>) -> Result<Self, AirInventoryError> {
123        let app_engine = BabyBearPoseidon2Engine::new(config.app_fri_params.fri_params);
124        let app_vm_pk = {
125            let vm_pk = config.app_vm_config.create_airs()?.keygen(&app_engine);
126            assert!(
127                vm_pk.max_constraint_degree
128                    <= config.app_fri_params.fri_params.max_constraint_degree()
129            );
130            check_max_constraint_degrees(
131                config.app_vm_config.as_ref(),
132                &config.app_fri_params.fri_params,
133            );
134            VmProvingKey {
135                fri_params: config.app_fri_params.fri_params,
136                vm_config: config.app_vm_config.clone(),
137                vm_pk,
138            }
139        };
140        check_recursive_verifier_size(
141            &app_vm_pk.vm_pk.get_vk(),
142            config.app_fri_params.fri_params,
143            config.leaf_fri_params.fri_params.log_blowup,
144        );
145        let leaf_committed_exe = {
146            let leaf_engine = BabyBearPoseidon2Engine::new(config.leaf_fri_params.fri_params);
147            let leaf_program = LeafVmVerifierConfig {
148                app_fri_params: config.app_fri_params.fri_params,
149                app_system_config: config.app_vm_config.as_ref().clone(),
150                compiler_options: config.compiler_options,
151            }
152            .build_program(&app_vm_pk.vm_pk.get_vk());
153            Arc::new(VmCommittedExe::commit(
154                leaf_program.into(),
155                leaf_engine.config().pcs(),
156            ))
157        };
158        Ok(Self {
159            leaf_committed_exe,
160            leaf_fri_params: config.leaf_fri_params.fri_params,
161            app_vm_pk: Arc::new(app_vm_pk),
162        })
163    }
164
165    pub fn num_public_values(&self) -> usize {
166        self.app_vm_pk.vm_config.as_ref().num_public_values
167    }
168
169    pub fn get_app_vk(&self) -> AppVerifyingKey {
170        AppVerifyingKey {
171            fri_params: self.app_vm_pk.fri_params,
172            vk: self.app_vm_pk.vm_pk.get_vk(),
173            memory_dimensions: self
174                .app_vm_pk
175                .vm_config
176                .as_ref()
177                .memory_config
178                .memory_dimensions(),
179        }
180    }
181
182    pub fn leaf_verifier_program_commit(&self) -> Com<SC> {
183        self.leaf_committed_exe.get_program_commit()
184    }
185
186    pub fn app_fri_params(&self) -> FriParameters {
187        self.app_vm_pk.fri_params
188    }
189
190    pub fn vm_config(&self) -> &VC {
191        &self.app_vm_pk.vm_config
192    }
193
194    pub fn app_config(&self) -> AppConfig<VC> {
195        AppConfig {
196            app_fri_params: self.app_fri_params().into(),
197            app_vm_config: self.vm_config().clone(),
198            leaf_fri_params: self.leaf_fri_params.into(),
199            compiler_options: Default::default(),
200        }
201    }
202}
203
204/// Try to determine statically if there will be an issue with the recursive verifier size and log
205/// a warning if so.
206///
207/// `next_log_blowup` refers to the `log_blowup` of the next verifier in the chain; this determines
208/// a maximum trace height.
209fn check_recursive_verifier_size<SC: StarkGenericConfig>(
210    vk: &MultiStarkVerifyingKey<SC>,
211    fri_params: FriParameters,
212    next_log_blowup: usize,
213) where
214    Val<SC>: PrimeField32 + TwoAdicField,
215{
216    let vk = &vk.inner;
217
218    // for each round we will compute the pair (total_width, num_airs, num_pts)
219    let mut rounds = vec![];
220
221    // Preprocessed rounds.
222    rounds.extend(
223        vk.per_air
224            .iter()
225            .filter_map(|vk| vk.params.width.preprocessed)
226            .map(|width| (width, 1, 2)),
227    );
228
229    let common_main_total_width = vk
230        .per_air
231        .iter()
232        .map(|vk| vk.params.width.common_main)
233        .sum();
234    rounds.push((common_main_total_width, vk.per_air.len(), 2));
235
236    for vk in vk.per_air.iter() {
237        for &cached_main_width in &vk.params.width.cached_mains {
238            rounds.push((cached_main_width, 1, 2));
239        }
240    }
241
242    let mut after_challenge_rounds = vec![];
243    for vk in vk.per_air.iter() {
244        let widths = &vk.params.width.after_challenge;
245        if widths.len() > after_challenge_rounds.len() {
246            after_challenge_rounds.resize(widths.len(), (0, 0, 2));
247        }
248        for (i, &width) in widths.iter().enumerate() {
249            after_challenge_rounds[i].0 += SC::Challenge::D * width;
250            after_challenge_rounds[i].1 += 1;
251        }
252    }
253    rounds.extend(after_challenge_rounds);
254
255    let quotient_round = (
256        vk.per_air
257            .iter()
258            .map(|vk| SC::Challenge::D * vk.quotient_degree as usize)
259            .sum(),
260        vk.per_air.len(),
261        1,
262    );
263    rounds.push(quotient_round);
264
265    // This computes the number of rows in the `FRI_REDUCED_OPENING` chip, which is the expected
266    // bottleneck of the recursive verifier.
267    let fri_reduced_opening_trace_height = fri_params.num_queries
268        * rounds
269            .iter()
270            .map(|(total_width, num_airs, total_pts)| total_pts * (total_width + 2 * num_airs))
271            .sum::<usize>();
272    // First check: is FriReducedOpening trace height too large?
273    if fri_reduced_opening_trace_height > (1 << (Val::<SC>::TWO_ADICITY - next_log_blowup)) {
274        tracing::warn!("recursive verifier size may be too large; FriReducedOpening height ({fri_reduced_opening_trace_height}) > {}", 1 << (Val::<SC>::TWO_ADICITY - next_log_blowup));
275    }
276    // Second check: static check for log up soundness constraints using FriReducedOpening trace
277    // height as proxy
278    if fri_reduced_opening_trace_height as u32 >= Val::<SC>::ORDER_U32 / 200 {
279        tracing::warn!(
280            "recursive verifier size may violate log up soundness constraints; {} > {}",
281            200 * fri_reduced_opening_trace_height,
282            Val::<SC>::ORDER_U32
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
