openvm_cuda_backend/
poly.rs

1use std::sync::Arc;
2
3use getset::Getters;
4use openvm_cuda_common::{
5    copy::{MemCopyD2D, MemCopyH2D},
6    d_buffer::DeviceBuffer,
7    error::CudaError,
8    stream::{cudaStream_t, GpuDeviceCtx},
9};
10use openvm_stark_backend::prover::MatrixDimensions;
11use p3_field::PrimeCharacteristicRing;
12
13use crate::{
14    base::DeviceMatrix,
15    cuda::{
16        batch_ntt_small::{batch_ntt_small, validate_gpu_l_skip},
17        mle_interpolate::{
18            mle_interpolate_fused_2d, mle_interpolate_shared_2d, mle_interpolate_stage_2d,
19            MLE_SHARED_TILE_LOG_SIZE,
20        },
21        poly::{
22            eq_hypercube_interleaved_stage_ext, eq_hypercube_nonoverlapping_stage_ext,
23            eq_hypercube_stage_ext, mobius_eq_hypercube_stage_ext,
24        },
25        sumcheck::fold_mle_column,
26        LOG_WARP_SIZE,
27    },
28    prelude::{EF, F},
29    KernelError,
30};
31
32#[derive(derive_new::new, Getters)]
33pub struct PleMatrix<F> {
34    /// Stores the coefficient form of the univariate polynomial `f(Z, \vect x)` for each `\vect x
35    /// in H_n`. Buffer size is the same as `evals`.
36    #[getset(get = "pub")]
37    pub(crate) mixed: DeviceBuffer<F>,
38    height: usize,
39    width: usize,
40}
41
42impl<F> MatrixDimensions for PleMatrix<F> {
43    fn width(&self) -> usize {
44        self.width
45    }
46
47    fn height(&self) -> usize {
48        self.height
49    }
50}
51
52impl PleMatrix<F> {
53    /// Creates a `PleMatrix`. This doubles the VRAM footprint to cache the `mixed` buffer.
54    pub fn from_evals(
55        l_skip: usize,
56        evals: DeviceBuffer<F>,
57        height: usize,
58        width: usize,
59        device_ctx: &GpuDeviceCtx,
60    ) -> Self {
61        validate_gpu_l_skip(l_skip).expect("GPU PleMatrix requires l_skip <= 9");
62        let mut mixed = evals;
63        if l_skip > 0 {
64            // For univariate coordinate, perform inverse NTT for each 2^l_skip chunk per column:
65            // (width cols) * (height / 2^l_skip chunks per col). Use natural ordering.
66            let num_uni_poly = width * (height >> l_skip);
67            unsafe {
68                batch_ntt_small(
69                    &mut mixed,
70                    l_skip,
71                    num_uni_poly,
72                    true,
73                    device_ctx.stream.as_raw(),
74                )
75                .unwrap();
76            }
77        }
78        Self {
79            mixed,
80            height,
81            width,
82        }
83    }
84
85    pub fn to_evals(
86        &self,
87        l_skip: usize,
88        device_ctx: &GpuDeviceCtx,
89    ) -> Result<DeviceMatrix<F>, KernelError> {
90        validate_gpu_l_skip(l_skip)?;
91        let width = self.width();
92        let height = self.height();
93        // D2D copy so we can do in-place NTT
94        let mut evals = self.mixed.device_copy_on(device_ctx)?;
95        if l_skip > 0 {
96            // For univariate coordinate, perform NTT for each 2^l_skip chunk per column: (width
97            // cols) * (height / 2^l_skip chunks per col). Use natural ordering.
98            let num_uni_poly = width * (height >> l_skip);
99            unsafe {
100                batch_ntt_small(
101                    &mut evals,
102                    l_skip,
103                    num_uni_poly,
104                    false,
105                    device_ctx.stream.as_raw(),
106                )
107                .unwrap();
108            }
109        }
110        Ok(DeviceMatrix::new(Arc::new(evals), height, width))
111    }
112}
113
114/// Assumes that `evals` is column-major matrix of evaluations on a hypercube `H_n`.
115/// In-place interpolates `evals` from evaluations to coefficient form.
116pub fn mle_evals_to_coeffs_inplace(
117    evals: &mut DeviceBuffer<F>,
118    n: usize,
119    device_ctx: &GpuDeviceCtx,
120) -> Result<(), CudaError> {
121    if n == 0 {
122        return Ok(());
123    }
124    debug_assert!(evals.len().is_multiple_of(1 << n));
125    let width = evals.len() >> n;
126    // SAFETY:
127    // - `evals` is valid buffer, representing a single column vector, so width = 1
128    // - we only use forward direction with bit_reversed = false
129    unsafe {
130        mle_interpolate_stages(
131            evals.as_mut_ptr(),
132            width,
133            1 << n,
134            0,
135            0,
136            n as u32 - 1,
137            true,
138            false,
139            device_ctx.stream.as_raw(),
140        )
141    }
142}
143
144/// Helper function to process MLE interpolation stages from `start_log_step` to `end_log_step`
145/// (both inclusive).
146///
147/// Automatically selects the best kernel(s):
148/// - Warp shuffle for steps 1-16 (log_step 0-4) when 2+ stages can be fused
149/// - Shared memory for steps up to 2048 (log_step up to 11)
150/// - Fallback individual kernels for larger steps
151///
152/// # Parameters
153/// - `log_blowup`: meaningful_count = padded_height >> log_blowup
154/// - `bit_reversed`: if true, physical index = tidx << log_blowup (strided access) if false,
155///   physical index = tidx (contiguous access)
156///
157/// # Assumptions
158/// - If `right_pad` is true, `end_log_step` must be `< MLE_SHARED_TILE_LOG_SIZE`.
159/// # Safety
160/// Same requirements as the individual kernel functions.
161#[allow(clippy::too_many_arguments)]
162pub unsafe fn mle_interpolate_stages(
163    buffer: *mut F,
164    width: usize,
165    padded_height: u32,
166    log_blowup: u32,
167    start_log_step: u32,
168    end_log_step: u32,
169    is_eval_to_coeff: bool,
170    right_pad: bool,
171    stream: cudaStream_t,
172) -> Result<(), CudaError> {
173    if start_log_step > end_log_step {
174        return Ok(());
175    }
176
177    let mut current_log_step = start_log_step;
178
179    // Phase 1: Use warp shuffle for steps where log_step < LOG_WARP_SIZE (step <= 16)
180    // Only if we can fuse at least 2 stages (otherwise shared memory is just as good)
181    let warp_end = end_log_step.min(LOG_WARP_SIZE as u32 - 1);
182    let warp_stages = warp_end.saturating_sub(current_log_step) + 1;
183    if current_log_step < LOG_WARP_SIZE as u32 && warp_stages >= 2 {
184        let num_stages = warp_stages;
185        let start_step = 1u32 << current_log_step;
186
187        mle_interpolate_fused_2d(
188            buffer,
189            width,
190            padded_height,
191            log_blowup,
192            start_step,
193            num_stages,
194            is_eval_to_coeff,
195            right_pad,
196            stream,
197        )?;
198
199        current_log_step = warp_end + 1;
200    }
201
202    if current_log_step > end_log_step {
203        return Ok(());
204    }
205
206    // Phase 2: Use shared memory for steps where log_step < MLE_SHARED_TILE_LOG_SIZE
207    if current_log_step < MLE_SHARED_TILE_LOG_SIZE {
208        let shared_end = end_log_step.min(MLE_SHARED_TILE_LOG_SIZE - 1);
209
210        mle_interpolate_shared_2d(
211            buffer,
212            width,
213            padded_height,
214            log_blowup,
215            current_log_step,
216            shared_end,
217            is_eval_to_coeff,
218            right_pad,
219            stream,
220        )?;
221
222        current_log_step = shared_end + 1;
223    }
224
225    // Phase 3: Fallback to individual kernel launches for very large steps
226    // Note: fallback kernel doesn't support bit_reversed mode, only use for forward
227    assert!(
228        current_log_step > end_log_step || !right_pad,
229        "bit_reversed mode not supported for log_step >= MLE_SHARED_TILE_LOG_SIZE"
230    );
231    let height = padded_height >> log_blowup;
232    while current_log_step <= end_log_step {
233        let step = 1u32 << current_log_step;
234        mle_interpolate_stage_2d(
235            buffer,
236            width,
237            height,
238            padded_height,
239            step,
240            is_eval_to_coeff,
241            stream,
242        )?;
243        current_log_step += 1;
244    }
245
246    Ok(())
247}
248
249/// Given vector `x` in `F^n`, populates `out` with `eq_n(x, y)` for `y` on hypercube `H_n`.
250///
251/// Note: This function launches `n` CUDA kernels.
252///
253/// The multilinear equality polynomial is defined as:
254/// ```text
255///     eq(x, z) = \prod_{i=0}^{n-1} (x_i z_i + (1 - x_i)(1 - z_i)).
256/// ```
257///
258/// # Safety
259/// - `n` is set to the length of `xs`.
260/// - `out` must have length `>= 2^n`.
261pub unsafe fn evals_eq_hypercube(
262    out: &mut DeviceBuffer<EF>,
263    xs: &[EF],
264    device_ctx: &GpuDeviceCtx,
265) -> Result<(), KernelError> {
266    let n = xs.len();
267    assert!(out.len() >= 1 << n);
268    // Use memcpy instead of memset since EF will be in Montgomery form.
269    [EF::ONE]
270        .copy_to_on(out, device_ctx)
271        .map_err(KernelError::MemCopy)?;
272
273    for (i, &x_i) in xs.iter().enumerate() {
274        let step = 1 << i;
275        eq_hypercube_stage_ext(out.as_mut_ptr(), x_i, step, device_ctx.stream.as_raw())
276            .map_err(KernelError::Kernel)?;
277    }
278    Ok(())
279}
280
281/// Given vector `u_tilde` in `F^n`, populates `out` with `mobius_eq(u_tilde, y)` for `y` on
282/// hypercube `H_n`.
283///
284/// The Möbius-adjusted equality kernel is defined as:
285/// ```text
286///     mobius_eq(u_tilde, y) = \prod_{i=0}^{n-1} ((1 - 2*u_tilde_i)(1 - y_i) + u_tilde_i * y_i).
287/// ```
288///
289/// Note: This function launches `n` CUDA kernels.
290///
291/// # Safety
292/// - `n` is set to the length of `omega`.
293/// - `out` must have length `>= 2^n`.
294pub unsafe fn evals_mobius_eq_hypercube(
295    out: &mut DeviceBuffer<EF>,
296    omega: &[EF],
297    device_ctx: &GpuDeviceCtx,
298) -> Result<(), KernelError> {
299    let n = omega.len();
300    assert!(out.len() >= 1 << n);
301    // Use memcpy instead of memset since EF will be in Montgomery form.
302    [EF::ONE]
303        .copy_to_on(out, device_ctx)
304        .map_err(KernelError::MemCopy)?;
305
306    for (i, &omega_i) in omega.iter().enumerate() {
307        let step = 1 << i;
308        mobius_eq_hypercube_stage_ext(out.as_mut_ptr(), omega_i, step, device_ctx.stream.as_raw())
309            .map_err(KernelError::Kernel)?;
310    }
311    Ok(())
312}
313
314/// For a fixed `x` in `F^max_n`, stores the evaluations of `eq_n(x[..n], -)` on hypercube `H_n` for
315/// `n = 0..=max_n`. The evaluations are stored contiguously in a single buffer of length `2^(max_n
316/// + 1)`. We define `eq_0 = 1` always.
317#[derive(Getters)]
318pub struct EqEvalSegments<F> {
319    /// Index 0 is unused and initialized to zero. Index 1 stores the length-1 `eq_0 = 1` layer.
320    #[getset(get = "pub")]
321    pub(crate) buffer: DeviceBuffer<F>,
322    #[getset(get_copy = "pub")]
323    max_n: usize,
324}
325
326impl<F> EqEvalSegments<F> {
327    /// Returns start pointer for buffer of length `2^n` corresponding to evaluations of `eq(x[..n],
328    /// -)` on `H_n`.
329    pub fn get_ptr(&self, n: usize) -> *const F {
330        assert!(n <= self.max_n);
331        // SAFETY: `buffer` has length `2^{max_n + 1}`
332        unsafe { self.buffer.as_ptr().add(1 << n) }
333    }
334
335    /// # Safety
336    /// Caller must ensure that `buffer` has length `2^{max_n + 1}` and is properly initialized.
337    pub unsafe fn from_raw_parts(buffer: DeviceBuffer<F>, max_n: usize) -> Self {
338        Self { buffer, max_n }
339    }
340}
341
342// Currently only implement kernels for EF.
343impl EqEvalSegments<EF> {
344    /// Creates a new `EqEvalSegments` instance with `max_n = x.len()`.
345    pub fn new(x: &[EF], device_ctx: &GpuDeviceCtx) -> Result<Self, KernelError> {
346        let max_n = x.len();
347        let mut buffer = DeviceBuffer::with_capacity_on(2 << max_n, device_ctx);
348        // Index 0 should never to be used, but we initialize it to zero.
349        // Index 1 is set to eq_0 = 1 for initial state
350        [EF::ZERO, EF::ONE]
351            .copy_to_on(&mut buffer, device_ctx)
352            .map_err(KernelError::MemCopy)?;
353        // At step i, we populate `eq_{i+1}` starting at offset `2^{i+1}`
354        for (i, &x_i) in x.iter().enumerate() {
355            let step = 1 << i;
356            // SAFETY:
357            // - `buffer` has length `2^{max_n + 1}`
358            // - `2 * 2^{i+1} <= 2^{max_n + 1}` ensures everything is in bounds
359            unsafe {
360                // The `eq_{i+1}` segment starts at offset `2^{i+1}`
361                let dst = buffer.as_mut_ptr().add(2 * step);
362                // The `eq_i` segment starts at offset `2^i`
363                let src = buffer.as_ptr().add(step);
364                eq_hypercube_nonoverlapping_stage_ext(
365                    dst,
366                    src,
367                    x_i,
368                    step as u32,
369                    device_ctx.stream.as_raw(),
370                )
371                .map_err(KernelError::Kernel)?;
372            }
373        }
374        Ok(Self { buffer, max_n })
375    }
376}
377
378/// Same as [EqEvalSegments] but keeping segment tree with buffers separated by layer to allow
379/// dropping layers.
380#[derive(Getters)]
381pub struct EqEvalLayers<F> {
382    /// Layers are reference counted so the length-1 seed layer can be shared across trees.
383    pub layers: Vec<Arc<DeviceBuffer<F>>>,
384}
385
386impl<F> EqEvalLayers<F> {
387    /// Returns start pointer for buffer of length `2^n` corresponding to evaluations of `eq(x[..n],
388    /// -)` on `H_n`.
389    pub fn get_ptr(&self, n: usize) -> *const F {
390        debug_assert_eq!(self.layers[n].len(), 1 << n);
391        self.layers[n].as_ptr()
392    }
393}
394
395// Currently only implement kernels for EF.
396impl EqEvalLayers<EF> {
397    pub(crate) fn one_layer(
398        device_ctx: &GpuDeviceCtx,
399    ) -> Result<Arc<DeviceBuffer<EF>>, KernelError> {
400        [EF::ONE]
401            .to_device_on(device_ctx)
402            .map(Arc::new)
403            .map_err(KernelError::MemCopy)
404    }
405
406    /// Creates a new `EqEvalLayers` instance with `layers.len() = x.len() + 1`.
407    ///
408    /// Inserts `x_i` from the front for each layer.
409    pub fn new_rev<'a>(
410        n: usize,
411        x: impl IntoIterator<Item = &'a EF>,
412        device_ctx: &GpuDeviceCtx,
413    ) -> Result<Self, KernelError> {
414        Self::new_rev_with_one(n, x, Self::one_layer(device_ctx)?, device_ctx)
415    }
416
417    pub(crate) fn new_rev_with_one<'a>(
418        n: usize,
419        x: impl IntoIterator<Item = &'a EF>,
420        layer_0: Arc<DeviceBuffer<EF>>,
421        device_ctx: &GpuDeviceCtx,
422    ) -> Result<Self, KernelError> {
423        let mut layers = Vec::with_capacity(n + 1);
424        layers.push(layer_0);
425        for (i, &x_i) in x.into_iter().enumerate() {
426            let step = 1 << i;
427            let buffer = DeviceBuffer::with_capacity_on(2 * step, device_ctx);
428            // SAFETY:
429            // - `buffer` has length `2^{i+1}`
430            unsafe {
431                let dst = buffer.as_mut_ptr();
432                let src = layers.last().unwrap().as_ptr();
433                eq_hypercube_interleaved_stage_ext(
434                    dst,
435                    src,
436                    x_i,
437                    step as u32,
438                    device_ctx.stream.as_raw(),
439                )
440                .map_err(KernelError::Kernel)?;
441            }
442            layers.push(Arc::new(buffer));
443        }
444        Ok(Self { layers })
445    }
446
447    /// Creates a new `EqEvalLayers` instance with `layers.len() = x.len() + 1`.
448    ///
449    /// Inserts `x_i` from the back for each layer. This matches behavior of
450    /// [`EqEvalSegments::new`].
451    pub fn new<'a>(
452        n: usize,
453        x: impl IntoIterator<Item = &'a EF>,
454        device_ctx: &GpuDeviceCtx,
455    ) -> Result<Self, KernelError> {
456        Self::new_with_one(n, x, Self::one_layer(device_ctx)?, device_ctx)
457    }
458
459    pub(crate) fn new_with_one<'a>(
460        n: usize,
461        x: impl IntoIterator<Item = &'a EF>,
462        layer_0: Arc<DeviceBuffer<EF>>,
463        device_ctx: &GpuDeviceCtx,
464    ) -> Result<Self, KernelError> {
465        let mut layers = Vec::with_capacity(n + 1);
466        layers.push(layer_0);
467        for (i, &x_i) in x.into_iter().enumerate() {
468            let step = 1 << i;
469            let buffer = DeviceBuffer::with_capacity_on(2 * step, device_ctx);
470            // SAFETY:
471            // - `buffer` has length `2^{i+1}`
472            unsafe {
473                let dst = buffer.as_mut_ptr();
474                let src = layers.last().unwrap().as_ptr();
475                eq_hypercube_nonoverlapping_stage_ext(
476                    dst,
477                    src,
478                    x_i,
479                    step as u32,
480                    device_ctx.stream.as_raw(),
481                )
482                .map_err(KernelError::Kernel)?;
483            }
484            layers.push(Arc::new(buffer));
485        }
486        Ok(Self { layers })
487    }
488}
489
490/// Square-root decomposition of hypercube equality buffer for memory optimization.
491///
492/// Instead of storing a single buffer of size 2^n, we store two buffers of size
493/// 2^(n/2), reducing memory usage from O(2^n) to O(2*2^(n/2)). The full value at
494/// index i is computed on-the-fly as: low[i % low_cap] * high[i / low_cap]
495pub struct SqrtHyperBuffer {
496    pub low: DeviceBuffer<EF>,
497    pub high: DeviceBuffer<EF>,
498    pub low_capacity: usize,
499    pub size: usize,
500}
501
502impl SqrtHyperBuffer {
503    /// Build a buffer from `xi`. Note that last elements of `xi` correspond to the lowest index
504    /// bits.
505    pub fn from_xi(xi: &[EF], device_ctx: &GpuDeviceCtx) -> Result<Self, KernelError> {
506        let low = {
507            let mut res = DeviceBuffer::with_capacity_on(1 << (xi.len() / 2), device_ctx);
508            unsafe { evals_eq_hypercube(&mut res, &xi[..xi.len() / 2], device_ctx)? };
509            res
510        };
511        let high = {
512            let mut res = DeviceBuffer::with_capacity_on(1 << xi.len().div_ceil(2), device_ctx);
513            unsafe { evals_eq_hypercube(&mut res, &xi[xi.len() / 2..], device_ctx)? };
514            res
515        };
516        Ok(Self {
517            low,
518            high,
519            low_capacity: 1 << (xi.len() / 2),
520            size: 1 << xi.len(),
521        })
522    }
523
524    pub fn fold_columns(&mut self, r: EF, device_ctx: &GpuDeviceCtx) -> Result<(), CudaError> {
525        assert!(self.size > 1);
526        if self.size > self.low_capacity {
527            unsafe {
528                fold_mle_column(
529                    &mut self.high,
530                    self.size / self.low_capacity,
531                    r,
532                    device_ctx.stream.as_raw(),
533                )?;
534            }
535        } else {
536            unsafe {
537                fold_mle_column(&mut self.low, self.size, r, device_ctx.stream.as_raw())?;
538            }
539        };
540        self.size /= 2;
541        Ok(())
542    }
543}
544
545/// Square-root decomposition using pre-computed eq evaluation layers.
546///
547/// Instead of folding columns on GPU, we pre-compute all layers of eq evaluations
548/// and simply drop the highest layer each round (no GPU work needed).
549pub struct SqrtEqLayers {
550    /// Layers for `xi[(n + 1) / 2..]`. Layers insert `xi_i` from the front, so `low` contains the
551    /// latter half of `xi`.
552    pub low: EqEvalLayers<EF>,
553    /// Layers for `xi[..(n + 1) / 2]`.
554    /// Layers: [ eq(xi[j..(n+1)/2], ..) ] for j = (0..(n+1)/2).rev()
555    pub high: EqEvalLayers<EF>,
556}
557
558impl SqrtEqLayers {
559    /// Build layers from `xi` values.
560    ///
561    /// This is meant to match behavior of [SqrtHyperBuffer::from_xi] but with layers.
562    ///
563    /// Example: This means `[a,b,c,d]` should be sent to `low: [[d], [d, c]], high: [[b], [b, a]]`.
564    pub fn from_xi(xi: &[EF], device_ctx: &GpuDeviceCtx) -> Result<Self, KernelError> {
565        let n = xi.len();
566        let low_n = n / 2;
567        let high_n = n - low_n;
568
569        let layer_0 = EqEvalLayers::one_layer(device_ctx)?;
570        let low = EqEvalLayers::new_with_one(
571            low_n,
572            xi[high_n..].iter().rev(),
573            Arc::clone(&layer_0),
574            device_ctx,
575        )?;
576        let high =
577            EqEvalLayers::new_with_one(high_n, xi[..high_n].iter().rev(), layer_0, device_ctx)?;
578
579        Ok(Self { low, high })
580    }
581
582    pub fn max_n(&self) -> usize {
583        self.low_n() + self.high_n()
584    }
585
586    pub fn low_n(&self) -> usize {
587        self.low.layers.len() - 1
588    }
589
590    pub fn high_n(&self) -> usize {
591        self.high.layers.len() - 1
592    }
593
594    /// Drop the highest layer. Drops from high first, then low. This corresponding to `pop_front`
595    /// from `xi`.
596    pub fn drop_layer(&mut self) {
597        if self.high.layers.len() > 1 {
598            self.high.layers.pop();
599        } else if self.low.layers.len() > 1 {
600            self.low.layers.pop();
601        }
602    }
603}
openvm_cuda_backend/poly.rs

openvm_cuda_backend/
poly.rs
