Code Generator：rustc-codegen-cuda#

每个 Rust 程序最终都会到达 codegen 后端——将优化后的 MIR 转换为机器码的编译器部分。通常，该后端是 LLVM。cuda-oxide 换入了自己的后端 rustc-codegen-cuda，它拦截此过程以提取设备代码，并通过 cuda-oxide 管道进行路由，然后将其他所有内容交还给 LLVM，就像什么都没发生过一样。

本页解释了该后端如何加载、当 rustc 调用它时做了什么，以及它如何找到属于 GPU 的每个函数。

rustc 如何加载自定义后端#

rustc 有一个大多数人从未见过的标志：

rustc -Z codegen-backend=path/to/libfoo.so

当你传递此标志时，rustc 执行以下操作：

  1. 对共享库调用 dlopen。

  2. 调用 dlsym("__rustc_codegen_backend") 来找到入口点。

  3. 期望该函数返回一个 Box<dyn CodegenBackend>。

就这些。没有插件注册表，没有配置文件，没有握手协议。一个符号，一个 trait object，你就拥有了 codegen 管道。

cuda-oxide 提供此入口点：

#[unsafe(no_mangle)]
pub fn __rustc_codegen_backend() -> Box<dyn CodegenBackend> {
    let config = CudaCodegenConfig::from_env();
    let llvm_backend = rustc_codegen_llvm::LlvmCodegenBackend::new();
    Box::new(CudaCodegenBackend { config, llvm_backend })
}

这里发生两件事。首先，CudaCodegenConfig::from_env() 读取环境变量（CUDA_OXIDE_VERBOSE、CUDA_OXIDE_DUMP_MIR 等——见下方 环境变量）来配置后端行为。其次，它创建标准 LLVM 后端并将其存储在 CudaCodegenBackend 内部。这就是使整个架构运作的包装模式：cuda-oxide 不替换 LLVM 后端，而是包装它。

备注

cargo oxide build 会为你设置 -Z codegen-backend 标志。你永远不需要自己输入 dlopen 咒语，除非你喜欢那种事情。

拦截：codegen_crate()#

CodegenBackend trait 有几个方法，但关键是 codegen_crate(tcx)。这是 rustc 交出完整的、经过类型检查、借用检查和单态化的 crate，并说"把这个变成机器码"的地方。

当 rustc 调用 CudaCodegenBackend::codegen_crate(tcx) 时，会发生以下情况：

第 1 步：收集单态化项#

let (items, cgus) = tcx.collect_and_partition_mono_items(());

这给了我们 crate 中的每一个已经单态化的函数，分组到codegen 单元（CGU）中。CGU 是 rustc 并行代码生成的单元——可以把它想象成将成为单个目标文件的函数桶。

第 2 步：扫描设备入口点#

后端遍历每个 CGU 中的每个函数，并检查魔术名称前缀：

  • cuda_oxide_kernel_<hash>_——由 #[kernel] 设置

  • cuda_oxide_device_<hash>_——由 #[device] 设置

这些前缀是 proc macro 与后端通信的方式。没有特殊的属性元数据，没有侧信道——只是一个在人群中脱颖而出的名称。确切的前缀字符串（以及匹配和剥离它们的辅助函数）位于工作区内部的 reserved-oxide-symbols crate 中；macro 端和 collector 端都从那里导入，因此契约保持在同一个地方。8 个十六进制字符的 <hash> 使得意外命中的冲突实际上不可能：没有人会意外地写出 cuda_oxide_kernel_246e25db_foo。

第 3 步：如果找到设备代码，构建并编译#

如果找到任何 kernel 或 device 函数，则按顺序发生两件事：

a) 收集设备调用图。 collector::collect_device_functions() 从每个 kernel 入口点执行广度优先遍历，发现所有从设备代码传递调用的函数。更多内容见 设备函数收集。

b) 生成设备代码。 device_codegen::generate_device_code() 将收集到的函数桥接到 stable MIR，并运行完整的 cuda-oxide 管道（dialect-mir -> mem2reg -> dialect-llvm -> .ll -> PTX）。

第 4 步：始终编译宿主代码#

无论是否找到设备代码：

self.llvm_backend.codegen_crate(tcx)

包装的 LLVM 后端编译所有宿主代码——main()、你的 CLI 解析器、你的异步运行时，一切。设备管道是一个支线任务；宿主管道始终运行。

