rust-cuda

The Rust CUDA Project aims to make Rust a tier-1 language for NVIDIA GPU computing. It provides a specialized toolchain to compile Rust to optimized PTX code and provides high-level Rust wrappers for existing CUDA libraries. The project is composed of several specialized crates targeting different parts of the CUDA ecosystem:

• rustc_codegen_nvvm: A rustc backend that targets NVVM IR to generate highly optimized PTX code for the CUDA Driver API.
• cuda_std: A library of GPU-side utilities (e.g., thread index queries, memory allocation, warp intrinsics) designed to make writing GPU kernels cleaner and more reliable.
• cust: A high-level CPU-side wrapper for the CUDA Driver API. It handles kernel launching, GPU memory allocation, and device queries using Rust features like RAII and Result.
• cudnn: GPU-accelerated primitives for deep neural networks.
• gpu_rand: GPU-friendly random number generation (currently implementing xoroshiro RNGs).
• optix: CPU-side hardware raytracing and denoising via the CUDA OptiX library.

OptiX uses several specialized program types to handle different stages of the ray tracing pipeline. Memorizing these shorthand terms is helpful as they are used throughout the API:

• RG (Ray generation): The entry point; responsible for creating and tracing rays.
• IS (Intersection): Provides intersections for custom, user-defined primitives.
• AH (Any-hit): Runs during traversal for each potential intersection; reports if an intersection is valid and if traversal should stop.
• CH (Closest-hit): Runs only for the closest intersection found; used for material shading and property interpolation.
• MS (Miss): Runs when a ray exits the scene without hitting anything.
• EX (Exception): Handles error conditions like stack overflows.
• DC (Direct callable): Can be called manually from another program but cannot call optixTrace.
• CC (Continuation callable): Can be called manually and is allowed to continue ray traversal.

The Rust CUDA project replaces the traditional NVCC compiler with a custom rustc backend. The workflow involves compiling GPU kernel code into NVVM IR, which is then converted to PTX via libNVVM. This PTX is typically embedded into a host binary using include_str!(). The host binary then uses the cust crate to interact with the CUDA Driver API, which performs JIT compilation to SASS (GPU machine code) for execution.

Key components of the pipeline:

• rustc_codegen_nvvm: A custom rustc backend that compiles GPU kernel crates to NVVM IR (LLVM bitcode).
• cuda_std: The GPU-side standard library providing thread indexing, shared memory, and intrinsics.
• cuda_builder: A build-script helper used in a host crate's build.rs to drive rustc_codegen_nvvm and produce the .ptx file.
• cust: A safe Rust wrapper around the CUDA Driver API used by the host to load modules, allocate memory, launch kernels, and synchronize results.

+---------------------------------------------------------------------+
|                        Rust CUDA Pipeline                           |
|                                                                     |
|  Host code (.rs)           GPU kernel code (.rs)                    |
|       |                          |                                  |
|       |                   rustc_codegen_nvvm                        |
|       |                   (custom rustc backend)                    |
|       |                          |                                  |
|       |                     NVVM IR (.bc)                           |
|       |                          |                                  |
|       |                      libNVVM                               |
|       |                          |                                  |
|       |                      PTX (.ptx)  <-- embedded via           |
|       |                          |           include_str!()         |
|       v                          v                                  |
|  Host binary ---- cust ------> Driver API                           |
|                  (Rust)         (CUDA)                              |
|                                  |                                  |
|                              JIT compile                            |
|                                  |                                  |
|                              SASS (GPU machine code)                |
|                                  |                                  |
|                              GPU execution                          |
+---------------------------------------------------------------------+

GPU computing utilizes the parallel nature of GPUs for tasks like fluid simulation, AI model training, and physically based rendering. This project leverages CUDA and Rust to provide a high-performance, safe environment for GPU programming.

Why use CUDA?

• Control: Deep control over kernel dispatch and memory management.
• Ecosystem: Access to libraries like cuRAND, cuBLAS, libNVVM, and OptiX.
• Performance: Unmatched performance due to its focus on computing.
• Note: CUDA is limited to NVIDIA GPUs.

