English | 中文版

12. 相关工作:Rust on GPU/NPU,以及与 NVIDIA 工具链的整合可能

摘要:把 Rust 带到异构计算上的项目已有不少,以“内存安全“作为卖点的也不只 ascend-rs 一家。本章把 ascend-rs 放进当下活跃的 Rust-on-accelerator 项目矩阵中——rust-cuda、rust-gpu、krnl、cudarc、wgpu,以及最近宣布的 OxiCUDA——再探索 NVIDIA 一侧具体的整合路径。我们 tree 里已经有 mlir_to_gpu 后端,它从驱动 Ascend 路径的同一份 ascend_tile_* MLIR 直接生成 CUDA C;因此真正值得问的问题,是“如果把 kernel、host runtime、以及安全卫士缝起来,相比各自重新实现一遍,整个生态能得到什么“。本章识别出四个整合机会;其中之一——把第 11 章的安全卫士跑在 OxiCUDA 这类 runtime-PTX 项目生成的 PTX 上——是一个真正新颖的联合贡献,因为现今的 Rust-GPU 项目没有谁在自己生成的低层 IR 上带编译期安全分析。

12.1 Rust on Accelerators 全景

第 1 章包含了一个五行的表格,概要介绍了开源生态的全景。那张表对一个背景小节来说够用,但对于一次诚实的对比就太粗了——它把非常不同的设计塞进同一个“approach“格子里。本节把矩阵沿着真正重要的几个轴展开:每个项目替换了厂商技术栈中的什么,kernel 如何到达设备,是否在 Rust 自身类型系统之外提供任何编译期安全保证,以及宿主侧运行时是什么形态。

从这个矩阵里可以看出几个 ch01 表格被压平掉的事实。

第一,只有两个项目尝试替换驱动之上的厂商技术栈:OxiCUDA 声称替换整个 CUDA 用户空间栈(从 cuBLAS 到 cuDNN);ascend-rs 替换 Ascend 的 kernel 编写语言并构建了自己的编译枢纽(ascend_compile),但运行时仍然调用 ACL / CANN。其它五个项目给你一种更安全的方式来表达 kernel,但保留了厂商的库。

第二,运行时模型干净地一分为二。rust-cuda、rust-gpu、krnl 和 ascend-rs 都走 AOT:kernel 在 build time 编译为机器码产物。cudarc、OxiCUDA 以及 wgpu 的 compute 路径都使用运行时编译(经由 nvrtc、运行时 PTX 发射,或第一次 draw 时的 shader 编译)。AOT 给你可复现性,并允许在 binary 构建之前跑安全分析;JIT 给你灵活性(形状特化、与目标相关的调优)。第 10 章和第 11 章只在 AOT 模型下才说得通——你没办法在还没发射的 PTX 上跑 check_all。

第三,这是让 ascend-rs 的位置与众不同的地方:没有别人在它生成的低层 IR 上提供编译期安全卫士。其他每个项目的安全叙事都是“Rust 借用检查器,加上对厂商运行时的安全 API 包装“。这是真实的贡献——CUDA C++ 的 aliasing 和 UAF bug 极为常见——但它对生成 kernel 的内存布局没有任何陈述。第 10 章和第 11 章对那个布局是有所陈述的:它们在 ptoas 产生的 stage-2 计划上跑 check_aliasing、check_capacity、check_dead_writes、check_slayout_consistency。这是一种厂商工具链本应自带、但实际并没有的分析;ascend-rs 把它作为旁路提供出来,作为厂商工具链以 rc=0 接受的 kernel 的一张“Rust 安全卡片“。

12.1.1 ascend-rs 落后于他人的部分

一节诚实的全景必须把差距点出来。在上面的矩阵里有三处可见的差距。

Library coverage. OxiCUDA 声称提供 cuBLAS、cuDNN、cuFFT、cuSPARSE、cuSOLVER、cuRAND 的等价物。ascend-rs 有一个 ascend_std 标准库和附录 E 中的 kernel 清单,但没有等价于完整 DNN 库的东西——我们的 DNN 故事是“把你需要的算子作为一个 Rust tile kernel 写出来“,而不是“use ascend_dnn::conv2d“。对于心智模型是“拿起 DNN 库就跑“的用户来说,ascend-rs 需要更多自行组装。

