Rust 单态化:Instance 类型与惰性实例化

“Rust 把 Vec<T> 单态化成 Vec<i32>、Vec<String> 各一份”——常见说法。但 rustc 内部到底什么时候、怎么做这件事?答案不是直觉的”一开始就给每个泛型生成 N 份 MIR”,而是一种惰性实例化机制,核心是 Instance 这个类型。这篇文章把它讲透。

1. 单态化的传统直觉 vs rustc 的真实做法

朴素理解:看到 Vec<i32>::push 调用,复制一份 Vec::push 的 MIR,把 T 替换成 i32,得到独立的 Vec_i32_push MIR。

rustc 不这么干。它用一种更轻量的方式:保留泛型 MIR 一份,用 Instance(def, args) 这种”定义 + 参数对”来代表实例化,只在 codegen 真正读 MIR 的那一刻才动态替换类型参数。

官方文档里说得直白:

Monomorphization happens on-the-fly and no monomorphized MIR is ever created. codegen and const eval will do all required instantiations as they run.

记住三个关键词:on-the-fly(即时)、no monomorphized MIR is ever created(从不预创建)、as they run(按需)。

2. `Instance` 类型

Instance 是 rustc 表达”一个具体可调用函数”的核心类型:

pub struct Instance<'tcx> {
    pub def: InstanceKind<'tcx>,      // 函数定义(可能是泛型)
    pub args: GenericArgsRef<'tcx>,   // 类型参数替换列表
}

两个字段的角色:

字段	含义	例子
`def`	指向”哪个函数”(`DefId` 或变体)	`Item(Vec::push)`
`args`	这个函数被哪些具体类型实例化	`[i32]`

合起来:Instance { def: Item(Vec::push), args: [i32] } = “Vec<i32>::push 这个具体函数”。

InstanceKind 有多个变体,常见几个:

变体	含义
`Item(DefId)`	普通函数/方法定义
`Intrinsic(DefId)`	编译器内建函数
`Virtual(DefId, usize)`	虚函数调用(从 vtable 来)
`CloneShim(DefId, Ty)`	自动生成的 `Clone::clone` 实现
`VTableShim` / `ReifyShim` / `FnPtrShim`	各种自动生成的胶水代码

所以 Instance 不只表示用户写的函数,也表示编译器自动合成的代码(shim、虚分发、Clone 实现等)。

3. 惰性单态化的设计

源码定义:
fn scale<T: Copy>(x: T) -> T { x }

调用点:
scale::<i32>(1)
scale::<f32>(2.0)

rustc 内部状态:

MIR 表里:
  fn scale<T>(x: T) -> T { x }     ← 唯一一份"泛型 MIR"

Instance 表里:
  Instance { def: Item(scale), args: [i32] }    ← 表达"scale::<i32>"
  Instance { def: Item(scale), args: [f32] }    ← 表达"scale::<f32>"

只有调用 tcx.instance_mir(instance) 时,rustc 才把那份泛型 MIR 取出来,根据 instance.args 把 T 替换成具体类型,临时生成一份单态化 MIR 喂给后续阶段。

这样设计换来什么

节省内存:不存 N 份 MIR,只存一份泛型 MIR + N 个 Instance。Instance 本身轻量(就两个字段)。

统一处理:不管你调的是非泛型函数(args 为空)还是泛型函数(args 有内容),都用 Instance 表示。后续 pass 不用关心”我处理的是泛型还是不是泛型”,只看 Instance。

灵活:可以在编译后期再决定如何实例化(比如 const eval、特化等都晚于一开始做)。

代价:codegen 时需要做参数替换,需要一个明确的”是否完全单态化”检查——下一节。

4. Instance 的三种状态

Instance 在编译过程中可能处于三种状态,只有”完全单态化”才能 codegen。

状态 1:完全泛型(无 args)

// 仅有定义,没有任何调用
fn scale<T>(x: T) -> T { x }

对应 Instance:

