openvm_cuda_backend/
gpu_device.rs

1use derivative::Derivative;
2use openvm_cuda_common::common::get_device;
3use openvm_stark_sdk::config::FriParameters;
4use p3_baby_bear::BabyBear;
5use p3_commit::TwoAdicMultiplicativeCoset;
6use p3_field::FieldAlgebra;
7use p3_util::log2_strict_usize;
8
9use crate::{cuda::kernels::ntt::sppark_init, fri_log_up::FriLogUpPhaseGpu};
10
11#[derive(Derivative, derive_new::new, Clone, Copy, Debug)]
12pub struct GpuConfig {
13    pub fri: FriParameters,
14    pub shift: BabyBear,
15}
16
17#[derive(Derivative, Clone, Debug)]
18pub struct GpuDevice {
19    pub config: GpuConfig,
20    pub id: u32,
21    rap_phase_seq: Option<FriLogUpPhaseGpu>,
22}
23
24impl GpuDevice {
25    pub fn new(config: GpuConfig, rap_phase_seq: Option<FriLogUpPhaseGpu>) -> Self {
26        let device_id = get_device().unwrap();
27        unsafe {
28            sppark_init().unwrap();
29        }
30        Self {
31            config,
32            id: device_id as u32,
33            rap_phase_seq,
34        }
35    }
36
37    pub fn rap_phase_seq(&self) -> &FriLogUpPhaseGpu {
38        self.rap_phase_seq
39            .as_ref()
40            .expect("FriLogUpPhaseGpu is not initialized")
41    }
42
43    pub fn natural_domain_for_degree(&self, degree: usize) -> TwoAdicMultiplicativeCoset<BabyBear> {
44        let log_n = log2_strict_usize(degree);
45        TwoAdicMultiplicativeCoset {
46            log_n,
47            shift: BabyBear::ONE,
48        }
49    }
50}