openvm_cuda_backend/
engine.rs

1use getset::MutGetters;
2use openvm_stark_backend::{
3    memory_metering::ProvingMemoryConfig, prover::Coordinator, StarkEngine, SystemParams,
4};
5#[cfg(feature = "baby-bear-bn254-poseidon2")]
6use openvm_stark_sdk::config::baby_bear_bn254_poseidon2::BabyBearBn254Poseidon2Config;
7
8#[cfg(feature = "baby-bear-bn254-poseidon2")]
9use crate::{
10    bn254_sponge::MultiFieldTranscriptGpu, gpu_backend::GenericGpuBackend,
11    hash_scheme::BabyBearBn254Poseidon2HashScheme,
12};
13use crate::{
14    hash_scheme::{DefaultHashScheme, GpuHashScheme},
15    prelude::SC,
16    sponge::DuplexSpongeGpu,
17    GpuBackend, GpuDevice,
18};
19
20/// Generic GPU proving engine parameterised by a hash scheme.
21///
22/// Use the [`BabyBearPoseidon2GpuEngine`] type alias for the default Poseidon2 engine.
23#[derive(MutGetters)]
24pub struct GpuEngine<HS: GpuHashScheme> {
25    #[getset(get_mut = "pub")]
26    device: GpuDevice,
27    #[getset(get_mut = "pub")]
28    config: HS::SC,
29}
30
31/// Concrete GPU engine using the default BabyBear-Poseidon2 hash scheme.
32pub type BabyBearPoseidon2GpuEngine = GpuEngine<DefaultHashScheme>;
33
34/// GPU engine using the BabyBear + BN254 Poseidon2 hash scheme (Groth16-friendly).
35#[cfg(feature = "baby-bear-bn254-poseidon2")]
36pub type BabyBearBn254Poseidon2GpuEngine = GpuEngine<BabyBearBn254Poseidon2HashScheme>;
37
38impl<HS: GpuHashScheme> GpuEngine<HS> {
39    pub fn new(params: SystemParams) -> Self {
40        let config = HS::default_config(params.clone());
41        Self {
42            device: GpuDevice::new(params).expect("failed to create GpuDevice"),
43            config,
44        }
45    }
46}
47
48impl StarkEngine for GpuEngine<DefaultHashScheme> {
49    type SC = SC;
50    type PB = GpuBackend;
51    type PD = GpuDevice;
52    type TS = DuplexSpongeGpu;
53
54    fn new(params: SystemParams) -> Self {
55        GpuEngine::new(params)
56    }
57
58    fn config(&self) -> &Self::SC {
59        &self.config
60    }
61
62    fn device(&self) -> &Self::PD {
63        &self.device
64    }
65
66    fn proving_memory_config(&self) -> ProvingMemoryConfig {
67        self.device
68            .prover_config()
69            .proving_memory_config(self.config())
70    }
71
72    fn initial_transcript(&self) -> Self::TS {
73        DuplexSpongeGpu::default()
74    }
75
76    fn prover_from_transcript(
77        &self,
78        transcript: DuplexSpongeGpu,
79    ) -> Coordinator<SC, Self::PB, Self::PD, Self::TS> {
80        Coordinator::new(GpuBackend::default(), self.device.clone(), transcript)
81    }
82}
83
84#[cfg(feature = "baby-bear-bn254-poseidon2")]
85impl StarkEngine for GpuEngine<BabyBearBn254Poseidon2HashScheme> {
86    type SC = BabyBearBn254Poseidon2Config;
87    type PB = GenericGpuBackend<BabyBearBn254Poseidon2HashScheme>;
88    type PD = GpuDevice;
89    type TS = MultiFieldTranscriptGpu;
90
91    fn new(params: SystemParams) -> Self {
92        GpuEngine::new(params)
93    }
94
95    fn config(&self) -> &Self::SC {
96        &self.config
97    }
98
99    fn device(&self) -> &Self::PD {
100        &self.device
101    }
102
103    fn proving_memory_config(&self) -> ProvingMemoryConfig {
104        self.device
105            .prover_config()
106            .proving_memory_config(self.config())
107    }
108
109    fn initial_transcript(&self) -> Self::TS {
110        MultiFieldTranscriptGpu::default()
111    }
112
113    fn prover_from_transcript(
114        &self,
115        transcript: MultiFieldTranscriptGpu,
116    ) -> Coordinator<BabyBearBn254Poseidon2Config, Self::PB, Self::PD, Self::TS> {
117        Coordinator::new(
118            GenericGpuBackend::default(),
119            self.device.clone(),
120            transcript,
121        )
122    }
123}
openvm_cuda_backend/engine.rs

openvm_cuda_backend/
engine.rs