Ecosystem. rust-cuda、rust-gpu、wgpu 和 cudarc 都有多年在 crates.io 上的存在、下游用户和 bug report。ascend-rs 处于开发中(第 1 章的 status 列已说明),并且驻留在私有仓库里;公开的 yijunyu/ascend-rs 只放产物(参见仓库说明)。“Widely used“是我们目前还做不出的声明。

Target breadth inside NVIDIA. rust-cuda / rust-gpu / OxiCUDA 都把 NVIDIA 当作一等公民。ascend-rs 有一个 mlir_to_gpu 后端发射 CUDA C,但它属于次级 codegen 出口,并非主攻方向。如果你的主目标就是 NVIDIA,那些 NVIDIA 优先的项目在 NVIDIA 特定打磨上仍然领先于我们。

本章主张整合而非竞争,正是因为上面的几个轴在很大程度上是正交的。OxiCUDA 的 library coverage 与 ascend-rs 的安全卫士本来就是两件不同的工作,完全可以叠加在一起。

12.2 与 OxiCUDA 的逐项对比

OxiCUDA 值得专门一节,因为在相关项目里,它是最容易和 ascend-rs 混淆的:两者的简介里都有“替换厂商技术栈“的说法,两者都把 Rust 作为编写表层。这种混淆是表层的,值得讲清楚。

12.2.1 目标硬件

OxiCUDA 是 NVIDIA 优先。v0.1 公告把 NVIDIA CUDA 列为主要后端,Metal / Vulkan / WebGPU / ROCm / Intel Level Zero 作为附加后端。共享抽象似乎是“一个 kernel + 一份运行时库“,由各厂商后端实现它。

ascend-rs 是 Ascend 优先。Ascend 910B2 / 310P 上的 CANN 8.5 是主目标;crates/rustc_codegen_mlir/src/ 中的 14 个 mlir_to_*.rs 是次级 codegen 出口,共享 ascend_tile_* MLIR dialect 作为公共 IR。这些出口按字母顺序为:aie、bang、cpp、csl、gaudi、gpu、hexagon、linalg、msl、musa、nki、pto、spirv。对一个 NVIDIA 用户而言,mlir_to_gpu 存在并能发射 CUDA C,但没人声称它和一个专门面向 NVIDIA 的项目一样打磨成熟。

12.2.2 替换了什么

这是两个项目在 headline 上看起来相似、在细节上分道扬镳的轴。

OxiCUDA 替换的是 CUDA SDK 运行时栈:你不安装 nvcc,你不安装 CUDA toolkit,你不写 .cu 文件。剩下的唯一运行时依赖是 libcuda.so(驱动)。PTX 在运行时从 Rust 数据结构生成并通过 Driver API 交给驱动。替换是横向的,跨越用户空间的库。

ascend-rs 替换的是 kernel 编写语言,并提供独立的 编译枢纽(ascend_compile),但不替换 CANN / ACL。替换是纵向的,在 kernel 编程切片之内:你的 kernel 源码是 Rust 而不是 AscendC C++,然后 rustc_codegen_mlir 把它经由 MLIR 下沉到 AscendC,再由 bisheng 产出 NPU binary。CANN 仍然是运行时。我们没有重写 cuBLAS 在 Ascend 一侧的对应物(aclBLAS)。

这种差异对部署是有影响的。OxiCUDA 用户的 NVIDIA 机器只配驱动加上编译好的 Rust binary。ascend-rs 用户的 910B2 机器需要安装 CANN 8.5,加上 Rust binary 和 codegen 产物库。前者更激进;后者更保守,在真实的厂商硬件上落地更快。

12.2.3 编译模型

OxiCUDA 在运行时生成 PTX。kernel 路径上没有 .cu 文件、磁盘上没有 .o 文件;PTX 字符串在进程内由 Rust 类型构建,提交给 Driver API,module 句柄存入缓存。这种方式在精神上与 cudarc 加 nvrtc 类似,但通过直接发射 PTX 砍掉了 nvrtc 依赖。

ascend-rs 是 AOT 编译。在 cargo build 时,rustc_codegen_mlir 产出 MLIR,mlir_to_cpp 或 mlir_to_pto 产出 C++ 或 PTO-MLIR,bisheng(经由 ascend_compile)产出 NPU 目标代码,最终链接到 kernels.so,host 在运行时通过 libloading 加载。第 7 章给出完整的图。

