openvm_cuda_backend/
ntt.rs

1use std::{
2    collections::BTreeSet,
3    sync::{Mutex, OnceLock},
4};
5
6use openvm_cuda_common::{
7    common::{device_reset_epoch, get_device},
8    d_buffer::DeviceBuffer,
9    stream::GpuDeviceCtx,
10};
11
12use crate::{cuda::ntt, prelude::F};
13
14const MAX_LG_DOMAIN_SIZE: usize = ntt::MAX_CUDA_NTT_LOG_DOMAIN_SIZE as usize;
15const LG_WINDOW_SIZE: usize = MAX_LG_DOMAIN_SIZE.div_ceil(5);
16const WINDOW_SIZE: usize = 1 << LG_WINDOW_SIZE;
17const WINDOW_NUM: usize = MAX_LG_DOMAIN_SIZE.div_ceil(LG_WINDOW_SIZE);
18const RADIX_TWIDDLES_SIZE: usize = 32 + 64 + 128 + 256 + 512;
19
20static INIT_FORWARD: OnceLock<Mutex<BTreeSet<(i32, u64)>>> = OnceLock::new();
21static INIT_INVERSE: OnceLock<Mutex<BTreeSet<(i32, u64)>>> = OnceLock::new();
22
23fn ensure_initialized(inverse: bool) -> Result<(), openvm_cuda_common::error::CudaError> {
24    let device_key = (get_device()?, device_reset_epoch());
25    let initialized = if inverse {
26        &INIT_INVERSE
27    } else {
28        &INIT_FORWARD
29    };
30    let initialized = initialized.get_or_init(|| Mutex::new(BTreeSet::new()));
31    let mut initialized = initialized.lock().unwrap();
32    if initialized.contains(&device_key) {
33        return Ok(());
34    }
35
36    {
37        let device_ctx = GpuDeviceCtx::for_device(device_key.0 as u32)?;
38        let partial_twiddles =
39            DeviceBuffer::<[F; WINDOW_SIZE]>::with_capacity_on(WINDOW_NUM, &device_ctx);
40        let twiddles = DeviceBuffer::<F>::with_capacity_on(RADIX_TWIDDLES_SIZE, &device_ctx);
41        unsafe {
42            // Both CUDA helpers upload into constant memory and synchronize before returning, so
43            // these staging buffers are safe to free once the calls complete.
44            ntt::generate_all_twiddles(&twiddles, inverse, device_ctx.stream.as_raw())?;
45            ntt::generate_partial_twiddles(&partial_twiddles, inverse, device_ctx.stream.as_raw())?;
46        }
47    }
48    initialized.insert(device_key);
49    Ok(())
50}
51
52struct NttImpl<'a> {
53    buffer: &'a DeviceBuffer<F>,
54    lg_domain_size: u32,
55    padded_poly_size: u32,
56    poly_count: u32,
57    is_intt: bool,
58    stage: u32,
59    device_ctx: GpuDeviceCtx,
60}
61
62impl<'a> NttImpl<'a> {
63    fn new(
64        buffer: &'a DeviceBuffer<F>,
65        lg_domain_size: u32,
66        padded_poly_size: u32,
67        poly_count: u32,
68        is_intt: bool,
69        device_ctx: &GpuDeviceCtx,
70    ) -> Self {
71        ensure_initialized(is_intt).expect("failed to initialize CUDA NTT twiddle tables");
72        Self {
73            buffer,
74            lg_domain_size,
75            padded_poly_size,
76            poly_count,
77            is_intt,
78            stage: 0,
79            device_ctx: device_ctx.clone(),
80        }
81    }
82
83    fn step(&mut self, iterations: u32) {
84        assert!(iterations <= 10);
85        let radix = if iterations < 6 { 6 } else { iterations };
86        unsafe {
87            ntt::ct_mixed_radix_narrow(
88                self.buffer,
89                radix,
90                self.lg_domain_size,
91                self.stage,
92                iterations,
93                self.padded_poly_size,
94                self.poly_count,
95                self.is_intt,
96                self.device_ctx.stream.as_raw(),
97            )
98            .expect("failed to launch CUDA mixed-radix NTT step");
99        }
100        self.stage += iterations;
101    }
102}
103
104/// Performs column-wise batch NTT on `buffer`, where `buffer` is assumed to be column-major with
105/// columns of height `2^(log_trace_height + log_blowup)`. The NTT are performed on the first
106/// `2^log_trace_height` elements of each column. If `bit_reverse` is true, then the input columns
107/// are assumed to be ordered in **natural** ordering, and a bit-reversal permutation is applied for
108/// the internal algorithm of the NTT. If `bit_reverse` is false, then the input columns are assumed
109/// to be in bit-reverse ordering. If `is_intt` is true, the inverse NTT is performed; otherwise,
110/// the forward NTT is performed.
111pub fn batch_ntt(
112    buffer: &DeviceBuffer<F>,
113    log_trace_height: u32,
114    log_blowup: u32,
115    width: u32,
116    bit_reverse: bool,
117    is_intt: bool,
118    device_ctx: &GpuDeviceCtx,
119) {
120    if log_trace_height == 0 {
121        return;
122    }
123    assert!(
124        log_trace_height <= ntt::MAX_CUDA_NTT_LOG_DOMAIN_SIZE,
125        "CUDA batch_ntt supports log_trace_height <= {}",
126        ntt::MAX_CUDA_NTT_LOG_DOMAIN_SIZE
127    );
128
129    let padded_poly_size = 1 << (log_trace_height + log_blowup);
130
131    if bit_reverse {
132        unsafe {
133            ntt::bit_rev(
134                buffer,
135                buffer,
136                log_trace_height,
137                padded_poly_size,
138                width,
139                device_ctx.stream.as_raw(),
140            )
141            .expect("failed to launch CUDA bit-reversal permutation");
142        }
143    }
144
145    let mut _impl = NttImpl::new(
146        buffer,
147        log_trace_height,
148        padded_poly_size,
149        width,
150        is_intt,
151        device_ctx,
152    );
153    if log_trace_height <= 10 {
154        _impl.step(log_trace_height);
155    } else if log_trace_height <= 17 {
156        let step = log_trace_height / 2;
157        _impl.step(step + log_trace_height % 2);
158        _impl.step(step);
159    } else if log_trace_height <= ntt::MAX_CUDA_NTT_LOG_DOMAIN_SIZE {
160        let step = log_trace_height / 3;
161        let rem = log_trace_height % 3;
162        _impl.step(step);
163        _impl.step(step);
164        _impl.step(step + rem);
165    } else {
166        unreachable!("log_trace_height is bounded above");
167    }
168}
openvm_cuda_backend/ntt.rs

openvm_cuda_backend/
ntt.rs
