Compiling Rust Directly to PTX via nvptx64-nvidia-cuda

Writing GPU Kernels in Pure Rust

Historically, developers had to write .cu files and compile them with NVIDIA's nvcc compiler. With the LLVM nvptx64-nvidia-cuda target, you can write GPU kernels directly in pure Rust:

// In your kernel crate (target = "nvptx64-nvidia-cuda")
#![no_std]
#![feature(abi_ptx)]

#[no_mangle]
pub unsafe extern "ptx-kernel" fn vector_add_kernel(
    a: *const f32,
    b: *const f32,
    c: *mut f32,
    n: usize,
) {
    let idx = core::arch::nvptx::_thread_idx_x() as usize
        + core::arch::nvptx::_block_idx_x() as usize * core::arch::nvptx::_block_dim_x() as usize;

    if idx < n {
        *c.add(idx) = *a.add(idx) + *b.add(idx);
    }
}

Building the PTX Kernel

Compiling with cargo build --target nvptx64-nvidia-cuda --release outputs human-readable .ptx assembly, ready to be loaded by the host Rust runtime.