介绍 CUDA Rust：用于编写 GPU 内核的两条轨道

cargo +nightly-2026-04-03 install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide

cargo oxide new vecadd_demo
cd vecadd_demo
cargo oxide doctor
cargo oxide run

use cuda_device::{kernel, launch_bounds, launch_contract, thread, DisjointSlice};
use cuda_host::cuda_module;
use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D};
 
// === DEVICE CODE - everything in here is compiled to PTX ===
// The macro also generates the host-side API used further down:
// `load`, `prepare_vecadd`, and the safe `vecadd` launch method.
#[cuda_module]
mod kernels {
    use super::*;
 
    #[kernel] // GPU entry point
    #[launch_bounds(256)] // max threads per block; lets the compiler budget registers
    #[launch_contract(domain = 1, block = (256, 1, 1))] // indexes in 1-D, 256-thread blocks
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let idx_raw = idx.get(); // the plain usize, for reading the inputs
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[idx_raw] + b[idx_raw];
        }
    }
}
 
fn main() -> Result<(), Box<dyn std::error::Error>> {
    // === HOST SETUP - device, stream, and buffers ===
    let ctx = CudaContext::new(0)?;
    let stream = ctx.default_stream();
 
    const N: usize = 1024;
    let a_host: Vec<f32> = (0..N).map(|i| i as f32).collect();
    let b_host: Vec<f32> = (0..N).map(|i| (i * 2) as f32).collect();
 
    let a_dev = DeviceBuffer::from_host(&stream, &a_host)?;
    let b_dev = DeviceBuffer::from_host(&stream, &b_host)?;
    let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N)?;
 
    // === LOAD, PREPARE, LAUNCH ===
    // SAFETY: this package owns the embedded device bundle produced for the
    // kernels module above.
    let module = unsafe { kernels::load(&ctx)? };
 
    // 4 blocks of 256 threads, 0 bytes of dynamic shared memory. `prepare_vecadd`
    // checks that against the contract above and against the live device limits.
    // The safe `vecadd` below takes that token where a raw config would go.
    let prepared = module.prepare_vecadd(LaunchConfig1D::new((N as u32).div_ceil(256), 256, 0))?;
    module.vecadd(&stream, &prepared, &a_dev, &b_dev, &mut c_dev)?;
 
    // === READ BACK AND VERIFY ===
    // Copies down and synchronizes, so the launch has finished by the time
    // `c_host` can be read.
    let c_host = c_dev.to_host_vec(&stream)?;
    let errors = (0..N)
        .filter(|&i| (c_host[i] - (a_host[i] + b_host[i])).abs() > 1e-5)
        .count();
 
    if errors == 0 {
        println!("PASSED: all {} elements correct", N);
    } else {
        eprintln!("FAILED: {} errors", errors);
        std::process::exit(1);
    }
    Ok(())
}

cargo new vecadd_demo
cd vecadd_demo
cargo add cutile

use cutile::prelude::*;
 
// The macro captures this module's AST into the host binary. The kernel is
// JIT-compiled through CUDA Tile IR the first time it is actually launched.
#[cutile::module]
mod kernel {
    use cutile::core::*;
 
    #[cutile::entry()]
    fn add<const B: i32>(
        // B is the tile width, a static dimension. A different B produces a
        // different specialization.
        z: &mut Tensor<f32, { [B] }>, // exclusive output, one sub-tensor of B elements
        x: &Tensor<f32, { [-1] }>,    // shared input; -1 is a dynamic dimension, resolved at launch
        y: &Tensor<f32, { [-1] }>,
    ) {
        // This body runs once per mut sub-tensor, as a single logical thread.
        // Tile kernels load tiles, not scalars, from x and y.
        let tx = load_tile_like(x, z); // the slice of x lining up with this sub-tensor of z
        let ty = load_tile_like(y, z);
        z.store(tx + ty); // elementwise across the whole tile
    }
}
 
fn main() -> Result<(), Error> {
    let device = Device::new(0)?;
    let stream = device.new_stream()?;
 
    // These are lazy. Nothing has touched the GPU yet.
    let x = api::ones::<f32>(&[1024]);
    let y = api::ones::<f32>(&[1024]);
 
    // Partitioning does three things at once: gives each tile exclusive
    // ownership of its own 128-element chunk, fixes the grid at 1024/128 = 8
    // tiles, and supplies B.
    let z = api::zeros::<f32>(&[1024]).partition([128]);
 
    let c: Vec<f32> = kernel::add(z, x, y) // takes ownership of all three tensors
        .first()                           // ...and returns them; pick the output back out
        .unpartition()                     // drop the host-side partition wrapper; no data moves
        .to_host_vec()                     // record the copy back
        .sync_on(&stream)?;                // and only now does any of it run
 
    let errors = c.iter().filter(|&&v| (v - 2.0).abs() > 1e-5).count();
    if errors == 0 {
        println!("PASSED: all {} elements correct", c.len());
    } else {
        eprintln!("FAILED: {errors} errors");
    }
    Ok(())
}

module.vecadd(&stream, &prepared, &c_dev, &b_dev, &mut c_dev)?;

let z = api::zeros::<f32>(&[1024]);
kernel::add(z.partition([128]), z, y)