iro_cuda_ffi_kernels/

lib.rs

//! Reference CUDA kernels for iro-cuda-ffi.
//!
//! This crate provides sample kernels that demonstrate proper iro-cuda-ffi usage patterns
//! and serve as integration tests for the iro-cuda-ffi core crate.
//!
//! # Available Kernels
//!
//! - `vector_add_f32`: Element-wise vector addition
//! - `fma_chain_f32`: Deep compute chain (FMA)
//! - `saxpy_f32`: Single-precision A*X + Y
//! - `daxpy_f64`: Double-precision A*X + Y
//! - `scale_f32`: Vector scaling
//! - `reduce_sum_f32`: Parallel sum reduction
//! - `reduce_max_f32`: Parallel max reduction
//!
//! # Example
//!
//! ```ignore
//! use iro_cuda_ffi::prelude::*;
//! use iro_cuda_ffi_kernels::vector_add_f32;
//!
//! let stream = Stream::new()?;
//!
//! let a = DeviceBuffer::from_slice_sync(&stream, &[1.0f32, 2.0, 3.0, 4.0])?;
//! let b = DeviceBuffer::from_slice_sync(&stream, &[5.0f32, 6.0, 7.0, 8.0])?;
//! let mut c = DeviceBuffer::<f32>::zeros(4)?;
//!
//! vector_add_f32(&stream, &a, &b, &mut c)?;
//!
//! let result = c.to_vec(&stream)?;
//! assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]);
//! ```
//!
//! ## Device-to-device copy
//!
//! ```ignore
//! use iro_cuda_ffi::prelude::*;
//!
//! let stream = Stream::new()?;
//! let src = DeviceBuffer::from_slice_sync(&stream, &[1.0f32, 2.0, 3.0])?;
//! let mut dst = DeviceBuffer::<f32>::alloc(src.len())?;
//!
//! dst.copy_from_device_sync(&stream, &src)?;
//! let result = dst.to_vec(&stream)?;
//! assert_eq!(result, vec![1.0, 2.0, 3.0]);
//! ```

#![warn(missing_docs)]
#![warn(clippy::all, clippy::pedantic)]

use iro_cuda_ffi::error::icffi_codes;
use iro_cuda_ffi::prelude::*;

// FFI declarations for kernel exports (Rust 2024: unsafe extern required)
mod ffi {
    use iro_cuda_ffi::abi::{InBufferDesc, LaunchParams, OutBufferDesc};

    // Kernel FFI declarations.
    //
    // All kernel wrapper functions follow the iro-cuda-ffi ABI:
    // - Return `cudaError_t` (i32 where 0 = success)
    // - Take `LaunchParams` by value as first argument
    // - Never synchronize internally
    unsafe extern "C" {
        pub fn icffi_vector_add_f32(
            p: LaunchParams,
            a: InBufferDesc<f32>,
            b: InBufferDesc<f32>,
            out: OutBufferDesc<f32>,
        ) -> i32;

        pub fn icffi_fma_chain_f32(
            p: LaunchParams,
            a: InBufferDesc<f32>,
            b: InBufferDesc<f32>,
            out: OutBufferDesc<f32>,
            iters: u32,
        ) -> i32;

        pub fn icffi_saxpy_f32(
            p: LaunchParams,
            x: InBufferDesc<f32>,
            y: OutBufferDesc<f32>,
            a: f32,
        ) -> i32;

        pub fn icffi_daxpy_f64(
            p: LaunchParams,
            x: InBufferDesc<f64>,
            y: OutBufferDesc<f64>,
            a: f64,
        ) -> i32;

        pub fn icffi_scale_f32(
            p: LaunchParams,
            x: InBufferDesc<f32>,
            y: OutBufferDesc<f32>,
            a: f32,
        ) -> i32;

        pub fn icffi_reduce_sum_f32(
            p: LaunchParams,
            input: InBufferDesc<f32>,
            output: OutBufferDesc<f32>,
        ) -> i32;

        pub fn icffi_reduce_max_f32(
            p: LaunchParams,
            input: InBufferDesc<f32>,
            output: OutBufferDesc<f32>,
        ) -> i32;

        // Ensure ABI asserts TU is linked
        pub fn icffi_abi_asserts_linked();
    }
}

