oxiproj-core 0.1.2

//! Boilerplate-reducing helpers for launching OxiProj's 1-D GPU kernels.

use oxicuda::{
    grid_size_for, CudaResult, DeviceBuffer, Kernel, KernelArgs, LaunchParams, PinnedBuffer, Stream,
};

/// Default threads-per-block for OxiProj's element-wise (one-thread-per-point)
/// kernels. 256 is a good occupancy default across Turing→Blackwell.
pub const BLOCK: u32 = 256;

/// Launch `kernel` over `n` threads (one per coordinate / point) with a 1-D grid
/// of [`BLOCK`]-thread blocks, then block until the stream completes.
///
/// `args` is the kernel argument tuple (device pointers as `u64`, the element
/// count as `u32`, then any `f64` scalars), in the order the kernel's `.entry`
/// parameters declare them.
///
/// # Errors
///
/// Propagates any launch or synchronisation error from the driver.
pub fn launch_1d<A: KernelArgs>(
    kernel: &Kernel,
    stream: &Stream,
    n: u32,
    args: &A,
) -> CudaResult<()> {
    let params = LaunchParams::new(grid_size_for(n, BLOCK), BLOCK);
    kernel.launch(&params, stream, args)?;
    stream.synchronize()
}

/// Allocate a device buffer of `src.len()` elements and enqueue the host→device
/// copy of `src` **onto `stream`**, returning the buffer to feed a kernel.
///
/// Use this — not [`DeviceBuffer::from_host`] — to stage a kernel's inputs.
/// `from_host` issues the *synchronous* `cuMemcpyHtoD`, which for pageable host
/// memory (an ordinary `&[T]`) merely stages into the driver's internal buffer
/// and then runs the DMA to the device on the **default (NULL) stream**. Every
/// OxiProj stream is created `CU_STREAM_NON_BLOCKING`, so it does **not**
/// implicitly synchronise with the NULL stream: a kernel launched on `stream`
/// would race that still-in-flight DMA and read not-yet-copied (zero) device
/// memory — yielding `f(0)` for whichever elements the DMA had not yet reached.
/// Enqueuing the copy on the *same* `stream` the kernel uses makes the kernel
/// correctly stream-ordered after it.
///
/// The copy is enqueued, not awaited: `src` must remain valid until `stream` is
/// synchronised (the [`launch_1d`] that consumes the returned buffer does this).
///
/// # Errors
///
/// Propagates any allocation or host→device copy error from the driver
/// (including `CudaError::InvalidValue` when `src` is empty).
pub fn upload<T: Copy>(src: &[T], stream: &Stream) -> CudaResult<DeviceBuffer<T>> {
    let mut buf = DeviceBuffer::<T>::alloc(src.len())?;
    buf.copy_from_host_async(src, stream)?;
    Ok(buf)
}

/// Run an **in-place** element-wise kernel over `data` (a flat `f64` buffer)
/// using page-locked (pinned) host staging for fast, consistent DMA.
///
/// Pageable host↔device copies fall back to a slow staged path inside the driver;
/// pinned memory enables true DMA, which on bandwidth-bound element-wise kernels
/// (e.g. coordinate reprojection) is several times faster and far less variable.
/// The host→device copy, kernel launch, and device→host copy are all enqueued on
/// one stream and awaited with a single synchronisation.
///
/// `n_threads` is the number of threads to launch (e.g. coordinate *pairs*, not
/// the `f64` element count). `make_args` builds the kernel argument tuple from the
/// device buffer pointer; its first parameter must be that pointer.
///
/// # Errors
///
/// Propagates any allocation, copy, launch, or synchronisation error.
pub fn run_inplace_pinned<A, F>(
    kernel: &Kernel,
    stream: &Stream,
    data: &mut [f64],
    n_threads: u32,
    make_args: F,
) -> CudaResult<()>
where
    A: KernelArgs,
    F: FnOnce(u64) -> A,
{
    let mut staged = PinnedBuffer::<f64>::from_slice(data)?;
    let mut dev = DeviceBuffer::<f64>::alloc(data.len())?;
    dev.copy_from_host_async(staged.as_slice(), stream)?;
    let params = LaunchParams::new(grid_size_for(n_threads, BLOCK), BLOCK);
    kernel.launch(&params, stream, &make_args(dev.as_device_ptr()))?;
    dev.copy_to_host_async(staged.as_mut_slice(), stream)?;
    stream.synchronize()?;
    data.copy_from_slice(staged.as_slice());
    Ok(())
}