openvm_recursion_circuit/
tracegen.rs

1use openvm_cpu_backend::CpuBackend;
2use openvm_stark_backend::{
3    keygen::types::MultiStarkVerifyingKey,
4    proof::Proof,
5    prover::{AirProvingContext, ProverBackend},
6    StarkProtocolConfig,
7};
8use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, F};
9use p3_matrix::dense::RowMajorMatrix;
10
11use crate::system::Preflight;
12
13/// Backend-generic trait to generate a proving context
14pub(crate) trait ModuleChip<PB: ProverBackend> {
15    /// Context needed for trace generation (e.g., VK, proofs, preflights).
16    type Ctx<'a>;
17
18    /// Generate an AirProvingContext. If required_height is Some(..), then the
19    /// resulting trace matrices must have height required_height. This function
20    /// should return None iff required_height is defined AND the matrix requires
21    /// more than required_height rows.
22    fn generate_proving_ctx(
23        &self,
24        ctx: &Self::Ctx<'_>,
25        required_height: Option<usize>,
26    ) -> Option<AirProvingContext<PB>>;
27}
28
29/// Trait to generate a CPU row-major common trace
30pub(crate) trait RowMajorChip<F> {
31    /// Context needed for trace generation (e.g., VK, proofs, preflights).
32    type Ctx<'a>;
33
34    /// Generate row major trace with the same semantics as TraceGenerator
35    fn generate_trace(
36        &self,
37        ctx: &Self::Ctx<'_>,
38        required_height: Option<usize>,
39    ) -> Option<RowMajorMatrix<F>>;
40}
41
42pub(crate) struct StandardTracegenCtx<'a> {
43    pub vk: &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
44    pub proofs: &'a [&'a Proof<BabyBearPoseidon2Config>],
45    pub preflights: &'a [&'a Preflight],
46}
47
48impl<SC: StarkProtocolConfig<F = F>, T: RowMajorChip<F>> ModuleChip<CpuBackend<SC>> for T {
49    type Ctx<'a> = T::Ctx<'a>;
50
51    #[tracing::instrument(level = "trace", skip_all)]
52    fn generate_proving_ctx(
53        &self,
54        ctx: &Self::Ctx<'_>,
55        required_height: Option<usize>,
56    ) -> Option<AirProvingContext<CpuBackend<SC>>> {
57        let common_main_rm = self.generate_trace(ctx, required_height);
58        common_main_rm.map(AirProvingContext::simple_no_pis)
59    }
60}
61
62#[cfg(feature = "cuda")]
63pub(crate) mod cuda {
64    use openvm_cuda_common::stream::GpuDeviceCtx;
65
66    use crate::cuda::{preflight::PreflightGpu, proof::ProofGpu, vk::VerifyingKeyGpu};
67
68    pub(crate) struct StandardTracegenGpuCtx<'a> {
69        pub vk: &'a VerifyingKeyGpu,
70        pub proofs: &'a [ProofGpu],
71        pub preflights: &'a [PreflightGpu],
72        pub device_ctx: &'a GpuDeviceCtx,
73    }
74}