Launching 3D Grid/Block Kernels from Safe Rust

Embedding & Dynamic Launching via cudarc

Once compiled, the PTX can be embedded into the host binary using include_str! or dynamically compiled at runtime using NVRTC (NVIDIA Runtime Compilation):

use cudarc::driver::CudaContext;

use cudarc::driver::{LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

const KERNEL_PTX: &str = include_str!("../kernels/vector_add.ptx");

pub fn execute_gpu_pipeline() -> Result<(), Box> {
    let ctx = CudaContext::new(0)?;
    let stream = ctx.default_stream();
    println!("Initialized CUDA Context on GPU: {}", ctx.name()?);

    // Load PTX module into GPU context
    let module = ctx.load_module(Ptx::from_src(KERNEL_PTX))?;
    let f = module.load_function("vector_add_kernel")?;
    // Kernel symbol resolved from loaded PTX module

    let n = 1_048_576; // 1M elements
    let a_dev = stream.clone_htod(&vec![2.0f32; n])?;
    let b_dev = stream.clone_htod(&vec![3.0f32; n])?;
    let mut c_dev = stream.alloc_zeros::(n)?;
    // Buffers allocated asynchronously on CudaStream
    // Zero-initialized output CudaSlice ready

    // Configure 3D Grid & Block Topology
    let threads_per_block = 256;
    let blocks = (n as u32 + threads_per_block - 1) / threads_per_block;
    let cfg = LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (threads_per_block, 1, 1),
        shared_mem_bytes: 0,
    };

    // Launch kernel asynchronously!
    let mut builder = stream.launch_builder(&f);
    builder.arg(&a_dev).arg(&b_dev).arg(&mut c_dev).arg(&n);
    unsafe { builder.launch(cfg) }?;

    // Copy result back to CPU
    let c_host = stream.clone_dtoh(&c_dev)?;
    assert_eq!(c_host[0], 5.0);
    println!("Successfully verified 1M GPU vector additions!");
    Ok(())
}