openvm_recursion_circuit/cuda/
preflight.rs

1use std::iter::once;
2
3use itertools::Itertools;
4use openvm_cuda_backend::prelude::EF;
5use openvm_cuda_common::{d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
6use openvm_stark_backend::{keygen::types::MultiStarkVerifyingKey, proof::Proof};
7use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, Digest};
8
9use crate::{
10    cuda::{
11        to_device_or_nullptr_on,
12        types::{TraceHeight, TraceMetadata},
13    },
14    proof_shape::proof_shape::compute_air_shape_lookup_counts,
15    system::Preflight,
16};
17
18/*
19 * Tracegen information (i.e. records) on a GPU device. Each field should
20 * be computable as soon as the verifier circuit runs preflight.
21 */
22#[derive(Debug)]
23pub struct PreflightGpu {
24    pub cpu: Preflight,
25    pub transcript: TranscriptLog,
26    pub proof_shape: ProofShapePreflightGpu,
27    pub gkr: GkrPreflightGpu,
28    pub batch_constraint: BatchConstraintPreflightGpu,
29    pub stacking: StackingPreflightGpu,
30    pub whir: WhirPreflightGpu,
31}
32
33#[derive(Debug)]
34pub struct TranscriptLog {
35    _dummy: usize,
36}
37
38#[derive(Debug)]
39pub struct ProofShapePreflightGpu {
40    pub sorted_trace_heights: DeviceBuffer<TraceHeight>,
41    pub sorted_trace_metadata: DeviceBuffer<TraceMetadata>,
42    pub sorted_cached_commits: DeviceBuffer<Digest>,
43
44    pub per_row_tidx: DeviceBuffer<usize>,
45    pub pvs_tidx: DeviceBuffer<usize>,
46    pub post_tidx: usize,
47
48    pub num_present: usize,
49    pub n_max: usize,
50    pub n_logup: usize,
51    pub final_cidx: usize,
52    pub final_total_interactions: usize,
53    pub main_commit: Digest,
54}
55
56#[derive(Debug)]
57pub struct GkrPreflightGpu {
58    _dummy: usize,
59}
60
61#[derive(Debug)]
62pub struct BatchConstraintPreflightGpu {
63    pub sumcheck_rnd: DeviceBuffer<EF>,
64}
65
66#[derive(Debug)]
67pub struct StackingPreflightGpu {
68    pub sumcheck_rnd: DeviceBuffer<EF>,
69}
70
71#[derive(Debug)]
72pub struct WhirPreflightGpu {
73    _dummy: usize,
74}
75
76impl PreflightGpu {
77    pub fn new(
78        vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
79        proof: &Proof<BabyBearPoseidon2Config>,
80        preflight: &Preflight,
81        device_ctx: &GpuDeviceCtx,
82    ) -> Self {
83        PreflightGpu {
84            cpu: preflight.clone(),
85            transcript: Self::transcript(preflight),
86            proof_shape: Self::proof_shape(vk, proof, preflight, device_ctx),
87            gkr: Self::gkr(preflight),
88            batch_constraint: Self::batch_constraint(preflight, device_ctx),
89            stacking: Self::stacking(preflight, device_ctx),
90            whir: Self::whir(preflight),
91        }
92    }
93
94    fn transcript(_preflight: &Preflight) -> TranscriptLog {
95        TranscriptLog { _dummy: 0 }
96    }
97
98    fn proof_shape(
99        vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
100        proof: &Proof<BabyBearPoseidon2Config>,
101        preflight: &Preflight,
102        device_ctx: &GpuDeviceCtx,
103    ) -> ProofShapePreflightGpu {
104        let mut sorted_cached_commits: Vec<Digest> = vec![];
105        let mut cidx = 1;
106        let mut total_interactions = 0;
107        let l_skip = vk.inner.params.l_skip;
108
109        let bc_air_shape_lookups = compute_air_shape_lookup_counts(vk);
110
111        let (sorted_trace_heights, sorted_trace_metadata): (Vec<_>, Vec<_>) = preflight
112            .proof_shape
113            .sorted_trace_vdata
114            .iter()
115            .map(|(air_idx, vdata)| {
116                let height = TraceHeight {
117                    air_idx: *air_idx,
118                    log_height: vdata.log_height.try_into().unwrap(),
119                };
120                let metadata = TraceMetadata {
121                    cached_idx: sorted_cached_commits.len(),
122                    starting_cidx: cidx,
123                    total_interactions,
124                    num_air_id_lookups: bc_air_shape_lookups[*air_idx],
125                };
126                cidx += vdata.cached_commitments.len()
127                    + vk.inner.per_air[*air_idx].preprocessed_data.is_some() as usize;
128                total_interactions += (1 << vdata.log_height.max(l_skip))
129                    * vk.inner.per_air[*air_idx].num_interactions();
130                sorted_cached_commits.extend_from_slice(&vdata.cached_commitments);
131                (height, metadata)
132            })
133            .chain(
134                proof
135                    .trace_vdata
136                    .iter()
137                    .enumerate()
138                    .filter_map(|(air_idx, vdata)| {
139                        if vdata.is_none() {
140                            Some((
141                                TraceHeight {
142                                    air_idx,
143                                    ..Default::default()
144                                },
145                                TraceMetadata {
146                                    num_air_id_lookups: bc_air_shape_lookups[air_idx],
147                                    ..Default::default()
148                                },
149                            ))
150                        } else {
151                            None
152                        }
153                    }),
154            )
155            .unzip();
156        let per_row_tidx = preflight
157            .proof_shape
158            .starting_tidx
159            .iter()
160            .copied()
161            .chain(once(preflight.proof_shape.post_tidx))
162            .collect_vec();
163
164        let mut pvs_tidx_by_air_id = vec![0usize; vk.inner.per_air.len()];
165        for ((air_idx, pvs), &starting_tidx) in proof
166            .public_values
167            .iter()
168            .enumerate()
169            .filter(|(_, per_air)| !per_air.is_empty())
170            .zip(&preflight.proof_shape.pvs_tidx)
171        {
172            debug_assert!(!pvs.is_empty());
173            pvs_tidx_by_air_id[air_idx] = starting_tidx;
174        }
175
176        ProofShapePreflightGpu {
177            sorted_trace_heights: to_device_or_nullptr_on(&sorted_trace_heights, device_ctx)
178                .unwrap(),
179            sorted_trace_metadata: to_device_or_nullptr_on(&sorted_trace_metadata, device_ctx)
180                .unwrap(),
181            sorted_cached_commits: to_device_or_nullptr_on(&sorted_cached_commits, device_ctx)
182                .unwrap(),
183            per_row_tidx: to_device_or_nullptr_on(&per_row_tidx, device_ctx).unwrap(),
184            pvs_tidx: to_device_or_nullptr_on(&pvs_tidx_by_air_id, device_ctx).unwrap(),
185            post_tidx: preflight.proof_shape.post_tidx,
186            num_present: preflight.proof_shape.sorted_trace_vdata.len(),
187            n_max: preflight.proof_shape.n_max,
188            n_logup: preflight.proof_shape.n_logup,
189            final_cidx: cidx,
190            final_total_interactions: total_interactions,
191            main_commit: proof.common_main_commit,
192        }
193    }
194
195    fn gkr(_preflight: &Preflight) -> GkrPreflightGpu {
196        GkrPreflightGpu { _dummy: 0 }
197    }
198
199    fn batch_constraint(
200        preflight: &Preflight,
201        device_ctx: &GpuDeviceCtx,
202    ) -> BatchConstraintPreflightGpu {
203        BatchConstraintPreflightGpu {
204            sumcheck_rnd: to_device_or_nullptr_on(
205                &preflight.batch_constraint.sumcheck_rnd,
206                device_ctx,
207            )
208            .unwrap(),
209        }
210    }
211
212    fn stacking(preflight: &Preflight, device_ctx: &GpuDeviceCtx) -> StackingPreflightGpu {
213        StackingPreflightGpu {
214            sumcheck_rnd: to_device_or_nullptr_on(&preflight.stacking.sumcheck_rnd, device_ctx)
215                .unwrap(),
216        }
217    }
218
219    fn whir(_preflight: &Preflight) -> WhirPreflightGpu {
220        WhirPreflightGpu { _dummy: 0 }
221    }
222}
openvm_recursion_circuit/cuda/preflight.rs

openvm_recursion_circuit/cuda/
preflight.rs