两种模型严格意义上没有谁更好。AOT 让我们能跑一个编译期安全卫士;JIT / 运行时发射让 OxiCUDA 能基于 Rust 代码只在运行时才知道的形状做特化。自然的综合方案——12.3 节会接上——是:即便 PTX 是运行时发射的,只要在 launch 之前有一个时刻能将其冻结,AOT 安全分析仍可以应用到这份 PTX 上。

12.2.4 安全叙事

两个项目在编写侧都说“类型安全、内存安全的 Rust“,并且在 API 表层都兑现了这个声明。Rust 的借用检查器在 host 一侧抓它能抓到的东西。

ascend-rs 通过第 10 章和第 11 章的安全卫士又往前迈了一步。那两章描述了在 ptoas 产生的 stage-2 计划上运行的六个 pass:check_aliasing、check_capacity、check_dead_writes、check_slayout_consistency,以及另外两个。它们抓的是 ptoas 自己以 rc=0 接受的 bug——也就是厂商工具链产出一个 binary、这个 binary 能跑、但悄悄破坏数据的那一类 case。这是一类没有“对运行时的安全包装“能抓到的 bug,因为 unsafety 存在于包装之下,在生成 kernel 的内存布局里。

截至本文撰写时,OxiCUDA 的 v0.1 还没有为它发射的 PTX 文档化对应的分析。这无意贬低——v0.1 就是 v0.1,该阶段的重点本就是 library coverage,而把分析往后排是合理的取舍。这恰恰是一个整合机会,12.3.3 节会讲。

12.2.5 范围

OxiCUDA 在框架层面更宽:其公告把计算图、GPU 训练、推理和强化学习都列入目标。一个有这种野心的 v0.1 是“先把表层占住,再慢慢打磨“。

ascend-rs 更窄:我们提供编译器后端、标准库、HAL、安全的运行时包装,以及一份经过验证的 kernel suite(DeepSeek-R1-Distill-Qwen-1.5B 在 910B2 上端到端 decode 180+ tok/s,跨 MultiKernelBench 类目 500+ 编译通过的 kernel——见第 9 章)。我们不提供训练循环或 RL 框架。有人可以在 HAL 之上构建这些,但那不是我们已经发布的东西。

12.2.6 一句话对比

OxiCUDA:删掉 CUDA SDK,运行时从 Rust 发射 PTX,覆盖 NVIDIA library surface。

ascend-rs:用安全 Rust 写 NPU kernel,经 MLIR 穿过厂商工具链,证明厂商工具链自己证不了的安全属性。

两者是各自独立的项目,彼此之间也没有直接的竞争关系。下一节论证它们可以在接缝处接到一起。

12.3 NVIDIA 侧的整合机会

crates/rustc_codegen_mlir/src/mlir_to_gpu.rs 自我们启动多 vendor 出口以来就在 tree 里,它已经是落地代码,而非纸面设想。给定驱动 Ascend 路径的同一份 ascend_tile_* MLIR,它发射一个可被 nvcc -arch=sm_80 或 clang++ --cuda-gpu-arch=sm_80 编译的 .cu 文件。该文件顶部的映射表覆盖了核心算子——ascend_tile_load_f32、_store_f32、_add_f32、_sub_f32、_mul_f32、_exp_f32、_softmax_f32、_reduce_max_f32、_reduce_sum_f32、_scale_f32——通过直接的 CUDA kernel 模式,其中 matmul_f32 当前作为 cuBLAS-SGEMM 的占位注释发出。

以此为起点,有四个整合机会自然浮现。它们按今天可落地的具体程度排序。

12.3.1 短期:mlir_to_gpu + cudarc host 运行时

今天 mlir_to_gpu 产出一个 .cu 字符串。这个字符串被喂给什么是用户自己的事。ascend-rs 不提供 NVIDIA host 运行时。

短期整合是把 .cu 输出连接到 cudarc 的安全 driver/runtime 绑定:

1. 用户用 ascend_std tile API 写 kernel,和 Ascend 完全一样。
2. rustc_codegen_mlir 配合 ACLRS_CODEGEN_PATH=gpu 发射 .cu。
3. nvcc(或 clang++ --cuda-gpu-arch=sm_80)编译为 .ptx 或共享库。
4. 在 host 一侧,一个小的 ascend_hal CUDA 后端(类似于 crates/ascend_hal/ 中已经存在的、被 CLAUDE.md 引用的 cuda 后端)使用 cudarc 进行设备初始化、分配、stream 创建和 launch。