其他 Trait 方法#

CodegenBackend trait 还要求 join_codegen() 和 link()。它们分别处理等待并行 codegen 线程和调用链接器。在 cuda-oxide 中，两者都是纯粹的委托：

fn join_codegen(&self, ongoing: Box<dyn Any>, sess: &Session) -> ... {
    self.llvm_backend.join_codegen(ongoing, sess)
}

fn link(&self, sess: &Session, codegen: ..., outputs: &OutputFilenames) -> ... {
    self.llvm_backend.link(sess, codegen, outputs)
}

没有拦截，没有修改。LLVM 后端像平常一样处理链接。

备注

对于依赖树中的大多数 crate——core、cuda-core、随机的工具 crate——没有找到 kernel，后端只是委托给 LLVM。设备管道只在实际包含 #[kernel] 函数的 crate 中激活。对于一个典型项目，这意味着几十个 crate 中只有一个触发设备编译。

设备函数收集#

找到 kernel 入口点是容易的部分。困难的部分是找出那些 kernel 调用的其他所有东西。一个 kernel 可能调用一个辅助函数，它调用 cuda-device 内部函数，它调用 core 数学函数——所有这些都需要出现在 PTX 中。

collector 模块通过 MIR 调用图的广度优先遍历来处理这个问题。

工作列表算法#

worklist = [所有 kernel 入口点]
visited  = {}
collected = []

while worklist 非空:
    fn = worklist.pop()
    如果 fn 在 visited 中: continue
    visited.add(fn)

    mir = tcx.instance_mir(fn)
    collected.push(fn)

    对于 mir 中的每个 basic_block:
        对于 [Call, Drop] 中的 terminator:
            callee = resolve_callee(terminator)
            如果 should_collect(callee):
                worklist.push(callee)

对于每个函数，collector 通过 tcx.instance_mir() 检索其 MIR，然后扫描每个基本块的终止器。Call 和 Drop 终止器引用被调用者，如果它们通过了 crate 过滤规则，则将其添加到工作列表中。

输出是一个 Vec<CollectedFunction>，其中每个条目携带：

  • 单态化后的 Instance（完全解析、无泛型的函数）。

  • 一个 is_kernel 标志（用于在 PTX 元数据中标记 GPU 入口点）。

  • export_name（在最终 PTX 中可见的符号名称）。

导出名称和 FQDN 对齐#

collector 必须生成与 MIR 翻译器为调用目标生成的相匹配的导出名称。两侧都使用完全限定域名（FQDN）——例如，helper_fn::cuda_oxide_device_<hash>_vecadd_device 而不是裸的 vecadd_device。

rustc_public API 的 CrateDef::name() 返回包含 crate 名称的 FQDN。在 collector 侧，def_path_str() 为本地项省略了 crate 名称，因此一个小辅助函数（fqdn()）为本地定义预先添加它，以产生相同的字符串。

在降级期间，定义侧和调用侧都将 :: 转换为 __，产生有效的 LLVM/PTX 标识符（例如，helper_fn__vecadd_device）。

对于泛型或复杂的名称（尖括号、闭包），则使用重整后的符号名称——这是唯一的，并且已经是一个有效的标识符。

Kernel 入口点遵循单独的路径：compute_kernel_export_name 获取 #[kernel] macro 的基本名称，并且对于泛型和闭包-泛型的实例化，追加 _TID_<hex32>，其中 <hex32> 是 rustc 128 位 type-id 哈希值（泛型参数元组的）。宿主启动器通过 core::intrinsics::type_id intrinsic（由 cuda_host::type_id_u128 包装）计算相同的哈希值，因此两侧逐字节一致。非泛型 kernel 保持其裸的基本名称。

备注

这个 FQDN 对齐策略将在框架升级时被 pliron 的 Legaliser 所取代。Legaliser 提供系统化的名称清理及冲突检测，使手动的 :: 到 __ 替换变得不必要。

Crate 过滤规则#

并非依赖树中的每个函数都属于 GPU。collector 强制执行关于哪些 crate 允许出现在设备代码中的规则：

Crate

状态

备注

本地 crate

允许

你的 kernel 代码

cuda_device

允许

GPU 内部函数（存根被过滤——见下文）

