sp1_gpu_sys/
dft.rs

1use sp1_primitives::SP1Field;
2
3use crate::runtime::{CudaRustError, CudaStreamHandle, DEFAULT_STREAM};
4
5/// # Safety
6///
7/// Initializes the selected GPU DFT backend.
8pub unsafe fn dft_init_default_stream() -> CudaRustError {
9    sppark_init(DEFAULT_STREAM)
10}
11
12/// # Safety
13///
14/// Initializes twiddle tables for the selected GPU DFT backend.
15pub unsafe fn dft_init_twiddles(max_log_size: u32, stream: CudaStreamHandle) -> CudaRustError {
16    dft_init_twiddles_sys(max_log_size, stream)
17}
18
19extern "C" {
20    pub fn sppark_init(stream: CudaStreamHandle) -> CudaRustError;
21
22    #[link_name = "dft_init_twiddles"]
23    fn dft_init_twiddles_sys(max_log_size: u32, stream: CudaStreamHandle) -> CudaRustError;
24
25    pub fn batch_coset_dft(
26        d_out: *mut SP1Field,
27        d_in: *mut SP1Field,
28        lg_domain_size: u32,
29        lg_blowup: u32,
30        shift: SP1Field,
31        poly_count: u32,
32        is_bit_rev: bool,
33        stream: CudaStreamHandle,
34    ) -> CudaRustError;
35
36    pub fn batch_lde_shift_in_place(
37        d_inout: *mut SP1Field,
38        lg_domain_size: u32,
39        lg_blowup: u32,
40        shift: SP1Field,
41        poly_count: u32,
42        is_bit_rev: bool,
43        stream: CudaStreamHandle,
44    ) -> CudaRustError;
45
46    pub fn batch_coset_dft_in_place(
47        d_inout: *mut SP1Field,
48        lg_domain_size: u32,
49        lg_blowup: u32,
50        shift: SP1Field,
51        poly_count: u32,
52        is_bit_rev: bool,
53        stream: CudaStreamHandle,
54    ) -> CudaRustError;
55
56    pub fn batch_NTT(
57        d_inout: *mut SP1Field,
58        lg_domain_size: u32,
59        poly_count: u32,
60        stream: CudaStreamHandle,
61    ) -> CudaRustError;
62
63    pub fn batch_iNTT(
64        d_inout: *mut SP1Field,
65        lg_domain_size: u32,
66        poly_count: u32,
67        stream: CudaStreamHandle,
68    ) -> CudaRustError;
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_dft_init() {
77        unsafe {
78            assert!(dft_init_default_stream() == crate::runtime::CUDA_SUCCESS_CSL);
79            assert!(dft_init_twiddles(15, DEFAULT_STREAM) == crate::runtime::CUDA_SUCCESS_CSL);
80        }
81    }
82}
sp1_gpu_sys/dft.rs

sp1_gpu_sys/
dft.rs