Instance {
  def: Item(scale),
  args: [] ← 空!
}

无法 codegen —— T 是什么都不知道,生成不了具体代码。

状态 2:部分单态化(嵌套泛型未解析)

fn outer<T: Clone>(x: T) {
    inner::<T>(x.clone());   // ← T 被原样传给 inner
}

外层 outer::<i32> 被调用时,内层那次 inner::<T> 创建的 Instance:

Instance {
  def: Item(inner),
  args: [T/#0] ← T/#0 是个"类型参数占位符",还不是具体类型
}

也无法 codegen —— 表面看 inner 有 args,但 args 里的 T/#0 还是泛型符号。需要等 outer 被实例化时才会进一步把 T/#0 替换掉。

状态 3:完全单态化(可以 codegen)

scale::<f32>(1.0)

Instance {
  def: Item(scale),
  args: [f32] ← 完全是具体类型
}

可以 codegen。

5. `is_fully_monomorphized` 的两个检查

判断一个 Instance 是否处于状态 3,rustc(以及 cuda-oxide)用一个辅助函数:

pub fn is_fully_monomorphized<'tcx>(
    tcx: TyCtxt<'tcx>,
    instance: Instance<'tcx>,
) -> bool {
    let generics = tcx.generics_of(instance.def_id());

    // 检查 1:args 中是否有未解析的类型参数
    for arg in instance.args.iter() {
        if let Some(ty) = arg.as_type()
            && ty.has_param()
        {
            return false;
        }
    }

    // 检查 2:定义本身有泛型但 args 为空
    if generics.count() > 0 && instance.args.is_empty() {
        return false;
    }

    true
}

两个独立的检查,排除状态 2 和状态 1。

检查 1:args 里是否还有类型参数

for arg in instance.args.iter() {
    if let Some(ty) = arg.as_type()
        && ty.has_param()
    {
        return false;
    }
}

GenericArg 有三种:

