openvm_cuda_backend/
device.rs

1use getset::{CopyGetters, Getters, MutGetters};
2use openvm_cuda_common::{common::get_device, stream::GpuDeviceCtx};
3use openvm_stark_backend::{
4    memory_metering::ProvingMemoryConfig, StarkProtocolConfig, SystemParams,
5};
6
7use crate::cuda::{
8    batch_ntt_small::{ensure_device_ntt_twiddles_initialized, validate_gpu_l_skip},
9    device_info::get_sm_count,
10};
11
12/// Pure device configuration — no stream, no CUDA runtime state.
13/// Renamed from the old `GpuDevice`.
14#[derive(Clone, Getters, CopyGetters, MutGetters)]
15pub struct GpuDeviceConfig {
16    #[getset(get = "pub")]
17    pub(crate) params: SystemParams,
18    #[getset(get = "pub", get_mut = "pub")]
19    pub(crate) prover_config: GpuProverConfig,
20    pub id: u32,
21    #[getset(get_copy = "pub")]
22    pub sm_count: u32,
23}
24
25#[derive(Clone, Copy)]
26pub struct GpuProverConfig {
27    pub cache_stacked_matrix: bool,
28    pub cache_rs_code_matrix: bool,
29    pub zerocheck_save_memory: bool,
30}
31
32impl GpuProverConfig {
33    pub fn proving_memory_config<SC: StarkProtocolConfig>(
34        &self,
35        config: &SC,
36    ) -> ProvingMemoryConfig {
37        // Interaction memory is estimated from the CUDA fractional-GKR buffer model in
38        // `openvm_stark_backend::memory_metering`. Update that estimate when changing GKR
39        // input layout, work-buffer sizing, or scratch allocations.
40        ProvingMemoryConfig::from_protocol_config(config, self.cache_rs_code_matrix)
41    }
42}
43
44/// Stream-owning device handle. Wraps [`GpuDeviceConfig`] with a
45/// [`GpuDeviceCtx`] that carries the explicit CUDA stream.
46#[derive(Clone)]
47pub struct GpuDevice {
48    pub config: GpuDeviceConfig,
49    pub device_ctx: GpuDeviceCtx,
50}
51
52impl GpuDevice {
53    pub fn new(params: SystemParams) -> Result<Self, openvm_cuda_common::error::CudaError> {
54        validate_gpu_l_skip(params.l_skip)
55            .expect("GPU backend requires l_skip <= 9 for current CUDA kernels");
56        ensure_device_ntt_twiddles_initialized()
57            .expect("failed to initialize small-NTT twiddles for current CUDA device");
58
59        let prover_config = GpuProverConfig {
60            zerocheck_save_memory: params.log_blowup == 1,
61            ..Default::default()
62        };
63        let id = get_device().unwrap() as u32;
64        let device_ctx = GpuDeviceCtx::for_device(id)?;
65        let sm_count = get_sm_count(id).expect("failed to get SM count");
66        let config = GpuDeviceConfig {
67            params,
68            prover_config,
69            id,
70            sm_count,
71        };
72
73        Ok(Self { config, device_ctx })
74    }
75
76    // Delegate accessors to inner config
77    pub fn params(&self) -> &SystemParams {
78        &self.config.params
79    }
80
81    pub fn prover_config(&self) -> &GpuProverConfig {
82        &self.config.prover_config
83    }
84
85    pub fn prover_config_mut(&mut self) -> &mut GpuProverConfig {
86        &mut self.config.prover_config
87    }
88
89    pub fn sm_count(&self) -> u32 {
90        self.config.sm_count
91    }
92
93    pub fn id(&self) -> u32 {
94        self.config.id
95    }
96
97    pub fn with_cache_rs_code_matrix(mut self, cache_rs_code_matrix: bool) -> Self {
98        self.config.prover_config.cache_rs_code_matrix = cache_rs_code_matrix;
99        self
100    }
101
102    pub fn set_cache_rs_code_matrix(&mut self, cache_rs_code_matrix: bool) {
103        self.config.prover_config.cache_rs_code_matrix = cache_rs_code_matrix;
104    }
105}
106
107/// Default GPU prover cache settings.
108impl Default for GpuProverConfig {
109    fn default() -> Self {
110        Self {
111            cache_stacked_matrix: false,
112            cache_rs_code_matrix: false,
113            zerocheck_save_memory: true,
114        }
115    }
116}
openvm_cuda_backend/device.rs

openvm_cuda_backend/
device.rs
