openvm-cuda-backend 2.0.0

#![allow(dead_code)]
#![allow(clippy::missing_safety_doc)]

use openvm_cuda_common::{d_buffer::DeviceBuffer, error::CudaError, stream::cudaStream_t};
use openvm_stark_backend::prover::fractional_sumcheck_gkr::Frac;

use crate::prelude::{EF, F};

pub mod batch_ntt_small;
#[cfg(feature = "baby-bear-bn254-poseidon2")]
pub mod bn254_merkle_tree;
pub mod device_info;
pub mod logup_zerocheck;
pub mod matrix;
pub mod merkle_tree;
pub mod mle_interpolate;
pub mod ntt;
pub mod poly;
pub mod sponge;
pub mod stacked_reduction;
pub mod whir;

/// Log of warp size (32)
pub const LOG_WARP_SIZE: usize = 5;

pub mod sumcheck {
    use std::ffi::c_void;

    use super::*;
    use crate::poly::EqEvalSegments;

    const MAX_SUMCHECK_MLE_ROUND_D: u32 = 5;

    extern "C" {
        fn _sumcheck_mle_round(
            input_matrices: *const *const EF,
            output: *mut EF,
            tmp_block_sums: *mut EF,
            widths: *const u32,
            num_matrices: u32,
            height: u32,
            d: u32,
            stream: cudaStream_t,
        ) -> i32;

        fn _fold_mle(
            input_matrices: *const *const EF,
            output_matrices: *const *mut EF,
            widths: *const u32,
            num_matrices: u16,
            output_height: u32,
            max_output_cells: u32,
            r_val: EF,
            stream: cudaStream_t,
        ) -> i32;

        fn _fold_mle_column(
            buffer: *mut std::ffi::c_void,
            size: usize,
            r: EF,
            stream: cudaStream_t,
        ) -> i32;

        fn _batch_fold_mle(
            input_matrices: *const *const EF,
            output_matrices: *const *mut EF,
            widths: *const u32,
            num_matrices: u16,
            log_output_heights: *const u8,
            max_output_cells: u32,
            r_val: EF,
            stream: cudaStream_t,
        ) -> i32;

        fn _reduce_over_x_and_cols(
            input: *const std::ffi::c_void,
            output: *mut std::ffi::c_void,
            num_x: u32,
            num_cols: u32,
            large_domain_size: u32,
            stream: cudaStream_t,
        ) -> i32;

        fn _fold_ple_from_coeffs(
            input_coeffs: *const std::ffi::c_void,
            output: *mut std::ffi::c_void,
            num_x: u32,
            width: u32,
            domain_size: u32,
            r: EF,
            stream: cudaStream_t,
        ) -> i32;

        fn _triangular_fold_mle(
            output: *mut EF,
            input: *const EF,
            r: EF,
            output_max_n: u32,
            stream: cudaStream_t,
        ) -> i32;
    }

    #[allow(clippy::too_many_arguments)]
    pub unsafe fn sumcheck_mle_round(
        input_matrices: &DeviceBuffer<*const EF>,
        output: &DeviceBuffer<EF>,
        tmp_block_sums: &DeviceBuffer<EF>,
        widths: &DeviceBuffer<u32>,
        num_matrices: u32,
        height: u32,
        d: u32,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        if d == 0 || d > MAX_SUMCHECK_MLE_ROUND_D {
            return Err(CudaError::new(1));
        }
        CudaError::from_result(_sumcheck_mle_round(
            input_matrices.as_ptr(),
            output.as_mut_ptr(),
            tmp_block_sums.as_mut_ptr(),
            widths.as_ptr(),
            num_matrices,
            height,
            d,
            stream,
        ))
    }

