openvm_sdk/keygen/
mod.rs

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

openvm_sdk/keygen/
mod.rs