core

允许

no_std 标准库

其他 no_std crate

允许

依赖树中的任何 crate

std

禁止

编译时错误

alloc

允许（如果 GPU 分配器可用）

实验性

规则很简单：如果一个 crate 可以想象在 GPU 上运行（无操作系统依赖，无文件 I/O，无网络），则允许。如果它拖入了 std，则不。

错误消息#

当有人意外地试图在设备代码中使用 std——也许是调试时遗留下来的一个散落的 println!——collector 会生成一个清晰的错误，而不是让构建在三个阶段后以难以理解的 PTX 汇编器消息失败：

CUDA-OXIDE: FORBIDDEN CRATE IN DEVICE CODE
Device code calls: std::io::_print
From crate: 'std'
Only these crates are allowed in device code:
  - Local crate (your kernel code)
  - cuda_device (GPU intrinsics)
  - core (no_std standard library)

这是在调用图层面捕获问题而不是希望 LLVM 能产生有用的诊断信息的优势之一。（它不会。）

内部函数存根过滤#

cuda_device 包含像 threadIdx_x() 这样的函数，其函数体只是 unreachable!()。这些是存根——为了 rustc 能对设备代码进行类型检查而存在的占位符，但它们永远不会被编译成真正的函数。它们稍后在管道中被 dialect-nvvm 操作替换（例如，threadIdx_x() 变成 nvvm.read_ptx_sreg_tid_x）。

collector 识别这些存根——它们来自 cuda_device 且除了单个 unreachable 终止器之外没有有意义的 MIR——并跳过它们。如果你在想"但函数体是 unreachable!()，在运行时不会 panic 吗？"：不会，因为 mir-importer 在任何代码生成发生之前将这些函数的调用替换为相应的 GPU 硬件指令。存根体永远不会执行。

跨 Crate 设备编译#

有一个听起来很简单但仔细想想你会发现并不简单的问题：当你在一个 kernel 中调用 cuda_device::thread::index_1d() 时，编译器如何获得该函数的 MIR？你的 crate 没有它的源代码。cuda_device 是几小时前由另一台机器编译的。

答案在于 rustc 如何存储元数据。

.rlib 流程#

当 rustc 编译 cuda_device 时，该 crate 中没有 #[kernel] 函数。cuda-oxide 后端发现没有任何东西可拦截，并完全委托给 LLVM。输出是一个标准的 .rlib 归档文件，包含：

  • 已编译的机器码——宿主端目标文件。

  • .rmeta 元数据 blob——类型信息、trait 实现，关键的是，某些函数的序列化优化后的 MIR。

之后，当你编译你的 vecadd crate 时，collector 找到一个 kernel 入口点并开始 BFS 遍历。它遇到对 cuda_device::thread::index_1d() 的调用。此时，tcx.instance_mir() 深入到 cuda_device 的 .rmeta blob 并检索该函数的 MIR——完全优化、单态化，可直接翻译。

compiler/images/cross-crate-rlib.svg

跨 crate 设备编译。阶段 1 将 cuda_device 编译为带有序列化 MIR 的 .rlib，存储在 .rmeta blob 中。阶段 2 编译 vecadd——collector 找到一个 kernel，BFS 遍历进入 cuda_device，从 .rmeta 读取 MIR，并将所有内容编译为 PTX。#

来自依赖 crate 的设备代码被延迟编译为 PTX——仅当你 crate 中的 kernel 间接引用了它时才会编译。如果你从不调用一个函数，它永远不会出现在 PTX 中，即使它存在于依赖中。

哪些函数的 MIR 会被编码在 .rmeta 中？#

并不是每个函数的 MIR 都会保留到 .rmeta blob 中。rustc 有关于什么被序列化的规则：

类别

序列化的原因

泛型函数

必须在下游单态化——MIR 是"模板"

#[inline] 函数

明确请求用于跨 crate 内联

#[inline(always)] 函数

相同但更坚持

小的叶函数

启发式：无调用，少量语句——内联成本低

大多数 cuda_device 函数属于其中至少一个类别。GPU 内部函数很小。辅助函数是 #[inline]。在元素类型上的泛型函数……嗯，是泛型的。结果是几乎所有设备端库代码的 MIR 都可以用于跨 crate 编译。

备注

