引入 CUDA Rust：编写 GPU 内核的两种途径

# CUDA Rust 入门：编写 GPU 内核的两条路径 来源：https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels *2026 年 9 月，NVIDIA 宣布将深度投入 Rust 语言原生 GPU 编程。CUDA C++ 与 CUDA Python 是成熟的企业级工具链，而 NVIDIA 将持续推动 CUDA Rust 在 2027 年及以后的发展与成熟* AI 的系统层涵盖推理引擎、服务基础设施、驱动程序及智能体运行时，且随着模型与技术的演进不断迭代。越来越多的系统层代码采用 Rust 编写，它能在编译时捕获整类错误且不牺牲性能。NVIDIA 也顺应这一趋势——Nova Linux 驱动程序即采用 Rust 编写。NVIDIA Dynamo（https://www.nvidia.com/en-us/ai/dynamo/）的核心基于 Rust 构建，NVTX 已提供 Rust 绑定。唯独 GPU 内核是个例外：你可以从 Rust 启动内核，但内核本身通常仍需用其他语言编写。NVIDIA CUDA Rust 填补了这一空白，使得 GPU 内核能直接用 Rust 编写并原生编译为 PTX，而非依赖外部代码的封装。 使用 Rust 有两条路径，与 CUDA 自身的双轨模式对应： **SIMT** 是你在 CUDA C++ 或 numba-cuda（https://nvidia.github.io/numba-cuda/）中已熟悉的模型——你定义单个线程的行为，然后启动数千个线程并行执行。 **Tile** 则是较新的编程模型，同样支持 C++（https://docs.nvidia.com/cuda/cuda-tile-cpp-api-reference/）和 Python（https://docs.nvidia.com/cuda/cutile-python/）。所有前端都允许你描述单个数据块的操作，由 Tile IR 编译器（https://docs.nvidia.com/cuda/tile-ir/latest/index.html）完成后续工作。选择构建基础时，建议优先考虑 Tile——编译器会决定数据块如何映射到不同架构，因此源代码无需包含架构特定的选择；当你需要更精细的控制或手动管理内存与线程时，再回退到 SIMT。 选择哪种语言与选择哪种模型是两个独立问题。请根据现有技术栈选择最合适的 CUDA 接口。以下两个项目适用于 Rust 技术栈。我们计划支持跨语言互操作，因此这种选择不会将你局限在单一方案中。 以下是在两条路径上实现相同功能的内核示例：对 1,024 个浮点数执行逐元素加法。两者都是完整程序，均可运行并输出相同结果，可对比阅读以观察差异。 ## SIMT 路径：cuda-oxide https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels#the_simt_track_cuda-oxide cuda-oxide（https://github.com/NVlabs/cuda-oxide）是一个定制的 `rustc` 代码生成后端。它拦截编译过程，将 `#[kernel]` 函数通过 Rust MIR、社区 Pliron（https://github.com/pliron-org/pliron）IR 框架和 LLVM IR 传递到 PTX，其余部分则交给标准后端处理。基于 Pliron 的 GPU 方言由我们开发，所有方言和转换均在 Rust 中完成，直到标准 LLVM 后端接手。 你需要 Linux 系统、计算能力 8.0 或更高的 GPU、CUDA 工具包（12.x 或更新版本）、带有 libclang 头文件的 clang，以及固定的 nightly 工具链。`cargo oxide doctor` 会检查所有依赖项，包括可选的系统 LLVM。 首先安装 `cargo-oxide`，这是用于驱动构建的 Cargo 子命令： ```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 ``` 首次运行 `cargo oxide build` 会构建代码生成后端，可能需要一些时间。后续运行将利用缓存。程序将输出 `PASSED: all 1024 elements correct`。以下是完整程序（即 `cargo oxide new` 生成的代码，此处添加了注释）： ```rust use cuda_device::{kernel, launch_bounds, launch_contract, thread, DisjointSlice}; use cuda_host::cuda_module; use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D}; // === 设备代码 - 此处所有内容均编译为 PTX === // 宏同时生成主机端 API，用于后续调用： // `load`、`prepare_vecadd` 和安全的 `vecadd` 启动方法。 #[cuda_module] mod kernels { use super::*; #[kernel] // GPU 入口点 #[launch_bounds(256)] // 每块最大线程数；帮助编译器分配寄存器 #[launch_contract(domain = 1, block = (256, 1, 1))] // 1 维索引，每块 256 线程 pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice) { let idx = thread::index_1d(); let idx_raw = idx.get(); // 原始 usize 值，用于读取输入 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>> { // === 主机设置 - 设备、流和缓冲区 === 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)?; // === 加载、准备、启动 === // SAFETY: 此包拥有为上述 kernels 模块生成的嵌入式设备包。 let module = unsafe { kernels::load(&ctx)? }; // 4 个块，每块 256 线程，动态共享内存 0 字节。`prepare_vecadd` // 会根据上述约定和实时设备限制进行验证。 // 下方安全的 `vecadd` 方法需要此令牌作为原始配置的替代。 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)?; // === 读取并验证 === // 复制回主机并同步，确保在读取 `c_host` 时启动已完成。 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(()) } ``` 主机与设备代码共存于同一文件，通过单一命令构建，无需单独的内核 crate。首先关注内核签名，因为它承载了完整的安全性论证： - `a` 和 `b` 是普通共享切片，所有线程均可读取。 - `c` 是 `DisjointSlice` 类型，确保每个线程仅能访问自身对应的元素。之所以需要该类型，是因为 `&mut [f32]` 的形态不适合此任务——若所有线程共享同一可变引用，Rust 会正确拒绝。`DisjointSlice` 将单一可变借用拆分为按线程分配的独立部分。 `thread::index_1d()` 返回索引类型而非裸整数，`c.get_mut(idx)` 仅接受该类型。返回值为 `Option`，因此边界外情况是显式分支处理而非后续的内存错误。 启动过程经过验证而非盲目信任：`#[launch_contract]` 声明该内核采用 1 维索引和 256 线程块。`prepare_vecadd` 会根据该声明和实时设备限制验证你的 `LaunchConfig1D`，并返回安全的 `vecadd` 方法所需的证明。没有约定的内核仅暴露原始的 unsafe 启动方法，因为裸 `LaunchConfig` 无法说明其启动的内核特性。 ## Tile 路径：cutile-rs https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels#the_tile_track_cutile-rs cutile-rs（https://github.com/NVlabs/cutile-rs）工作在更高层级——你在数据块上执行计算而非标量。每个块作为单一逻辑线程运行内核主体，处理数据的一个子张量，由编译器决定由多少真实 GPU 线程支持。 `#[cutile::module]` 宏会将内核的 AST 嵌入主机二进制文件，并在内核首次需要时通过 CUDA Tile IR（NVIDIA 的块级编译器 IR）进行 JIT 编译。 要求比 SIMT 路径更轻量：需要计算能力 8.0 或更高的 GPU、CUDA 13.3、stable Rust 1.89 或更高版本及 Linux，无需 nightly 工具链和自备 LLVM。 `cutile` 已发布，无需克隆仓库： ```cargo new vecadd_demo cd vecadd_demo cargo add cutile ``` 以下是相同的逐元素加法实现，针对 Tile 编写。将其粘贴到 `src/main.rs` 并运行 `cargo run`： ```rust use cutile::prelude::*; // 宏将此模块的 AST 捕获到主机二进制文件中。内核首次实际启动时 // 通过 CUDA Tile IR 进行 JIT 编译。 #[cutile::module] mod kernel { use cutile::core::*; #[cutile::entry()] fn add( // B 是数据块宽度，为静态维度。不同的 B 会生成不同的特化版本。 z: &mut Tensor, // 独占输出，B 个元素的子张量 x: &Tensor, // 共享输入；-1 是动态维度，启动时解析 y: &Tensor, ) { // 此主体每次对 mut 子张量执行一次，作为单一逻辑线程运行。 // 块内核从 x 和 y 加载数据块而非标量。 let tx = load_tile_like(x, z); // 与 z 子张量对齐的 x 切片 let ty = load_tile_like(y, z); z.store(tx + ty); // 跨整个块执行逐元素操作 } } fn main() -> Result<(), cutile::Error> { let device = Device::new(0)?; let stream = device.new_stream()?; // 这些是惰性的。目前尚未接触 GPU。 let x = api::ones(&[1024]); let y = api::ones(&[1024]); // 分区同时完成三件事：为每个块分配其独有的 128 元素块独占所有权， // 确定网格大小为 1024/128 = 8 个块，并提供 B 值。 let z = api::zeros(&[1024]).partition([128]); let c: Vec<f32> = kernel::add(z, x, y) // 获取所有三个张量的所有权 .first() // ...然后返回它们；提取输出 .unpartition() // 丢弃主机端分区包装器；无数据移动 .to_host_vec() // 记录回拷操作 .sync_on(&stream)?; // 此时才实际执行 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(()) } ``` 输出：`PASSED: all 1024 elements correct` Tile 路径在稳定版 Rust 上达成相同结果，其签名同样体现了安全性论证。此次无需 `DisjointSlice`——主机端分区仅用于可变张量，它为每个块分配一个其他块不可重叠的可写子张量。这种独占性正是 `&mut` 已经保证的特性。 输入形状中的 `-1` 是哨兵值而非具体尺寸，该维度在启动时从张量读取，因此形状变化无需重新编译。主机端的关键行是 `.partition([128])`，它同时承担三项职责： 1. 使独占性生效：每个块拥有其 128 元素块，其他块无法访问。 2. 确定启动几何结构：1024 除以 128 得到 8 个块的网格，网格大小由分区推导而非单独计算并与内核索引对比验证。 3. 提供 `B` 值：该值在调用点未明确写出，因为启动器从分区中读取块宽度。这就是为什么 `&mut` 输出必须先分区才能传递。 观察启动返回值：主机调用的 `add` 是宏生成的启动器而非上方的设备函数。它获取所有三个张量的所有权，在 GPU 完成后将它们作为元组返回——这就是 `.first()` 的作用：从元组中提取输出。在 `.sync_on(&stream)` 之前不会执行任何操作，之前的所有内容都是惰性描述，仅记录而非提交。这包括 `ones`、`zeros`、内核调用甚至回拷操作。整个程序是一个包含单一同步点的调用链。 ## 编译器能捕获什么 https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels#what_the_compiler_catches 两个内核对内存做出相同声明：输入是共享的，输出仅属于一个写入者。它们仅在实现层级和是否需要专用类型上存在差异。这很重要，因为数千个线程以无保证的顺序访问相同缓冲区——当两个线程命中同一地址且其中一个在写入时，执行顺序将决定结果。这类 bug 很难按需复现，常在测试中通过却在生产环境中失败。 将 SIMT 内核的输出缓冲区同时作为输入传递时无法编译（无论该内核是否会产生竞争）： ```rust module.vecadd(&stream, &prepared, &c_dev, &b_dev, &mut c_dev)?; ``` 错误信息：`error[E0502]: cannot borrow 'c_dev' as mutable because it is also borrowed as immutable` Tile 侧的相同别名错误同样无法编译： ```rust let z = api::zeros(&[1024]); kernel::add(z.partition([128]), z, y) ``` 错误信息：`error[E0382]: use of moved value: 'z'` 两个示例都在编译时捕获了经典的别名错误，只是划分界限的位置不同。cuda-oxide 在每次启动调用时检查，而 cutile-rs 的所有权机制跨越了启动边界——这是更强的保证。Tile 不提供可能出错的共享内存或线程索引机制，因为编译器完全掌控这些。块是单一逻辑线程，因此不存在竞争风险——这使其天生安全，但也意味着你必须放弃这些控制能力。SIMT 保留了这种控制，目前其共享内存路径仍需 `unsafe`。共享内存是快速 SIMT 内核的基石，使该路径安全化是正在进行的工作。 ## 项目现状 https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels#where_the_projects_stand 两个项目均处于早期阶段，尚未达到生产就绪状态。cuda-oxide 是早期 alpha 版，cutile-rs 进度更快，已发布到 crates.io 并在 NVIDIA 外部使用——包括 HuggingFace 的 Grout（https://github.com/huggingface/grout）推理引擎和 mistral.rs（https://github.com/EricLBuehler/mistral.rs）。功能覆盖尚不完整，API 可能变动。遇到问题时请告知我们。 Cargo 和 crates 建立了“易于上手”的预期，而 GPU 编程历史上恰恰相反。缩小这种差距是工作的一部分。SIMT 路径目前仍需要固定 nightly 工具链，我们正努力摆脱这种依赖。 Rust 在 GPU 上的应用并非新概念。我们的工作之前已有该领域的优秀实践并持续并行发展。cuda-oxide 文档（https://nvlabs.github.io/cuda-oxide/）的生态系统附录（https://nvlabs.github.io/cuda-oxide/appendix/ecosystem.html）绘制了我们与 Rust-GPU、rust-cuda、CubeCL 等项目的相对位置图谱。我们与 rust-cuda 维护者持续合作，共同推动项目成熟。真正创新的是我们投入的工程能力以及清晰的发展路线图。 ## 当前可尝试的内容 https://developer.nvidia.com/blog/introducing-cuda-rust-two-tracks-for-writing-gpu-kernels#what_you_can_do_today - **运行 SIMT 示例**：在 cuda-oxide（https://github.com/NVlabs/cuda-oxide）中使用 `cargo oxide new` 和 `cargo oxide run`。 - **运行 Tile 示例**：克隆 cutile-rs（https://github.com/NVlabs/cutile-rs），执行 `cargo run -p cutile-examples --example hello_world`。 - **阅读文档**：cuda-oxide 手册（https://nvlabs.github.io/cuda-oxide/）和 cuTile Rust 文档（https://nvlabs.github.io/cutile-rs/main/）。 - **阅读论文**：《GPU 上的无畏并发》（https://arxiv.org/abs/2606.15991）。 - **提交问题**：在 cuda-oxide（https://github.com/NVlabs/cuda-oxide/issues）或 cutile-rs（https://github.com/NVlabs/cutile-rs/issues）反馈遇到的问题。 - **加入**...