openvm_cuda_backend/
gpu_device.rs1use 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}