openvm-cuda-backend 2.0.0

OpenVM CUDA prover backend for the SWIRL proof system
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use getset::{CopyGetters, Getters, MutGetters};
use openvm_cuda_common::{common::get_device, stream::GpuDeviceCtx};
use openvm_stark_backend::{
    memory_metering::ProvingMemoryConfig, StarkProtocolConfig, SystemParams,
};

use crate::cuda::{
    batch_ntt_small::{ensure_device_ntt_twiddles_initialized, validate_gpu_l_skip},
    device_info::get_sm_count,
};

/// Pure device configuration — no stream, no CUDA runtime state.
/// Renamed from the old `GpuDevice`.
#[derive(Clone, Getters, CopyGetters, MutGetters)]
pub struct GpuDeviceConfig {
    #[getset(get = "pub")]
    pub(crate) params: SystemParams,
    #[getset(get = "pub", get_mut = "pub")]
    pub(crate) prover_config: GpuProverConfig,
    pub id: u32,
    #[getset(get_copy = "pub")]
    pub sm_count: u32,
}

#[derive(Clone, Copy)]
pub struct GpuProverConfig {
    pub cache_stacked_matrix: bool,
    pub cache_rs_code_matrix: bool,
    pub zerocheck_save_memory: bool,
}

impl GpuProverConfig {
    pub fn proving_memory_config<SC: StarkProtocolConfig>(
        &self,
        config: &SC,
    ) -> ProvingMemoryConfig {
        // Interaction memory is estimated from the CUDA fractional-GKR buffer model in
        // `openvm_stark_backend::memory_metering`. Update that estimate when changing GKR
        // input layout, work-buffer sizing, or scratch allocations.
        ProvingMemoryConfig::from_protocol_config(config, self.cache_rs_code_matrix)
    }
}

/// Stream-owning device handle. Wraps [`GpuDeviceConfig`] with a
/// [`GpuDeviceCtx`] that carries the explicit CUDA stream.
#[derive(Clone)]
pub struct GpuDevice {
    pub config: GpuDeviceConfig,
    pub device_ctx: GpuDeviceCtx,
}

impl GpuDevice {
    pub fn new(params: SystemParams) -> Result<Self, openvm_cuda_common::error::CudaError> {
        validate_gpu_l_skip(params.l_skip)
            .expect("GPU backend requires l_skip <= 9 for current CUDA kernels");
        ensure_device_ntt_twiddles_initialized()
            .expect("failed to initialize small-NTT twiddles for current CUDA device");

        let prover_config = GpuProverConfig {
            zerocheck_save_memory: params.log_blowup == 1,
            ..Default::default()
        };
        let id = get_device().unwrap() as u32;
        let device_ctx = GpuDeviceCtx::for_device(id)?;
        let sm_count = get_sm_count(id).expect("failed to get SM count");
        let config = GpuDeviceConfig {
            params,
            prover_config,
            id,
            sm_count,
        };

        Ok(Self { config, device_ctx })
    }

    // Delegate accessors to inner config
    pub fn params(&self) -> &SystemParams {
        &self.config.params
    }

    pub fn prover_config(&self) -> &GpuProverConfig {
        &self.config.prover_config
    }

    pub fn prover_config_mut(&mut self) -> &mut GpuProverConfig {
        &mut self.config.prover_config
    }

    pub fn sm_count(&self) -> u32 {
        self.config.sm_count
    }

    pub fn id(&self) -> u32 {
        self.config.id
    }

    pub fn with_cache_rs_code_matrix(mut self, cache_rs_code_matrix: bool) -> Self {
        self.config.prover_config.cache_rs_code_matrix = cache_rs_code_matrix;
        self
    }

    pub fn set_cache_rs_code_matrix(&mut self, cache_rs_code_matrix: bool) {
        self.config.prover_config.cache_rs_code_matrix = cache_rs_code_matrix;
    }
}

/// Default GPU prover cache settings.
impl Default for GpuProverConfig {
    fn default() -> Self {
        Self {
            cache_stacked_matrix: false,
            cache_rs_code_matrix: false,
            zerocheck_save_memory: true,
        }
    }
}