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::{CudaDevice, LaunchAsync, LaunchConfig};

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

pub fn execute_gpu_pipeline() -> Result<(), Box> {
    let dev = CudaDevice::new(0)?;

    // Load PTX module into GPU context
    dev.load_ptx(KERNEL_PTX.into(), "vector_add_module", &["vector_add_kernel"])?;
    let f = dev.get_func("vector_add_module", "vector_add_kernel").unwrap();

    let n = 1_048_576; // 1M elements
    let a_dev = dev.htod_copy(vec![2.0f32; n])?;
    let b_dev = dev.htod_copy(vec![3.0f32; n])?;
    let mut c_dev = dev.alloc_zeros::(n)?;

    // 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!
    unsafe { f.launch(cfg, (&a_dev, &b_dev, &mut c_dev, n)) }?;

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