const BLOCK_SIZE: usize = 256;
// Must match reduction kernels (each thread loads 2 elements).
const REDUCTION_ELEMENTS_PER_THREAD: usize = 2;
const REDUCTION_ELEMENTS_PER_BLOCK: usize = BLOCK_SIZE * REDUCTION_ELEMENTS_PER_THREAD;
#[allow(clippy::cast_possible_truncation)]
const BLOCK_SIZE_U32: u32 = BLOCK_SIZE as u32; // Safe: 256 fits in u32

#[inline]
fn ensure_len_eq(name: &str, left_label: &str, left: usize, right_label: &str, right: usize) -> Result<()> {
    if left != right {
        return Err(IcffiError::with_location(
            icffi_codes::LENGTH_MISMATCH,
            format!("{name}: length mismatch ({left_label}={left} != {right_label}={right})"),
        ));
    }
    Ok(())
}

/// Maximum grid_x dimension per CUDA spec (2^31 - 1).
///
/// This matches the limit in `LaunchParams::is_valid()`. The CUDA spec caps
/// grid_x at this value, not u32::MAX. Using a larger value would pass the
/// u32 check but fail the actual launch.
pub const MAX_GRID_X: usize = 0x7FFF_FFFF;

#[inline]
#[allow(clippy::cast_possible_truncation)]
fn grid_u32(name: &str, grid: usize) -> Result<u32> {
    if grid > MAX_GRID_X {
        return Err(IcffiError::with_location(
            icffi_codes::GRID_TOO_LARGE,
            format!("{name}: grid size exceeds MAX_GRID_X ({grid} > {MAX_GRID_X})"),
        ));
    }
    // The truncation is safe because we've already checked the value fits in i32
    Ok(grid as u32)
}

#[inline]
fn div_ceil_usize(n: usize, denom: usize) -> usize {
    debug_assert!(denom != 0);
    let q = n / denom;
    if n % denom == 0 {
        q
    } else {
        q + 1
    }
}

/// Computes the number of blocks needed for a 1D launch.
#[inline]
fn blocks_for(n: usize, block_size: usize) -> usize {
    div_ceil_usize(n, block_size)
}

#[inline]
fn reduce_sum_f32_len(
    stream: &Stream,
    input: &DeviceBuffer<f32>,
    input_len: usize,
    output: &mut DeviceBuffer<f32>,
) -> Result<usize> {
    if input_len == 0 {
        return Ok(0);
    }

    if input_len > input.len() {
        return Err(IcffiError::with_location(
            icffi_codes::INVALID_ARGUMENT,
            format!(
                "reduce_sum_f32: input_len exceeds buffer length ({} > {})",
                input_len,
                input.len()
            ),
        ));
    }

    let grid = reduction_output_size(input_len);
    if output.len() < grid {
        return Err(IcffiError::with_location(
            icffi_codes::OUTPUT_TOO_SMALL,
            format!(
                "reduce_sum_f32: output too small ({} < {})",
                output.len(),
                grid
            ),
        ));
    }

    let grid_u32 = grid_u32("reduce_sum_f32", grid)?;
    let params = LaunchParams::new_1d(grid_u32, BLOCK_SIZE_U32, stream.raw());
    let input_desc = InBufferDesc::new(input.as_ptr(), input_len as u64);
    let output_desc = OutBufferDesc::new(output.as_mut_ptr(), grid as u64);

    check(unsafe { ffi::icffi_reduce_sum_f32(params, input_desc, output_desc) })?;

    Ok(grid)
}

#[inline]
fn reduce_max_f32_len(
    stream: &Stream,
    input: &DeviceBuffer<f32>,
    input_len: usize,
    output: &mut DeviceBuffer<f32>,
) -> Result<usize> {
    if input_len == 0 {
        return Ok(0);
    }

    if input_len > input.len() {
        return Err(IcffiError::with_location(
            icffi_codes::INVALID_ARGUMENT,
            format!(
                "reduce_max_f32: input_len exceeds buffer length ({} > {})",
                input_len,
                input.len()
            ),
        ));
    }

    let grid = reduction_output_size(input_len);
    if output.len() < grid {
        return Err(IcffiError::with_location(
            icffi_codes::OUTPUT_TOO_SMALL,
            format!(
                "reduce_max_f32: output too small ({} < {})",
                output.len(),
                grid
            ),
        ));
    }

    let grid_u32 = grid_u32("reduce_max_f32", grid)?;
    let params = LaunchParams::new_1d(grid_u32, BLOCK_SIZE_U32, stream.raw());
    let input_desc = InBufferDesc::new(input.as_ptr(), input_len as u64);
    let output_desc = OutBufferDesc::new(output.as_mut_ptr(), grid as u64);

    check(unsafe { ffi::icffi_reduce_max_f32(params, input_desc, output_desc) })?;

    Ok(grid)
}