    /// # Safety
    /// - `input_matrices` must consist of pointers to device memory locations.
    /// - `output_matrices` must consist of pointers to device memory locations.
    #[allow(clippy::too_many_arguments)]
    pub unsafe fn fold_mle(
        input_matrices: &DeviceBuffer<*const EF>,
        output_matrices: &DeviceBuffer<*mut EF>,
        widths: &DeviceBuffer<u32>,
        num_matrices: u16,
        output_height: u32,
        max_output_cells: u32,
        r_val: EF,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_fold_mle(
            input_matrices.as_ptr(),
            output_matrices.as_ptr(),
            widths.as_ptr(),
            num_matrices,
            output_height,
            max_output_cells,
            r_val,
            stream,
        ))
    }

    pub unsafe fn fold_mle_column(
        buffer: &mut DeviceBuffer<EF>,
        size: usize,
        r: EF,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_fold_mle_column(buffer.as_mut_raw_ptr(), size, r, stream))
    }

    /// # Safety
    /// - `input_matrices` must consist of pointers to device memory locations.
    /// - `output_matrices` must consist of pointers to device memory locations.
    #[allow(clippy::too_many_arguments)]
    pub unsafe fn batch_fold_mle(
        input_matrices: &DeviceBuffer<*const EF>,
        output_matrices: &DeviceBuffer<*mut EF>,
        widths: &DeviceBuffer<u32>,
        num_matrices: u16,
        log_output_heights: &DeviceBuffer<u8>,
        max_output_cells: u32,
        r_val: EF,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_batch_fold_mle(
            input_matrices.as_ptr(),
            output_matrices.as_ptr(),
            widths.as_ptr(),
            num_matrices,
            log_output_heights.as_ptr(),
            max_output_cells,
            r_val,
            stream,
        ))
    }

    pub unsafe fn fold_ple_from_coeffs(
        input_coeffs: *const F, // Base field (F)
        output: *mut EF,        // Extension field (EF)
        num_x: u32,
        width: u32,
        domain_size: u32,
        r: EF,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_fold_ple_from_coeffs(
            input_coeffs as *const c_void,
            output as *mut c_void,
            num_x,
            width,
            domain_size,
            r,
            stream,
        ))
    }

    pub unsafe fn reduce_over_x_and_cols<T>(
        input: &DeviceBuffer<T>,
        output: &DeviceBuffer<T>,
        num_x: u32,
        num_cols: u32,
        large_domain_size: u32,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_reduce_over_x_and_cols(
            input.as_raw_ptr(),
            output.as_mut_raw_ptr(),
            num_x,
            num_cols,
            large_domain_size,
            stream,
        ))
    }

    /// Folds the segments of `input` onto `output` using random element `r`.
    ///
    /// # Safety
    /// - `output` must have max `n` equal to `output_max_n`, for total length `2 * 2^output_max_n`.
    /// - `input` must have length `2 * 2^{output_max_n + 1}`.
    pub unsafe fn triangular_fold_mle(
        output: &mut EqEvalSegments<EF>,
        input: &EqEvalSegments<EF>,
        r: EF,
        output_max_n: usize,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        debug_assert_eq!(output.buffer.len(), 2 << output_max_n);
        debug_assert_eq!(input.buffer.len(), 4 << output_max_n);
        CudaError::from_result(_triangular_fold_mle(
            output.buffer.as_mut_ptr(),
            input.buffer.as_ptr(),
            r,
            output_max_n as u32,
            stream,
        ))
    }
}

// relate to prefix.cu
pub mod prefix {
    use super::*;

    extern "C" {
        fn _prefix_scan_block_ext(
            d_inout: *mut std::ffi::c_void,
            length: u64,
            round_stride: u64,
            block_num: u64,
            stream: cudaStream_t,
        ) -> i32;

        fn _prefix_scan_block_downsweep_ext(
            d_inout: *mut std::ffi::c_void,
            length: u64,
            round_stride: u64,
            stream: cudaStream_t,
        ) -> i32;

        fn _prefix_scan_epilogue_ext(
            d_inout: *mut std::ffi::c_void,
            length: u64,
            stream: cudaStream_t,
        ) -> i32;
    }

    pub unsafe fn prefix_scan_block_ext<T>(
        d_inout: &DeviceBuffer<T>,
        length: u64,
        round_stride: u64,
        block_num: u64,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_prefix_scan_block_ext(
            d_inout.as_mut_raw_ptr(),
            length,
            round_stride,
            block_num,
            stream,
        ))
    }

    pub unsafe fn prefix_scan_block_downsweep_ext<T>(
        d_inout: &DeviceBuffer<T>,
        length: u64,
        round_stride: u64,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_prefix_scan_block_downsweep_ext(
            d_inout.as_mut_raw_ptr(),
            length,
            round_stride,
            stream,
        ))
    }

    pub unsafe fn prefix_scan_epilogue_ext<T>(
        d_inout: &DeviceBuffer<T>,
        length: u64,
        stream: cudaStream_t,
    ) -> Result<(), CudaError> {
        CudaError::from_result(_prefix_scan_epilogue_ext(
            d_inout.as_mut_raw_ptr(),
            length,
            stream,
        ))
    }
}