用户得到一份 Rust 源码:设 ACLRS_CODEGEN_PATH=pto 跑 910B2,设 ACLRS_CODEGEN_PATH=gpu 跑 NVIDIA。host 一侧是 cudarc 的事——一个 API 维护良好、有测试套、有数千下游用户的 crate——而不是我们的事。

这是能跑通真实端到端路径的最小整合,也是有人最可能先做原型的那条路。瓶颈很平凡:mlir_to_gpu 的 matmul 路径需要把 cuBLAS 调用补上(今天还是 TODO 注释),ascend_hal CUDA 后端需要更多测试覆盖(cudarc 依赖已经存在,只是藏在 feature gate 之后)。两者都属于纯工程问题,无关研究。

12.3.2 中期:把运行时 PTX 发射作为 nvcc 的替代评估

上一步骤 3 中的 nvcc 依赖,与 CANN 的 bisheng 在 Ascend 一侧带来的同类构建期依赖性质一致。这种依赖是站得住脚的——两者都是厂商的官方编译器——但它也是最大的整合摩擦:CI 机器需要 toolkit,Docker 镜像膨胀,nightly build 在厂商升级时挂掉。

OxiCUDA 的运行时 PTX 发射是自然的替代。如果 mlir_to_gpu 后端不发射 CUDA C,而是直接发射 PTX(经由一个 mlir_to_ptx,或穿过 LLVM 的 NVPTX target),那么 nvcc 依赖消失,部署画像就匹配 OxiCUDA 已经在自己一侧做的事。

有两条路径:

1. 扩展 mlir_to_gpu,通过 mlir-sys 调用 LLVM 的 NVPTX 后端,直接把 MLIR 下沉为 PTX。这把 codegen 完全留在 ascend-rs 内部。
2. 与 OxiCUDA 合作:ascend-rs 把 tile MLIR 下沉到一个中间形态(LLVM dialect,或 OxiCUDA 的 Rust-AST 数据结构),然后由 OxiCUDA 现有的 PTX 发射器接管。

方案 2 值得考虑,因为 OxiCUDA 已经解决了“为现代架构生成有效 PTX“这个问题,并在维护那个 PTX 发射器。方案 1 更自洽,但意味着我们要扛起 PTX 生成的复杂度。

两条路径之间没有明显赢家;正确答案取决于 OxiCUDA 的 API 有多稳定,以及 mlir_to_gpu 实际需要 NVPTX 表面积的多大部分。一次探针——为 softmax kernel 用两种方式生成 PTX,diff 输出,在真实 NVIDIA GPU 上做基准——是自然的下一步。

12.3.3 中期:把第 11 章的安全卫士跑在 PTX 上

这是新颖度最高的机会。第 11 章安全卫士的 pass——check_aliasing、check_capacity、check_dead_writes、check_slayout_consistency——逻辑上并不是 Ascend-specific 的。它们运行在一个 stage-2 计划之上:一份扁平的 tile 列表,每个 tile 携带 (space, offset, rows, cols, dtype, blayout, slayout) 元组,加上其上的依赖图。check_aliasing 中没有任何东西知道它是跑在 PTO-MLIR 上而不是别的 IR 上。Ascend-specific 的部分都在 parser(parse_stage2)里,它从 ptoas --print-after-all 的输出产出该计划。

一个能从 PTX 产出等价 stage-2 计划的 parser——从 shared-memory 分配、从 ld.global / st.global 访问、从 warp 级 shuffle 模式——能让同样的六个 pass 跑在 NVIDIA kernel 上。Ascend 安全卫士抓到的每一种“我的 kernel 能跑但答案不对“的 bug,在 NVIDIA 一侧都有对应物(shared-memory bank 冲突、aliased __shared__ 数组、对每 SM 48 KB 或 100 KB 上限的容量超界)。

具体地:

1. 写一个 parse_ptx_stage2,产出与今天 parse_stage2 产出的同款 Plan 结构,但来源是 PTX 而非 PTO-MLIR。
2. 在结果上跑现有的 check_all(plan)。
3. 把它接到一个环境变量(镜像 ACLRS_PTO_SAFETY)——例如 ACLRS_PTX_SAFETY=error / =warn / 不设。

