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)]

#[unsafe(no_mangle)]
pub unsafe extern "ptx-kernel" fn vector_add_kernel(
    a: *const f32,
    b: *const f32,
    c: *mut f32,
    n: usize,
) {
    // Rust 2024 unsafe_op_in_unsafe_fn: explicit inner unsafe {} required
    unsafe {
        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.

32-Thread SIMT Warps & Shared Memory Bank Conflicts

NVIDIA Streaming Multiprocessors (SMs) schedule threads in hardware lockstep groups of 32 threads called a Warp. On-chip Shared Memory (SMEM) is divided into 32 four-byte banks where bank index equals (byte_addr / 4) % 32. If threads in a warp access distinct addresses mapping to the same bank (such as reading down a column of a [f32; 32] matrix with stride 32), the SM serializes the access into a 32-way bank conflict. Padding rows by +1 element ([[f32; 33]; 32]) or applying XOR swizzling eliminates bank conflicts completely.