/// Element-wise vector addition: out = a + b
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `a` - First input vector
/// * `b` - Second input vector (must have same length as `a`)
/// * `out` - Output vector (must have same length as `a`)
///
/// # Errors
///
/// Returns an error if the vectors have mismatched lengths or kernel launch fails.
///
/// # Example
///
/// ```ignore
/// let stream = Stream::new()?;
/// let a = DeviceBuffer::from_slice_sync(&stream, &[1.0f32, 2.0, 3.0])?;
/// let b = DeviceBuffer::from_slice_sync(&stream, &[4.0f32, 5.0, 6.0])?;
/// let mut c = DeviceBuffer::<f32>::zeros(3)?;
///
/// vector_add_f32(&stream, &a, &b, &mut c)?;
/// ```
#[track_caller]
pub fn vector_add_f32(
    stream: &Stream,
    a: &DeviceBuffer<f32>,
    b: &DeviceBuffer<f32>,
    out: &mut DeviceBuffer<f32>,
) -> Result<()> {
    let n = a.len();
    ensure_len_eq("vector_add_f32", "b", b.len(), "a", n)?;
    ensure_len_eq("vector_add_f32", "out", out.len(), "a", n)?;

    if n == 0 {
        return Ok(());
    }

    let grid = blocks_for(n, BLOCK_SIZE);
    let grid = grid_u32("vector_add_f32", grid)?;
    let params = LaunchParams::new_1d(grid, BLOCK_SIZE_U32, stream.raw());

    check(unsafe { ffi::icffi_vector_add_f32(params, a.as_in(), b.as_in(), out.as_out()) })
}

/// Deep compute chain: `out = fma_chain(a, b, iters)`
///
/// Each element performs `iters` iterations of: `acc = acc * b[i] + 1.0`.
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `a` - First input vector
/// * `b` - Second input vector (must have same length as `a`)
/// * `out` - Output vector (must have same length as `a`)
/// * `iters` - Number of multiply-add iterations per element
///
/// # Errors
///
/// Returns an error if the vectors have mismatched lengths or kernel launch fails.
#[track_caller]
pub fn fma_chain_f32(
    stream: &Stream,
    a: &DeviceBuffer<f32>,
    b: &DeviceBuffer<f32>,
    out: &mut DeviceBuffer<f32>,
    iters: u32,
) -> Result<()> {
    let n = a.len();
    ensure_len_eq("fma_chain_f32", "b", b.len(), "a", n)?;
    ensure_len_eq("fma_chain_f32", "out", out.len(), "a", n)?;

    if n == 0 {
        return Ok(());
    }

    let grid = blocks_for(n, BLOCK_SIZE);
    let grid = grid_u32("fma_chain_f32", grid)?;
    let params = LaunchParams::new_1d(grid, BLOCK_SIZE_U32, stream.raw());

    check(unsafe { ffi::icffi_fma_chain_f32(params, a.as_in(), b.as_in(), out.as_out(), iters) })
}

/// SAXPY operation: y = a * x + y (in-place)
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `a` - Scalar multiplier
/// * `x` - Input vector
/// * `y` - Input/output vector (modified in place)
///
/// # Errors
///
/// Returns an error if the vectors have mismatched lengths or kernel launch fails.
#[track_caller]
pub fn saxpy_f32(
    stream: &Stream,
    a: f32,
    x: &DeviceBuffer<f32>,
    y: &mut DeviceBuffer<f32>,
) -> Result<()> {
    let n = x.len();
    ensure_len_eq("saxpy_f32", "y", y.len(), "x", n)?;

    if n == 0 {
        return Ok(());
    }

    let grid = blocks_for(n, BLOCK_SIZE);
    let grid = grid_u32("saxpy_f32", grid)?;
    let params = LaunchParams::new_1d(grid, BLOCK_SIZE_U32, stream.raw());

    check(unsafe { ffi::icffi_saxpy_f32(params, x.as_in(), y.as_out(), a) })
}