pub enum GenericArg<'tcx> {
    Type(Ty<'tcx>),           // 类型参数
    Lifetime(Region<'tcx>),   // 生命周期参数
    Const(Const<'tcx>),       // 常量参数(如 [T; N] 的 N)
}

arg.as_type() 把 GenericArg 转成 Ty(如果不是类型就返回 None)。ty.has_param() 判断这个类型里是否还含有泛型参数符号(T/#0 之类)。

任何一个 args 里的类型还带泛型符号 → 状态 2,返回 false。

注意:常量参数和生命周期参数不参与此检查——它们各自有自己的处理路径。array_len::<5>(...) 这种 const generic 因为 arg.as_type() 返回 None,跳过检查 1。

检查 2:定义有泛型但 args 为空

if generics.count() > 0 && instance.args.is_empty() {
    return false;
}

捕捉的是状态 1。如果一个函数定义里有泛型参数,但 Instance 的 args 是空的,说明这只是泛型定义本身,没被实例化。

为什么会有这种 Instance?因为 CGU 扫描时 rustc 可能把泛型定义本身也作为一个候选项放进来——但它不该被生成代码。检查 2 把这种情况拦下来。

流程

输入: Instance
   │
   ▼
检查 1:遍历 args
        每个 Type 类型的 arg
        有 has_param() 返回 true 的吗?
   │
   ├── 有 → return false(状态 2)
   │
   ▼ 全部都是具体类型
检查 2:def 有泛型参数 && args 为空?
   │
   ├── 是 → return false(状态 1)
   │
   ▼
return true(状态 3,可以 codegen)

6. cuda-oxide 怎么用这个机制

cuda-oxide 的 collector(就是 crates/rustc-codegen-cuda/src/collector.rs 里那块)在两处用 is_fully_monomorphized。

6.1 数 kernel

pub fn count_kernels_in_cgus<'tcx>(
    tcx: TyCtxt<'tcx>,
    cgus: &[CodegenUnit<'tcx>],
) -> usize {
    let mut count = 0;
    for cgu in cgus {
        for (item, _data) in cgu.items() {
            if let MonoItem::Fn(instance) = item
                && is_kernel_function(tcx, instance.def_id())
                && is_fully_monomorphized(tcx, *instance)    // ← 只数完全单态化的
            {
                count += 1;
            }
        }
    }
    count
}

为什么要 is_fully_monomorphized 过滤?因为对泛型 kernel,CGU 里通常同时包含:

泛型定义本身(状态 1,args 空)
每个被实例化的具体版本(状态 3)

不过滤的话,泛型 scale<T> 会被数成一个 kernel,但它生成不了 PTX——T 都不知道是什么。

6.2 收集 device 函数 + 跳过未完全实例化的 callee

collect_device_functions 的 worklist BFS 里,每碰到一个调用点的 callee,都先检查:

if !is_fully_monomorphized(tcx, *instance) {
    // 跳过泛型定义,只处理具体实例化
    eprintln!(
        "[collector] Skipping non-monomorphized kernel: {} (needs type instantiation)",
        name
    );
    continue;
}

跳过的原因跟 6.1 一样——状态 1 / 状态 2 都不能生成 PTX。

6.3 单态化的便利:backend 拿到的 MIR 全是具体类型

当 cuda-oxide 进入 mir-importer 翻译 MIR 时,它调:

let body = func.instance.body()?;

这一步背后调的就是 tcx.instance_mir(instance.def)(stable MIR 包装过的版本)。rustc 在这一刻才把泛型 MIR 取出来,按 instance.args 把所有 T 替换成具体类型,返回单态化后的 Body。

所以 mir-importer 看到的 MIR 完全不带泛型符号——所有类型都是具体的 i32、*mut f32、MyStruct。这就是为什么 cuda-oxide 翻译 MIR 时完全不需要处理泛型替换逻辑。

这个机制也是 rustc backend 写起来比 GCC / clang frontend 容易的原因之一——单态化在 backend 介入之前就由 rustc 帮你做好了。

7. 三个状态在真实 example 里的体现

跑 cargo oxide run hello_constant 时 CGU 扫描的两次结果:

kernel	是否泛型	Instance 状态	出现次数
`hello_constant`	不是泛型	状态 3(args 空,但 generics.count = 0)	1
`hello_kernel`	不是泛型	状态 3(同上)	1

注意非泛型函数的 Instance 也是状态 3——因为 args 是空,generics.count() 也是 0,两个检查都过得了。这是统一处理的好处:非泛型函数 = 没有泛型参数要替换的”特殊单态化”。

如果 example 里有 scale<T>(...) 这种泛型 kernel,扫描结果会变:

Instance	状态	`is_fully_monomorphized`	被收集吗
`scale<T>`(定义本身)	1	false	跳过
`scale<f32>`	3	true	收集
`scale<i32>`	3	true	收集

最终 PTX 里会有 scale_TID_<hashf32>、scale_TID_<hashi32> 这种带哈希后缀的 mangled 名字——这是 cuda-oxide 给单态化版本起的导出名,把不同实例化映射到不同 PTX 符号。

8. 一句话总结

Rust 的单态化是惰性的:rustc 不预先生成单态化 MIR,而是用 Instance(def, args) 这种”定义 + 参数”的轻量结构表示每个具体实例,只在 codegen 真正读 MIR 那一刻才动态替换类型参数。Instance 有三种状态——完全泛型(状态 1)、部分单态化(状态 2)、完全单态化(状态 3),用 is_fully_monomorphized 的两个检查区分。只有状态 3 才能生成 PTX——cuda-oxide 在 count_kernels_in_cgus 和 collect_device_functions 里用这个谓词过滤掉前两种,backend 拿到的 MIR 全是具体类型,完全不用处理泛型替换。

系列上一篇: MLIR 方言的高低层:看意图还是看机制

评论区

评论功能即将上线, 敬请期待。