Why use Rust with CUDA?

• Safety: Applies Rust's safety guarantees to traditionally unsafe GPU tasks (e.g., managing shared memory, thread block layouts, and data indexing).
• RAII: The cust library uses RAII (via Drop implementations) to automatically manage memory freeing and handle returning, reducing manual memory management errors.
• Error Handling: Instead of using unreliable C-style macros (like CUDA_SAFE_CALL), this project uses Rust's Result type. This allows for safe error propagation using the ? operator or explicit unwrapping.

To achieve optimal results with the NVIDIA AI Denoiser, follow these integration guidelines:

• Avoid Pre-denoising Post-processing: Do not apply image filters (like blurring or sharpening) or reconstruction filters (like box, triangle, or Gaussian filters) to a noisy image before passing it to the denoiser. Custom post-processing can smear high-frequency noise across multiple pixels, making it harder for the deep learning model to identify and remove.
• Order of Operations: Always perform post-processing operations after the denoising process is complete.
• Reconstruction Filters: If using reconstruction filters, implement them using filter importance-sampling.
• Color Space: The input image pixel color space should ideally match the color space used during the denoiser's training. The included general-purpose model was trained using images output directly as HDR data. While slight variations (e.g., substituting sRGB with a simple gamma curve) are generally acceptable, matching the training data characteristics is recommended.

To implement a ray tracing system using NVIDIA OptiX 7, follow these four steps:

1. Create acceleration structures: Build one or more acceleration structures over geometry meshes and instances in your scene.
2. Create a program pipeline: Define a pipeline containing all programs (shaders) that will be invoked during a ray tracing launch.
3. Create a shader binding table (SBT): Build a table that includes references to your programs and their parameters. Ensure the data layout matches the implicit shader binding table record selection of the instances and geometries in your acceleration structures.
4. Launch a device-side kernel: Execute a kernel that invokes a Ray Generation (RG) program. This kernel uses multiple threads calling optixTrace to initiate traversal and the execution of other programs.

Note: Ray tracing work can be interleaved with other CUDA work, but the application is responsible for coordinating all GPU work as OptiX 7 does not perform internal synchronization.

For manual memory management and buffer reuse, use the unsafe functions accel_compute_memory_usage and accel_build. This approach allows you to allocate your own output and temporary buffers, which is more efficient for repeated builds.

Workflow:

1. Call accel_compute_memory_usage to determine the required sizes for output_size_in_bytes, temp_size_in_bytes, and temp_update_size_in_bytes.
2. Allocate device memory for these buffers. Note: Pointers to these buffers must be aligned to a 128-byte boundary.
3. Call accel_build using the allocated buffers.

The build is asynchronous on the device; you must use stream.synchronize() or CUDA events to ensure the build is complete before using the resulting traversable handle.

// ... setup ctx, stream, vertices, indices ...

let buf_vertex = cu::DeviceBuffer::from_slice(&vertices)?;
let buf_indices = cu::DeviceBuffer::from_slice(&indices)?;
let geometry_flags = ox::GeometryFlags::None;

let build_inputs = [ox::IndexedTriangleArray::new(
    &[&buf_vertex],
    &buf_indices,
    &[geometry_flags]
)];

let accel_options = ox::AccelBuildOptions::new(
    ox::BuildFlags::ALLOW_COMPACTION,
    ox::BuildOperation::Build
);

// Get the storage requirements
let sizes = accel_compute_memory_usage(ctx, accel_options, build_inputs)?;

// Allocate temporary and output buffers (must be 128-byte aligned)
let mut output_buffer = unsafe { DeviceBuffer::<u8>::uninitialized(sizes.output_size_in_bytes)? };
let mut temp_buffer = unsafe { DeviceBuffer::<u8>::uninitialized(sizes.temp_size_in_bytes)? };

// Build the accel
let hnd = unsafe {
    accel_build(
        ctx,
        stream,
        accel_options,
        build_inputs,
        &mut temp_buffer,
        &mut output_buffer,
        &mut properties,
    )?
};

stream.synchronize()?;