/// DAXPY operation: y = a * x + y (in-place, double precision)
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `a` - Scalar multiplier
/// * `x` - Input vector
/// * `y` - Input/output vector (modified in place)
///
/// # Errors
///
/// Returns an error if the vectors have mismatched lengths or kernel launch fails.
#[track_caller]
pub fn daxpy_f64(
    stream: &Stream,
    a: f64,
    x: &DeviceBuffer<f64>,
    y: &mut DeviceBuffer<f64>,
) -> Result<()> {
    let n = x.len();
    ensure_len_eq("daxpy_f64", "y", y.len(), "x", n)?;

    if n == 0 {
        return Ok(());
    }

    let grid = blocks_for(n, BLOCK_SIZE);
    let grid = grid_u32("daxpy_f64", grid)?;
    let params = LaunchParams::new_1d(grid, BLOCK_SIZE_U32, stream.raw());

    check(unsafe { ffi::icffi_daxpy_f64(params, x.as_in(), y.as_out(), a) })
}

/// Scale vector: out = a * x
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `a` - Scalar multiplier
/// * `x` - Input vector
/// * `out` - Output vector
///
/// # Errors
///
/// Returns an error if the vectors have mismatched lengths or kernel launch fails.
#[track_caller]
pub fn scale_f32(
    stream: &Stream,
    a: f32,
    x: &DeviceBuffer<f32>,
    out: &mut DeviceBuffer<f32>,
) -> Result<()> {
    let n = x.len();
    ensure_len_eq("scale_f32", "out", out.len(), "x", n)?;

    if n == 0 {
        return Ok(());
    }

    let grid = blocks_for(n, BLOCK_SIZE);
    let grid = grid_u32("scale_f32", grid)?;
    let params = LaunchParams::new_1d(grid, BLOCK_SIZE_U32, stream.raw());

    check(unsafe { ffi::icffi_scale_f32(params, x.as_in(), out.as_out(), a) })
}

/// Parallel sum reduction (first pass).
///
/// Reduces input to per-block partial sums. For a complete reduction,
/// call this function multiple times until the output has a single element.
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `input` - Input vector
/// * `output` - Output vector (must have at least `blocks_needed(input.len())` elements)
///
/// # Returns
///
/// Returns the number of elements written to output.
///
/// # Errors
///
/// Returns an error if output buffer is too small or kernel launch fails.
#[track_caller]
pub fn reduce_sum_f32(
    stream: &Stream,
    input: &DeviceBuffer<f32>,
    output: &mut DeviceBuffer<f32>,
) -> Result<usize> {
    reduce_sum_f32_len(stream, input, input.len(), output)
}

/// Parallel max reduction (first pass).
///
/// Reduces input to per-block partial maxima.
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `input` - Input vector
/// * `output` - Output vector (must have at least `blocks_needed(input.len())` elements)
///
/// # Returns
///
/// Returns the number of elements written to output.
///
/// # Errors
///
/// Returns an error if output buffer is too small or kernel launch fails.
#[track_caller]
pub fn reduce_max_f32(
    stream: &Stream,
    input: &DeviceBuffer<f32>,
    output: &mut DeviceBuffer<f32>,
) -> Result<usize> {
    reduce_max_f32_len(stream, input, input.len(), output)
}

/// Returns the number of output elements needed for reduction.
#[inline]
#[must_use]
pub const fn reduction_output_size(input_len: usize) -> usize {
    let q = input_len / REDUCTION_ELEMENTS_PER_BLOCK;
    let r = input_len % REDUCTION_ELEMENTS_PER_BLOCK;
    if r == 0 {
        q
    } else {
        q + 1
    }
}

