openvm_continuations/circuit/inner/
trace.rs

1use itertools::Itertools;
2#[cfg(feature = "cuda")]
3use openvm_circuit_primitives::hybrid_chip::cpu_proving_ctx_to_gpu;
4use openvm_cpu_backend::CpuBackend;
5#[cfg(feature = "cuda")]
6use openvm_cuda_backend::GpuBackend;
7#[cfg(feature = "cuda")]
8use openvm_cuda_common::stream::GpuDeviceCtx;
9use openvm_stark_backend::{
10    proof::Proof,
11    prover::{AirProvingContext, ProverBackend},
12};
13use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, F};
14use openvm_verify_stark_host::pvs::{DeferralPvs, VkCommit};
15
16use crate::circuit::{SingleAirTraceData, SubCircuitTraceData};
17
18#[derive(Copy, Clone)]
19pub enum ProofsType {
20    Vm,
21    Deferral,
22    Mix,
23    Combined,
24}
25
26// Trait that inner provers use to remain generic in PB
27pub trait InnerTraceGen<PB: ProverBackend, DC: Clone + Send + Sync> {
28    fn new(deferral_enabled: bool) -> Self;
29    fn generate_pre_verifier_subcircuit_ctxs(
30        &self,
31        proofs: &[Proof<BabyBearPoseidon2Config>],
32        proofs_type: ProofsType,
33        absent_trace_pvs: Option<(DeferralPvs<F>, bool)>,
34        child_is_app: bool,
35        child_vk_commit: VkCommit<F>,
36        device_ctx: &DC,
37    ) -> SubCircuitTraceData<PB>;
38    fn generate_post_verifier_subcircuit_ctxs(
39        &self,
40        proofs: &[Proof<BabyBearPoseidon2Config>],
41        proofs_type: ProofsType,
42        child_is_app: bool,
43        device_ctx: &DC,
44    ) -> Vec<AirProvingContext<PB>>;
45}
46
47pub struct InnerTraceGenImpl {
48    pub deferral_enabled: bool,
49}
50
51impl InnerTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()> for InnerTraceGenImpl {
52    fn new(deferral_enabled: bool) -> Self {
53        Self { deferral_enabled }
54    }
55
56    fn generate_pre_verifier_subcircuit_ctxs(
57        &self,
58        proofs: &[Proof<BabyBearPoseidon2Config>],
59        proofs_type: ProofsType,
60        absent_trace_pvs: Option<(DeferralPvs<F>, bool)>,
61        child_is_app: bool,
62        child_vk_commit: VkCommit<F>,
63        _device_ctx: &(),
64    ) -> SubCircuitTraceData<CpuBackend<BabyBearPoseidon2Config>> {
65        let SingleAirTraceData {
66            air_proving_ctx: verifier_pvs_ctx,
67            mut poseidon2_compress_inputs,
68            poseidon2_permute_inputs,
69            mut range_check_inputs,
70        } = super::verifier::generate_proving_ctx(
71            proofs,
72            proofs_type,
73            child_is_app,
74            child_vk_commit,
75            self.deferral_enabled,
76        );
77        let vm_pvs_ctx = super::vm_pvs::generate_proving_ctx(
78            proofs,
79            proofs_type,
80            child_is_app,
81            self.deferral_enabled,
82        );
83
84        let idx2_ctx = if self.deferral_enabled {
85            let (def_pvs_ctx, def_poseidon2_inputs, def_range_check_inputs) =
86                super::def_pvs::generate_proving_ctx(
87                    proofs,
88                    proofs_type,
89                    child_is_app,
90                    absent_trace_pvs,
91                );
92            poseidon2_compress_inputs.extend_from_slice(&def_poseidon2_inputs);
93            range_check_inputs.extend(def_range_check_inputs);
94            def_pvs_ctx
95        } else {
96            super::unset::generate_proving_ctx(&[], child_is_app)
97        };
98
99        SubCircuitTraceData {
100            air_proving_ctxs: vec![verifier_pvs_ctx, vm_pvs_ctx, idx2_ctx],
101            poseidon2_compress_inputs,
102            poseidon2_permute_inputs,
103            range_check_inputs,
104        }
105    }
106
107    fn generate_post_verifier_subcircuit_ctxs(
108        &self,
109        proofs: &[Proof<BabyBearPoseidon2Config>],
110        proofs_type: ProofsType,
111        child_is_app: bool,
112        _device_ctx: &(),
113    ) -> Vec<AirProvingContext<CpuBackend<BabyBearPoseidon2Config>>> {
114        if !self.deferral_enabled {
115            return vec![];
116        }
117
118        let (vm_unset, def_unset) = match proofs_type {
119            ProofsType::Vm => (
120                vec![],
121                proofs.iter().enumerate().map(|(i, _)| i).collect_vec(),
122            ),
123            ProofsType::Deferral => (
124                proofs.iter().enumerate().map(|(i, _)| i).collect_vec(),
125                vec![],
126            ),
127            ProofsType::Mix => (vec![1], vec![0]),
128            ProofsType::Combined => (vec![], vec![]),
129        };
130        vec![
131            super::unset::generate_proving_ctx(&vm_unset, child_is_app),
132            super::unset::generate_proving_ctx(&def_unset, child_is_app),
133        ]
134    }
135}
136
137#[cfg(feature = "cuda")]
138impl InnerTraceGen<GpuBackend, GpuDeviceCtx> for InnerTraceGenImpl {
139    fn new(deferral_enabled: bool) -> Self {
140        Self { deferral_enabled }
141    }
142
143    fn generate_pre_verifier_subcircuit_ctxs(
144        &self,
145        proofs: &[Proof<BabyBearPoseidon2Config>],
146        proofs_type: ProofsType,
147        absent_trace_pvs: Option<(DeferralPvs<F>, bool)>,
148        child_is_app: bool,
149        child_vk_commit: VkCommit<F>,
150        device_ctx: &GpuDeviceCtx,
151    ) -> SubCircuitTraceData<GpuBackend> {
152        let data =
153            <Self as InnerTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()>>::generate_pre_verifier_subcircuit_ctxs(
154                self,
155                proofs,
156                proofs_type,
157                absent_trace_pvs,
158                child_is_app,
159                child_vk_commit,
160                &(),
161            );
162        SubCircuitTraceData {
163            air_proving_ctxs: data
164                .air_proving_ctxs
165                .into_iter()
166                .map(|air_ctx| cpu_proving_ctx_to_gpu(air_ctx, device_ctx))
167                .collect_vec(),
168            poseidon2_compress_inputs: data.poseidon2_compress_inputs,
169            poseidon2_permute_inputs: data.poseidon2_permute_inputs,
170            range_check_inputs: data.range_check_inputs,
171        }
172    }
173
174    fn generate_post_verifier_subcircuit_ctxs(
175        &self,
176        proofs: &[Proof<BabyBearPoseidon2Config>],
177        proofs_type: ProofsType,
178        child_is_app: bool,
179        device_ctx: &GpuDeviceCtx,
180    ) -> Vec<AirProvingContext<GpuBackend>> {
181        let cpu_ctxs =
182            <Self as InnerTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()>>::generate_post_verifier_subcircuit_ctxs(
183                self,
184                proofs,
185                proofs_type,
186                child_is_app,
187                &(),
188            );
189        cpu_ctxs
190            .into_iter()
191            .map(|air_ctx| cpu_proving_ctx_to_gpu(air_ctx, device_ctx))
192            .collect_vec()
193    }
194}
openvm_continuations/circuit/inner/trace.rs

openvm_continuations/circuit/inner/
trace.rs
