GPU Offload in Rust: Portable, Safe, and Fast
Manuel S. Drehwald, Marcelo Domínguez, Kevin Sala, Alán Aspuru-Guzik, Johannes Doerfert
cs.PL
2026-08-14
多伦多大学与 LLNL 把跨厂商 GPU 卸载框架直接做进 rustc 与 LLVM,靠类型系统自动推导数据传输,RAJAPerf 内核时间追平手写 CUDA/HIP,整轮跑分最快快 32%、最慢慢 46%。
HPC 和科学计算仍然被 C、C++、Fortran 加 CUDA/HIP 统治。这些语言没有安全边界,显存拷贝多少字节、哪个缓冲区归哪个线程,全靠程序员自己记。Rust 本来能在编译期兜住内存错误和数据竞争,但真要把 Rust 代码放到 GPU 上跑,过去的路都不好走:rust-gpu 走 SPIR-V 路线,不支持一般指针;rust-cuda 所有内核都得写 unsafe;NVIDIA 新推的 cuda-oxide 安全性做得不错,却只绑自家硬件。安全、可移植、高性能三样,此前最多同时拿两样。
这篇的赌注是:把 GPU 卸载直接做进 rustc 主线和 LLVM 的 Offload 基础设施,而不是做成外挂 DSL。作者是多伦多大学、LLNL 与西班牙 URJC 的团队,工作受 Rust Foundation 资助,Rust 语言团队有人在里面推动。
整个框架围绕三个接口展开,按「编译器管多少」排序:
这套设计之所以可行,靠的是 Rust 类型的红利:可变引用在 LLVM IR 层天然带 noalias 属性,不需要手写 restrict 注解,给后端优化提供了很好的信息底子。
内核本身的安全问题用「分区策略」解。GPU 上多线程共享 &mut 切片是未定义行为,传统做法是全退回裸指针。这篇的做法是把「哪个线程能碰哪些元素」从内核逻辑里拆出来:每个 Region 绑定一个 PartitioningStrategy,由它保证线程间访问区域互不相交,内核代码就能用普通安全 Rust 写。分区策略本身是 unsafe trait,但都是纯 Rust 实现,用户可以自己写新的。
数据传输方向同样从类型推导出来:&T 和 const T 生成单向 MapTo,&mut T 和 mut T 生成双向 MapToFrom,64 位以内的标量按值直传,绕开指针间接寻址。OpenMP 里手写的 data-mapping pragma,这里由 MIR 层的分析自动生成。
工具链是三趟编译:第一趟收集主机侧内核单态化元数据,第二趟用 GPU 目标编译设备位码并打包,第三趟编译主机侧并把这个 fat binary 嵌进去。之所以不用 cuda-oxide 那样的单趟方案,是因为 Rust 的 #[cfg(targetarch)] 会按目标挑实现,单趟意味着设备代码里可能残留 x86 内联汇编,作者认为把这类目标特定代码翻译成 GPU IR 不可行。代价是跨趟要显式传信息,他们用一个 monomorphization 元数据查询解决了单态化断裂。
把 RAJAPerf 的 13 个内核移植成纯 Rust,在 AMD MI250X 和 NVIDIA H100 上跑,对比 RAJA 框架的手写 CUDA/HIP 后端(rustc 基于 LLVM 23.1.0-rc1):
| 维度 | Rust | 对照(CUDA/HIP) |
| 整轮跑分,MI250X | 最快快 32% | 最慢时 Rust 慢 43% |
| 整轮跑分,H100 | 最快快 11% | FIR/LTIMES 上慢 44%/46% |
| H100 H2D 传输 | 53 次、423 MB | RAJA 55 次、468 MB |
| H100 D2H 传输 | 9 次、69 MB | RAJA 9 次、99 MB |
| H100 传输总耗时 | 46 ms | RAJA 16 ms |
| RTX 2070 寄存器均值 | 33 个 | RAJA-CUDA 28 个 |
几个值得盯的点。裸内核时间与 RAJA 基本持平,差距集中在 FIR/LTIMES 这类只有几条指令的小内核,作者归因于三个编译器的展开决策不同,不是框架税。更扎眼的是传输数据量更少、耗时反而近三倍(46 ms 对 16 ms),作者怀疑是内存种类与异步传输的差异,还没有定论。易用性上还有一个更大的坑:接口 A 不做优化直接跑,可以比显式接口慢 400 倍以上,所以他们在 LLVM OpenMP-opt 里原型化了传输预取、循环外提和中间传输消除,声称足以让易用接口在 RAJAPerf 上追平显式接口,这部分是原型,未进评测表。
浮点上有个 Rust 特有的取舍:C++ 圈习惯开 fast-math,Rust 不让开,因为 nnan/ninf 假设会在安全代码里触发 UB。实验性的 algebraic float 提供了大部分优化机会且不含这两个假设,在 RTX A2000 上让 FIR 提速 2 倍、三个内核提速约 20%,在 MI250X 上没有显著收益。
这是把 Rust GPU 编程从「社区外挂」推进到「编译器主线基础设施」的一步。NVIDIA 同期在做 cuda-oxide,AMD 和 Intel 用户却不在它的射程内;这篇站在 LLVM Offload 上,现在同时出 NVIDIA 和 AMD 的原生码,Intel 目标等 LLVM 上游成熟即可接入。对 HPC 社区,这是 Fortran/C++ 之外一条带编译期内存安全的路;对普通 Rust 开发者,意味着以后可能不需要 unsafe 块就能写内核。前提是它能真合并进上游。
作者自己摆出来的:同一份代码两趟编译后,主机与设备端的类型布局和 ABI 可能分叉,他们在基础切片类型上就撞到了实况,x8664 和 amdgcn 后端把切片降成两个标量,nvptx64 却降成 [i64; 2] 数组,靠人工对表才统一;标准库上 GPU 也还没做。评测侧的问题同样明确:传输耗时比 RAJA 慢近三倍没解释透,只给了怀疑方向;FIR/LTIMES 的差距用「展开决策不同」带过,没有拆开验证;整轮跑分最差落后 46%,对追求确定性能的 HPC 用户不是小数。400 倍差距靠编译器优化抹平这条最关键的主张,原型做了、系统评测没跟上。另外这是原型阶段的工作,「集成进上游 rustc」目前是设计意图加部分实现,LLVM 23 的 Offload 拆分本身也还在进行中。