上一篇(codegen_crate 入口接管)讲完 backend 怎么进入 device 路径。这一篇看进入之后做的两件大事:找出所有需要编进 PTX 的函数(collector)+ 把 rustc internal MIR 转成 stable MIR(rustc_public)。重点拆
count_device_fns_in_cgus的扫描流程、246e25db这串魔数的由来、以及 stable MIR 这个解耦设计。
1. 概览:Step 5 干两件事
codegen_crate 检测到 device code
│
▼
1. collector::collect_device_functions
├── 找所有 #[kernel] 入口(命名约定 + 单态化检查)
├── worklist BFS 遍历调用图
└── 输出 Vec<CollectedFunction>(internal Instance)
│
▼
2. device_codegen::generate_device_code
├── rustc_internal::run 进入 stable 上下文
├── rustc_internal::stable(internal_instance) → stable Instance
└── 喂给 mir_importer::run_pipeline两件事其实是”找谁”+“换格式”。找谁难点在识别 kernel 和追踪调用图;换格式核心是跟 rustc 内部 API 解耦。
2. 一个反直觉的设计选择:命名约定 vs Attribute
直觉答案”看 #[kernel] 属性”。但 cuda-oxide 完全不读属性——它在 is_kernel_function 里只看函数的 path string:
pub fn is_kernel_function(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
is_kernel_symbol(&tcx.def_path_str(def_id))
}
pub fn is_kernel_symbol(name: &str) -> bool {
name.contains(KERNEL_PREFIX) // "cuda_oxide_kernel_246e25db_"
}为什么走命名约定?
| 原因 | 说明 |
|---|---|
| mangled 名字会进 LLVM IR | host 端要通过符号名找到 kernel 在 PTX 里的入口,字符串是天然的桥梁 |
| 跨 crate 调用照样工作 | 用户在 kernel_lib::scale_f32 这样调,FQDN 里仍然带前缀,识别不变 |
| 属性需要 tcx 查询 | 命名约定直接看 path,简单到不需要任何 rustc API 上下文 |
代价:用户写一个名叫 cuda_oxide_kernel_xxx 的函数,可能被误识别。这就是 246e25db 这串魔数登场的原因。
3. 246e25db 是哪里来的
reserved-oxide-symbols crate 里写得明明白白:
/// Magic suffix appended to every prefix to defend against accidental
/// name collisions in user code.
///
/// `sha256("cuda_oxide_ + rust")` truncated to 8 hex chars. The exact
/// value is irrelevant — what matters is that it's fixed forever and
/// no human would ever type it as part of a regular function name.
pub const HASH_SUFFIX: &str = "246e25db";计算过程:
输入: "cuda_oxide_ + rust" (一串字面字符串)
│
▼ sha256
输出: 246e25db<剩下 56 个 hex 字符> (256 bit / 64 hex)
│
▼ 截前 8 个 hex
最终: "246e25db"自己在 shell 里验证一下:
$ echo -n "cuda_oxide_ + rust" | sha256sum
246e25db0ab4c8caa753a365376b2f54a543a0463865ad484d7be5475095ba58 -前 8 个 hex 246e25db 就是 HASH_SUFFIX 的值。注意 echo -n(不要尾部换行)——加了 \n 整个 hash 就变了。
注意几点:
- 输入字符串是任意的——团队选了一个有纪念意义的串,换成其它的也行
- 取 8 个 hex(32 bit 哈希空间)就够了——人类不会随手写出这串字符
- 常量是”钉死”的——一旦发布就不能改,改了所有已编译产物都失效
有个测试专门把这个值 pin 住:
#[test]
fn hash_value_is_pinned() {
assert_eq!(HASH_SUFFIX, "246e25db");
assert_eq!(KERNEL_PREFIX, "cuda_oxide_kernel_246e25db_");
// ...
}只要这个 hash 改了,这个测试就挂——防止有人不小心改动。
4. 四个前缀:cuda-oxide 的命名空间
246e25db 这串后缀被嵌进四个固定前缀,分别给不同种类的代码生成对象用:
| 前缀常量 | 实际字符串 | 用途 |
|---|---|---|
KERNEL_PREFIX | cuda_oxide_kernel_246e25db_ | #[kernel] 函数(GPU 入口) |
DEVICE_PREFIX | cuda_oxide_device_246e25db_ | #[device] 函数(只能被 device 端调) |
DEVICE_EXTERN_PREFIX | cuda_oxide_device_extern_246e25db_ | #[device] extern "C" { ... } 声明(链接外部 LTOIR) |
INSTANTIATE_PREFIX | cuda_oxide_instantiate_246e25db_ | 泛型 kernel 的闭包单态化 helper |
reserved-oxide-symbols 这个 crate 是唯一定义这些前缀的地方——proc macro side(生成名字)和 collector side(消费名字)都从这里 import,保证两端永远一致。
5. #[kernel] proc macro 怎么改函数名
#[kernel] 处理非泛型函数的核心逻辑(cuda-macros crate):
fn generate_simple_kernel(mut input: ItemFn) -> TokenStream {
inject_thread_index_scope(&mut input);
let fn_name = input.sig.ident.clone(); // 原始函数名
let new_name = format_ident!("{}{}", KERNEL_PREFIX, fn_name); // 加前缀
input.sig.ident = new_name; // 改名
let expanded = quote! {
#[unsafe(no_mangle)] // 防 mangle
#input
// ... 配套的 CudaKernel impl 等 ...
};
TokenStream::from(expanded)
}逐步翻译:
// 用户写的
#[kernel]
pub unsafe fn hello_constant(out: *mut i32) {
*out = 42;
}
// 宏展开后(rustc 看到的)
#[unsafe(no_mangle)]
pub unsafe fn cuda_oxide_kernel_246e25db_hello_constant(out: *mut i32) {
*out = 42;
}
// 加上配套的 CudaKernel impl(host 端调用时用)
impl cuda_host::CudaKernel for hello_constant_kernel {
fn ptx_name() -> &'static str { "hello_constant" }
// ...
}两件事同时发生:
- 函数名字加前缀 → 这是 backend 识别的依据
#[unsafe(no_mangle)]→ 防止 rustc 进一步 mangle,确保符号名在 LLVM IR / PTX 里就是这个完整字符串
#[device] 宏走同样套路,只是用 DEVICE_PREFIX。
如果用户傻到精确写了 cuda_oxide_kernel_246e25db_evil? macro side 有防御。proc macro 处理 #[kernel] / #[device] 时调 reject_reserved_name,如果用户函数名以 RESERVED_ROOT = "cuda_oxide_" 开头直接报编译错:
error: function name reserved for cuda-oxide internal use6. count_device_fns_in_cgus 的扫描流程
到了 backend 这边,扫描代码就一段双重循环:
pub fn count_device_fns_in_cgus<'tcx>(
tcx: TyCtxt<'tcx>,
cgus: &[CodegenUnit<'tcx>],
) -> usize {
let mut count = 0;
for cgu in cgus { // ① 遍历所有 CGU
for (item, _data) in cgu.items() { // ② 遍历 CGU 里的 mono item
if let MonoItem::Fn(instance) = item // ③ 只看 Fn item
&& is_device_function(tcx, instance.def_id()) // ④ 命名约定匹配
&& is_fully_monomorphized(tcx, *instance) // ⑤ 单态化检查
{
count += 1;
}
}
}
count
}五步谓词,每一步过滤目的:
| 步 | 谓词 | 过滤掉什么 |
|---|---|---|
| ① | for cgu in cgus | (无)遍历所有 CGU |
| ② | for item in cgu.items() | (无)遍历每个 mono item |
| ③ | MonoItem::Fn(instance) | Static、vtable 等非函数 item |
| ④ | is_device_function → is_device_symbol 子串匹配 | 名字不带 cuda_oxide_device_246e25db_ 的所有函数 |
| ⑤ | is_fully_monomorphized | 泛型定义本身、嵌套泛型未解析的实例 |
count_kernels_in_cgus 的结构完全一样,只是 ④ 换成 is_kernel_function(检查 KERNEL_PREFIX 子串)。两个函数加起来回答一个问题:这个 crate 有没有任何需要走 device 路径的代码?
关于
is_fully_monomorphized的两个检查为什么必要,见 compiler 系列第 04 篇。
跑 cargo oxide run hello_constant,这两个谓词的真实值:
[PHASE 3/9] ... crate_name=hello_constant cgus=16 kernels=2 device_fns=016 个 CGU 里数出来 kernels=2(hello_constant、hello_kernel),device_fns=0(example 没单独的 #[device] 函数,所有 device 代码都是 kernel 入口)。
7. 互斥子串保证:为什么 is_device_symbol 不会误识 device-extern
注意 DEVICE_PREFIX 和 DEVICE_EXTERN_PREFIX 字符串前段有重叠:
DEVICE_PREFIX: cuda_oxide_device_246e25db_
DEVICE_EXTERN_PREFIX: cuda_oxide_device_extern_246e25db_
^^^^^^^
多了 "extern_"仔细看:DEVICE_PREFIX 里 device_ 之后就紧跟 246e25db_。一个 device-extern 名字里 device_ 之后是 extern_246e25db_,不包含完整的 device_246e25db_ 子串。
这就是 reserved-oxide-symbols 里那个测试守护的不变量:
#[test]
fn device_and_device_extern_are_mutually_exclusive_substrings() {
assert!(!DEVICE_EXTERN_PREFIX.contains(DEVICE_PREFIX));
assert!(!DEVICE_PREFIX.contains(DEVICE_EXTERN_PREFIX));
}hash 后缀让两个前缀变成”互斥子串”——一个名字里要么有 device_246e25db_,要么有 device_extern_246e25db_,不可能同时有。所以 is_device_symbol 不用写”如果是 extern 就排除”的特殊逻辑。这是一个非常优雅的设计副产品。
8. base name 提取:从 mangled 名字回原始名
backend 拿到一个 mangled 名字 cuda_oxide_kernel_246e25db_hello_constant,要在 PTX 里写没有前缀的入口名(hello_constant)。reserved-oxide-symbols 提供 strip 函数:
pub fn kernel_base_name(name: &str) -> Option<&str> {
name.find(KERNEL_PREFIX)
.map(|pos| &name[pos + KERNEL_PREFIX.len()..])
}也支持跨 crate 调用的 FQDN:
"cuda_oxide_kernel_246e25db_hello_constant" → "hello_constant"
"kernel_lib::cuda_oxide_kernel_246e25db_scale" → "scale"这就是为什么 cuda-oxide 同时支持 in-crate 和 cross-crate kernel 调用——base name 提取对路径前缀完全透明。
9. 找到 kernel 之后:Worklist BFS 遍历调用图
光找 kernel 入口不够——kernel 里调的所有函数也得编进 PTX,否则 device 端就缺符号。比如 hello_constant 调了 thread::xxx()、gpu_printf! 间接调了 vprintf 等。
collect_device_functions 用经典 worklist BFS:
struct DeviceCollector<'tcx> {
seen: HashSet<String>, // 去重,防止无限循环
worklist: VecDeque<CollectedFunction<'tcx>>,
result: Vec<CollectedFunction<'tcx>>,
}
// 算法
while let Some(func) = worklist.pop_front() {
if !seen.insert(func.mangled_name) {
continue; // 已访问过
}
// 看 func 的 MIR,找出所有它调用的函数
for callee in walk_mir_calls(tcx, func.instance) {
worklist.push_back(callee);
}
result.push(func);
}对 hello_constant 来说,真实 BFS 走完的调用图(docs/run.log):
[collector] Found kernel: kernels::cuda_oxide_kernel_246e25db_hello_constant -> hello_constant
[collector] Found kernel: hkernel::cuda_oxide_kernel_246e25db_hello_kernel -> hello_kernel
[collector] Processing hello_constant (3 basic blocks)
[collector] Skipping extern/intrinsic (no MIR): cuda_device::xxx
[collector] Skipping extern/intrinsic (no MIR): cuda_device::debug::__gpu_vprintf
[collector] Processing hello_kernel (1 basic blocks)两个 kernel 全部收集完,中间遇到的 intrinsic(cuda_device::xxx)和 extern device 函数(__gpu_vprintf)被跳过——它们没有 MIR body,后续会在 mir-importer 阶段被 intrinsic dispatcher 直接替换成 NVVM op,不需要进 collector 的 result。
10. 输出:CollectedFunction
pub struct CollectedFunction<'tcx> {
pub instance: Instance<'tcx>, // 单态化后的具体实例
pub is_kernel: bool, // 是不是 kernel 入口
pub export_name: String, // PTX 里的函数名(base name)
}is_kernel 这个 bit 后面会决定 PTX 输出:
is_kernel | PTX 输出 | 谁能调 |
|---|---|---|
true | .visible .entry 函数 | host 端通过 cuLaunchKernel |
false | .func 函数 | 只能被 device 端调用 |
11. stable MIR 转换:跟 rustc 内部 API 解耦
collector 返回的是 rustc internal 类型 Instance<'tcx>。直接传给 mir-importer 会有两个问题:
- mir-importer 不应该依赖 rustc 内部 API——rustc internal 是 unstable 的,每个 nightly 都可能改 API
'tcx生命周期病毒式传播——一旦签名带<'tcx>,从 importer 到 dialect 到 lower 全都得带,工程上很恶心
解决方案:用 rustc 团队官方的 rustc_public API(命名空间叫 stable_mir):
let result = rustc_internal::run(tcx, || {
let stable_functions: Vec<mir_importer::CollectedFunction> = functions
.iter()
.zip(export_names.iter())
.map(|(func, (export_name, is_kernel))| {
// 关键这一行:internal Instance → stable Instance
let stable_instance = rustc_internal::stable(func.instance);
mir_importer::CollectedFunction {
instance: stable_instance,
is_kernel: *is_kernel,
export_name: export_name.clone(),
}
})
.collect();
// 在 stable 上下文里调用 pipeline
mir_importer::run_pipeline(&stable_functions, &stable_device_externs, &pipeline_config)
});rustc_internal::run(tcx, closure) 干两件事:
- 进入 stable MIR 上下文:设置 thread-local,让
rustc_public::*API 可用 - 执行 closure 并清理
rustc_internal::stable(instance) 是 internal → stable 的转换桥。这样 mir-importer crate 整个签名都不带 'tcx,只依赖 rustc_public,跟 rustc 内部解耦。
类型对应表:
Internal (rustc_middle) | Stable (rustc_public) |
|---|---|
Instance<'tcx> | Instance |
Body<'tcx> | Body |
BasicBlock | BasicBlock |
Statement<'tcx> | Statement |
Terminator<'tcx> | Terminator |
Place<'tcx> | Place |
Rvalue<'tcx> | Rvalue |
Ty<'tcx> | Ty |
没有 'tcx,没有 arena 绑定——这些类型可以自由克隆、跨线程、跨模块传递。代价是访问数据要通过函数调用(比如 body.blocks())而不是直接字段访问,性能略差但完全可接受。
12. 完整流程串起来
用户写源码
┌─────────────────────────┐
│ #[kernel] │
│ fn hello_constant(...) │
└─────────────────────────┘
│
▼ proc macro 展开(编译期源码变换)
┌──────────────────────────────────────────────────────┐
│ #[unsafe(no_mangle)] │
│ fn cuda_oxide_kernel_246e25db_hello_constant(...) {} │
└──────────────────────────────────────────────────────┘
│
▼ rustc 前端 + 单态化
Instance,带上面那个 mangled 名字
│
▼ codegen_crate → count_kernels_in_cgus
for cgu / for item / MonoItem::Fn / is_kernel_symbol / is_fully_monomorphized
│
▼ collect_device_functions
worklist BFS 遍历调用图 → Vec<CollectedFunction<'tcx>>
│
▼ device_codegen::generate_device_code
┌──────────────────────────────────────────────────┐
│ rustc_internal::run(tcx, || { │
│ stable = rustc_internal::stable(internal) │ ← 关键解耦
│ mir_importer::run_pipeline(stable_functions) │
│ }) │
└──────────────────────────────────────────────────┘
│
▼
进入下一站:mir-importer 翻译 MIR(下一篇讲)13. 一句话总结
cuda-oxide 用命名约定 + 魔数后缀代替属性来识别 kernel/device 函数。
#[kernel]proc macro 把用户函数名前缀加上cuda_oxide_kernel_246e25db_(其中246e25db是sha256("cuda_oxide_ + rust")截前 8 字符),collector 在 CGU 里按子串匹配找出所有 kernel 入口,然后 worklist BFS 遍历调用图收集所有可达的 device 函数。最后通过rustc_internal::run进入 stable MIR 上下文,把 internalInstance<'tcx>转成 stableInstance,让 mir-importer 跟 rustc 内部解耦——后续整个 IR 翻译链路都不带'tcx生命周期。