mircuda 0.1.0

Native, explicit-stream Rust gateway to NVIDIA CUDA
#![cfg(target_os = "linux")]

use mircuda::{
    CompileOptions, Compiler, DeviceBuffer, Driver, KernelNode, LaunchConfig, Stream, TypedKernel,
    cuda_export, cuda_kernel,
};

cuda_export!(Increment = "increment"(value: &mut DeviceBuffer<u32>, delta: u32));

struct Sequence<'a> {
    kernel: &'a TypedKernel<Increment>,
    stream: &'a Stream,
    value: &'a mut DeviceBuffer<u32>,
    delta: u32,
    node: Option<KernelNode<Increment>>,
}

fn capture_increment(sequence: &mut Sequence<'_>) -> mircuda::Result<()> {
    sequence.node = Some(sequence.kernel.launch_captured(
        sequence.stream,
        LaunchConfig::for_elements(1, 1)?,
        (&mut *sequence.value, sequence.delta),
    )?);
    Ok(())
}

const fn triple<'borrow>(
    sequence: &'borrow mut Sequence<'_>,
) -> (&'borrow mut DeviceBuffer<u32>, u32) {
    sequence.delta = 3;
    (&mut *sequence.value, sequence.delta)
}

#[test]
fn captures_and_replays_on_the_same_stream() -> mircuda::Result<()> {
    let driver = Driver::initialize()?;
    let device = driver.devices()?.into_iter().next().ok_or(mircuda::Error::InvalidLaunch)?;
    let context = driver.create_context(device)?;
    let stream = context.create_stream()?;
    let pool = context.default_memory_pool()?;
    let compiler = Compiler::new(context.clone())?;
    let source = cuda_kernel!(
        r#"extern "C" __global__ void increment(unsigned int* value, unsigned int delta) {
            if (threadIdx.x == 0) value[0] += delta;
        }"#
    );
    let module = compiler.compile(source, &CompileOptions::default())?;
    let kernel = module.kernel::<Increment>()?;
    let mut value = pool.allocate_zeroed::<u32>(&stream, 1)?;
    {
        let sequence = Sequence {
            kernel: &kernel,
            stream: &stream,
            value: &mut value,
            delta: 1,
            node: None,
        };
        let mut graph = stream.capture(sequence, capture_increment)?;
        let node = graph.resources().node.ok_or(mircuda::Error::InvalidLaunch)?;
        graph.update_kernel(&node, &kernel, LaunchConfig::for_elements(1, 1)?, triple)?;
        graph.launch(&stream)?;
        graph.launch(&stream)?;
    }
    let mut host = context.allocate_pinned::<u32>(1)?;
    stream.copy_to_host(&value, &mut host)?;
    assert_eq!(host.to_vec()?, [6]);
    Ok(())
}