/// Computes the sum of all elements in the input vector.
///
/// This is a convenience function that handles multi-pass reduction internally.
/// It allocates temporary buffers from the CUDA memory pool and performs multiple
/// reduction passes until a single value remains.
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `input` - Input vector to reduce
///
/// # Returns
///
/// The sum of all elements.
///
/// # Errors
///
/// Returns an error if allocation fails or kernel launch fails.
///
/// # Memory Pool Usage
///
/// This function uses `zeros_async` for pool-based allocation and properly
/// returns buffers to the pool via `free_async` for optimal performance.
///
/// # Example
///
/// ```ignore
/// let stream = Stream::new()?;
/// let data = DeviceBuffer::from_slice_sync(&stream, &[1.0f32, 2.0, 3.0, 4.0])?;
///
/// let sum = reduce_sum_full(&stream, &data)?;
/// assert_eq!(sum, 10.0);
/// ```
#[track_caller]
pub fn reduce_sum_full(stream: &Stream, input: &DeviceBuffer<f32>) -> Result<f32> {
    if input.is_empty() {
        return Ok(0.0);
    }

    if input.len() == 1 {
        let result = input.to_vec(stream)?;
        return Ok(result[0]);
    }

    // Allocate primary working buffer from pool
    let mut current_len = input.len();
    let output_size = reduction_output_size(current_len);
    let mut buf_a = DeviceBuffer::<f32>::zeros_async(stream, output_size)?;

    // First pass: reduce input to partial sums
    current_len = reduce_sum_f32_len(stream, input, input.len(), &mut buf_a)?;

    // Additional passes until we have a single element
    // Use Option to track whether buf_b was allocated
    let mut buf_b: Option<DeviceBuffer<f32>> = None;

    while current_len > 1 {
        // Lazily allocate buf_b on first multi-pass iteration using Option::insert
        // which sets the value and returns a mutable reference in one operation.
        let b = match &mut buf_b {
            Some(b) => b,
            None => {
                let size = reduction_output_size(current_len);
                buf_b.insert(DeviceBuffer::<f32>::zeros_async(stream, size)?)
            }
        };

        current_len = reduce_sum_f32_len(stream, &buf_a, current_len, b)?;
        core::mem::swap(&mut buf_a, b);
    }

    // Copy final result to host (this synchronizes the stream)
    let result = buf_a.to_vec(stream)?;
    let value = result[0];

    // Return buffers to the memory pool
    // Note: Stream is already synchronized after to_vec, but free_async is
    // still correct and maintains pool hygiene for subsequent allocations.
    buf_a.free_async(stream)?;
    if let Some(b) = buf_b {
        b.free_async(stream)?;
    }

    Ok(value)
}

/// Computes the maximum of all elements in the input vector.
///
/// This is a convenience function that handles multi-pass reduction internally.
/// It allocates temporary buffers from the CUDA memory pool and performs multiple
/// reduction passes until a single value remains.
///
/// # Arguments
///
/// * `stream` - CUDA stream for the operation
/// * `input` - Input vector to reduce
///
/// # Returns
///
/// The maximum element value.
///
/// # Errors
///
/// Returns an error if allocation fails or kernel launch fails.
///
/// # Memory Pool Usage
///
/// This function uses `zeros_async` for pool-based allocation and properly
/// returns buffers to the pool via `free_async` for optimal performance.
///
/// # Example
///
/// ```ignore
/// let stream = Stream::new()?;
/// let data = DeviceBuffer::from_slice_sync(&stream, &[3.0f32, 1.0, 4.0, 1.5])?;
///
/// let max = reduce_max_full(&stream, &data)?;
/// assert_eq!(max, 4.0);
/// ```
#[track_caller]
pub fn reduce_max_full(stream: &Stream, input: &DeviceBuffer<f32>) -> Result<f32> {
    if input.is_empty() {
        return Ok(f32::NEG_INFINITY);
    }

    if input.len() == 1 {
        let result = input.to_vec(stream)?;
        return Ok(result[0]);
    }

    // Allocate primary working buffer from pool
    let mut current_len = input.len();
    let output_size = reduction_output_size(current_len);
    let mut buf_a = DeviceBuffer::<f32>::zeros_async(stream, output_size)?;

    // First pass: reduce input to partial maxima
    current_len = reduce_max_f32_len(stream, input, input.len(), &mut buf_a)?;

    // Additional passes until we have a single element
    // Use Option to track whether buf_b was allocated
    let mut buf_b: Option<DeviceBuffer<f32>> = None;

    while current_len > 1 {
        // Lazily allocate buf_b on first multi-pass iteration using Option::insert
        // which sets the value and returns a mutable reference in one operation.
        let b = match &mut buf_b {
            Some(b) => b,
            None => {
                let size = reduction_output_size(current_len);
                buf_b.insert(DeviceBuffer::<f32>::zeros_async(stream, size)?)
            }
        };

        current_len = reduce_max_f32_len(stream, &buf_a, current_len, b)?;
        core::mem::swap(&mut buf_a, b);
    }

    // Copy final result to host (this synchronizes the stream)
    let result = buf_a.to_vec(stream)?;
    let value = result[0];

    // Return buffers to the memory pool
    // Note: Stream is already synchronized after to_vec, but free_async is
    // still correct and maintains pool hygiene for subsequent allocations.
    buf_a.free_async(stream)?;
    if let Some(b) = buf_b {
        b.free_async(stream)?;
    }

    Ok(value)
}

/// Ensures the ABI asserts translation unit is linked.
///
/// Call this function once at startup to verify the ABI asserts compiled
/// successfully.
pub fn verify_abi_linked() {
    unsafe {
        ffi::icffi_abi_asserts_linked();
    }
}

#[cfg(test)]
mod lib_test;