parser 是难的部分;check 直接复用。这是面向任何 PTX 发射型 Rust-GPU 项目(OxiCUDA 或其它)的一份干净的联合贡献。今天没有任何这类项目在它生成的 PTX 上提供编译期安全分析。

12.3.4 长期:共享一份 tile IR

在栈的顶端,ascend_tile_* intrinsic 与 OxiCUDA(或 rust-cuda、或 rust-gpu)用作 tile 抽象的东西在解决同一个问题:用一种可被下沉到具体厂商的形式,描述一块矩形的数据 tile 以及其上的一个操作。我们在 crates/rustc_codegen_mlir/src/mlir_to_*.rs 中有 15 vendor backends;每一个都消费同一份 ascend_tile_* MLIR。

长期整合是把这份 tile dialect 做成一个独立产物——一个 MLIR dialect 加一份参考下沉——并邀请其它项目下沉 进入 它(OxiCUDA 的 Rust AST → tile MLIR → PTX)或下沉 离开 它(tile MLIR → SPIR-V → Vulkan,正如 mlir_to_spirv 已经做的)。到那时图景是:一份 tile IR、N 个 frontend、M 个 backend,以及一个夹在两端之间的安全卫士。

这是四个机会里最具推测性的一个,也是最需要看到真实用户需求才值得承担维护成本的那个。把它列在这里,是因为前三个机会自然把方向推向这里——如果你已经在共享一份 host 运行时(12.3.1)、共享一份 PTX 发射器(12.3.2),并共享一份安全卫士(12.3.3),那么共享所有这些一起下沉的源头 IR,就只是把这幅画补完了。

12.4 ascend-rs 能反哺生态的部分

如果整合发生,贡献是双向的。ascend-rs 中可以脱离 Ascend 上下文复用的部分,按就绪度递减排序:

1. 安全卫士。 六个 pass 在 pass 这一层已经与 Ascend specifics 解耦。一个 PTX 的 stage-2 parser 即可让它们对 NVIDIA 解锁。
2. ascend_compile 编译枢纽模式。 今天它分发到 bisheng,但其中的三个验证 pass 与双标志基础设施是可以一般化的:第 7 章描述了一个 C++-to-binary 的枢纽,而 CUDA C、SYCL、HIP 与 AscendC 都共享这条编译流水线的形态。一个把 ascend_compile 结构因式分解出来的多 vendor compile crate,本身就是一个合理的 Rust crate。
3. MLIR tile dialect。 已经有 15 个 vendor 后端共享它。把它从 ascend-rs 仓库拆出来属于工程任务,无关研究。
4. Rust 侧的 kernel 语料。 跨 MultiKernelBench 类目 500+ 个 kernel,从 softmax 到 DeepSeek 的 MLA attention,都用安全的 tile-API Rust 写出。对任何想要测试套的项目,这是一个起点。

不能搬迁的部分:CANN 相关的细节(pipe barrier、L0/L1/UB 内存布局、SoC 版本守卫)、ACL 运行时包装、aiv_kernel 宏中 AscendC 特定的 ABI,以及任何提到 910B2 的部分。这些占代码库的 ~30%;另外 ~70% 是通用有用的。

12.5 本章不是什么

值得明确说出本章刻意不主张的几点。

它不主张 ascend-rs 在某个基准上赢过了 OxiCUDA——它们没有提供可比的基准,即便有,目标硬件也不同。

它不主张 OxiCUDA 在运行时生成 PTX 是错的。这本就是一个有真实优势的合理设计选择;12.3.2 节把它当作整合机会来讨论,从不视其为设计缺陷。

它不主张第 11 章的安全卫士今天就能在 PTX 上工作。12.3.3 节把它列为中期机会是有原因的:PTX 的 stage-2 parser 还没写。本章所主张的范围只到 pass 可以移植这一层;parser 本身的工作量是另一回事。

它不主张 rust-cuda、rust-gpu、cudarc 或 krnl 已经过时。它们各自占据设计空间中一个独特的点(见 12.1),对于它们当初被构建出来要服务的用例,它们仍然是正确答案。

本章主张的是窄而(我希望)立得住的:ascend-rs 与 NVIDIA 一侧的 Rust 项目,正在用互补的强项解决重叠的问题;整合面已经在今天的代码里可见(mlir_to_gpu + cudarc + 第 11 章的安全卫士);一次认真把它们接到一起的努力,会产生比任何一个单独项目都更强的故事。