这不是 cuda-oxide 的特性——它是标准的 rustc 行为，旨在支持跨 crate 内联和泛型实例化。cuda-oxide 只是碰巧从中获益良多，因为使 rustc 能够在跨 crate 内联 Vec::push 调用的相同机制，也让我们能够将 cuda_device::thread::index_1d() 调用编译为 PTX。

桥接到 Stable MIR#

collector 完成后，我们有一个 Vec<CollectedFunction>，其中包含作为内部 rustc Instance 的每个设备函数。但 mir-importer 不懂 rustc 内部——它使用 rustc_public 提供的 stable MIR API。我们需要一座桥。

device_codegen 模块处理翻译：

rustc_internal::run(tcx, || {
    let stable_instances: Vec<_> = collected
        .iter()
        .map(|f| rustc_internal::stable(f.instance))
        .collect();
    mir_importer::run_pipeline(&stable_instances, &config)
})

这里发生两件事：

  1. rustc_internal::run(tcx, || { ... }) 设置 rustc_public 所需的线程局部上下文。在这个闭包内部，stable MIR 查询可用。在它之外，它们会 panic——stable API 需要访问编译器会话，而此函数提供了它。

  2. rustc_internal::stable(f.instance) 将每个内部 rustc_middle::ty::Instance 转换为其稳定等价物 stable_mir::mir::mono::Instance。这是单向转换——内部类型可以变成稳定类型，反之则不行。

转换完成后，mir_importer::run_pipeline() 接手。它读取每个 instance 的 stable MIR，将其转换为 pliron 的 MIR dialect，并运行完整的降级管道直到 PTX。该过程在 MIR Importer 中讲述。

环境变量#

CudaCodegenConfig::from_env() 读取以下环境变量来控制后端行为。所有都是可选的——默认值会产生一个安静的、面向生产的构建。

变量

效果

CUDA_OXIDE_VERBOSE

打印编译进度（找到了哪些 kernel、管道阶段、计时）

CUDA_OXIDE_DUMP_MIR

导入后将 dialect-mir 模块 dump 到 stderr（以及 mem2reg 之后）

CUDA_OXIDE_DUMP_LLVM

降级后将 dialect-llvm 模块 dump 到 stderr

CUDA_OXIDE_PTX_DIR

覆盖 .ptx 文件的输出目录（默认：宿主二进制文件旁边）

CUDA_OXIDE_TARGET

覆盖 GPU 目标架构（例如，Hopper 使用 sm_90a）

CUDA_OXIDE_SHOW_RUSTC_MIR

在转换为 pliron 之前 dump 原始 rustc MIR（用于调试导入 bug）

这些有意地设为环境变量而不是命令行标志。codegen 后端从 rustc 的参数解析器接收的信息非常有限——环境变量是传递配置的最简单方式，无需与编译器驱动程序的标志管道做争斗。

CUDA_OXIDE_TARGET 位于 cargo oxide 遵循的优先级链中：显式的 --arch 优先，然后是 CUDA_OXIDE_TARGET，然后是 cargo oxide run 的主机 CC 自动检测（仅适用于 run，不适用于 build 或 pipeline），最后是后端基于功能特性的默认值。

备注

CUDA_OXIDE_VERBOSE=1 cargo oxide build 是你在调试编译器时最好的朋友。它会显示确切收集了哪些函数、它们来自哪个 crate，以及每个管道阶段花费了多长时间。

汇总#

以下是 cargo oxide build 到 .ptx 文件的完整流程，从 codegen 后端的视角：

compiler/images/codegen-flow.svg

完整的 codegen 后端流程。当 codegen_crate 被调用时，CGU 被扫描以寻找 kernel/device 符号。如果找到，collector 执行 BFS 遍历，设备 codegen 桥接层转换为 stable MIR，mir_importer 生成 PTX。LLVM 宿主后端无论如何都会始终运行。#

包装模式的美妙之处在于 cuda-oxide 对编译管道的其余部分是不可见的。链接器看到正常的对象文件。构建系统看到正常的产物。任何不寻常事情发生的唯一证据就是一个 .ptx 文件躺在宿主二进制文件旁边，准备在运行时被 CUDA 驱动加载。

接下来去哪里#

有了收集好的设备函数和到手的 stable MIR，下一个阶段是翻译——在 MIR Importer 中讲述。