mem2reg:从 alloca-load-store 回到 SSA

几乎所有编译器前端都用同一个套路:每个 local 变量翻译成 alloca + load + store。简单粗暴,翻译时不用考虑 phi 节点(SSA 里的汇合节点,下面会解释)。代价是 IR 里塞满冗余的栈访问。mem2reg(memory to register,内存到寄存器)就是把这些冗余消掉、恢复 SSA 形式的优化 —— 经典到 LLVM 把它 hardcode 在标准 pipeline 里、cuda-oxide 直接从 pliron 复用一份。本文讲清楚它做什么、怎么做、什么时候做不了。

0. 几个名词先说清楚

缩写 / 术语	英文全称	中文	含义
SSA	Static Single Assignment	静态单赋值	IR 设计约定:每个值只被赋值一次,任何使用都能直接追溯到唯一定义
alloca	allocate(在栈上分配)	栈分配	LLVM IR 指令,在函数栈帧上开一块空间,返回这块空间的指针
phi(φ)节点	phi node (希腊字母 φ)	汇合节点	SSA 形式里,多条控制流汇合时表达”值从哪条路来”的特殊指令
promotable	promotable(可提升的)	可提升	mem2reg 能处理的 alloca——其地址没逃逸,只被直接 load/store
reaching definition	reaching definition	可达定值	程序某一点处,某变量的”当前值”由哪条赋值指令产生
dominance frontier	dominance frontier	支配边界	控制流图分析概念,告诉你”一个 block 的定义会影响哪些汇合 block”
.local	local state space	本地状态空间	PTX 概念,逻辑上是”线程本地内存”,物理上落在 GPU 外置 DRAM 显存里

1. 问题:alloca-load-store 模型的代价

编译器前端拿到源码,要把 local 变量翻译成 IR。最直白的做法是把每个 local 当作一个栈槽:

源码:                      生成的 IR:
let x = 1;        →        x_slot = alloca i32
                            store 1, x_slot
let y = x + 2;    →        t1 = load x_slot
                            t2 = add t1, 2
                            y_slot = alloca i32
                            store t2, y_slot
return y;         →        t3 = load y_slot
                            return t3

源码 3 行,IR 8 条。这就是 alloca-load-store 模型——简单粗暴,每个变量都对应一个内存位置。

为什么这么挫?因为这样不用考虑 phi 节点。后面这种代码:

let x;
if cond {
    x = 1;
} else {
    x = 2;
}
println!("{}", x);

x 在两个 block 被赋两次。但 SSA(静态单赋值,Static Single Assignment)规定每个值只能定义一次。最后 println! 用的 x 是 1 还是 2?只有运行时根据 cond 才知道。SSA 用一种叫 phi 节点(希腊字母 φ,中文常译”汇合节点”)的特殊指令表达”值从哪条控制流来”:

merge_block:
    %x = phi i32 [1, %then_block], [2, %else_block]
                   ↑ ↑              ↑ ↑
                   值 |来自哪个 block 值 |来自哪个 block

读法:“如果是从 %then_block 来的,%x 是 1;如果是从 %else_block 来的,%x 是 2”。phi 节点只在 block 入口出现,用来给”跨 block 汇合”的 SSA value 命名。

前端如果直接生成 SSA,就要做跨 block 数据流分析才能正确插 phi——工程上很麻烦,还容易出 bug。alloca-load-store 模型把这个难题推后:变量都当成内存,phi 节点不用想,后面让优化 pass(mem2reg)收拾。

2. mem2reg 一句话定义

mem2reg:把那些”行为像寄存器、只是被翻译成内存”的 alloca 消掉,直接用 SSA value 串联,需要的地方自动插 phi 节点。

mem 意思是 memory(那些 alloca 栈槽),reg 意思是 register(SSA value)。“把内存升级回寄存器”是这个名字的字面意思——LLVM 内部就叫这个名字。

跑完后上面那段 IR 变成:

mem2reg 前:                  mem2reg 后:
x_slot = alloca i32           
store 1, x_slot               
t1 = load x_slot              
t2 = add 1, 2                t2 = add 1, 2        ← x 直接换成它的"定义"1
y_slot = alloca i32           
store t2, y_slot              
t3 = load y_slot              
return t3                    return t2            ← y 换成它的定义 t2

8 条 → 2 条,所有冗余 alloca / load / store 全消掉。

3. 什么 alloca 是 “promotable” (可消的)

不是所有 alloca 都能消掉。mem2reg 只处理promotable 的 alloca,判定条件很严:

必要条件	反例(不 promotable)
alloca 的地址只被 load / store 直接用	地址被传给别的函数(fn(&x))、被存进别的指针、被 cast 成其它类型
load / store 的指针操作数就是 alloca 本身	store v, getelementptr(x_slot, 1)(算术后的指针)
store 的 value 不是 alloca 的地址	store x_slot, ptr_to_ptr(把”x 的地址”存起来)

简单说:alloca 的”地址语义”绝对不能被用到。一旦地址跑出来,就有可能被别的代码间接修改,mem2reg 没法保证安全消除。

举几个不能 promote 的例子:

let mut x = 1;
let r = &mut x;      // ← x 的地址被借走,alloca 的地址逃逸到 r
*r = 2;

x 的 alloca 在 IR 里地址被传给 r,mem2reg 不能 promote——因为它没法看出 *r = 2 会修改 x,如果 promote 掉就会得到错误结果。

unsafe { *(&mut x as *mut i32) = 2; }

同样,地址被 cast 出来用——不 promotable。

4. 跨 block:phi 节点怎么插

最容易理解的例子:

let x;
if cond {
    x = 1;
} else {
    x = 2;
}
return x;

前端生成的 alloca 版 IR:

entry:
    x_slot = alloca i32
    br cond, %then, %else

then:
    store 1, x_slot
    br %merge

else:
    store 2, x_slot
    br %merge

merge:
    t = load x_slot
    return t

mem2reg 消掉 x_slot,但 merge 里 load 没法直接换——因为 merge 有两个前驱,x 的”reaching definition”有两个候选(1 来自 then,2 来自 else)。这时候必须插 phi:

entry:
    br cond, %then, %else

then:
    br %merge

else:
    br %merge

merge:
    t = phi i32 [1, %then], [2, %else]    ← 插一个 phi
    return t

mem2reg 自动算出在哪些 block 入口需要插 phi——靠支配边界(dominance frontier)算法。

支配边界这个词解释一下:在控制流图里,如果 block A 一定先于 block B 执行,就说 “A 支配 B”。但如果 A 不再支配 C,但 C 又能从 A 路径走过来——C 就在 “A 的支配边界” 上。支配边界正好就是需要 phi 节点的地方——因为这里有多条控制流汇合,SSA value 需要”选一条来路”。

这是 Cytron 等人(Ron Cytron et al.)1991 年那篇经典论文里提出的 SSA 构造算法核心,LLVM 的 mem2reg 实现照这个套路。

5. 算法骨架

LLVM 的 mem2reg(以及 pliron / MLIR 的同名 pass)算法分三步:

① 扫所有 alloca,判定哪些 promotable
   (检查每个 use 是不是符合上一节的条件)
                ↓
② 对每个 promotable alloca:
   ├── 算出"用到 x 的 load 在哪些 block"
   ├── 用支配边界算出"在哪些 block 入口需要插 phi"
   └── 插 phi(节点先放空,稍后填操作数)
                ↓
③ 改写 IR:
   ├── 每个 store 不再写内存,直接成为某个 SSA value 的"定义"
   ├── 每个 load 改成"读 reaching definition"
   └── phi 的操作数填上每个前驱 block 的 reaching definition
                ↓
所有 alloca / store / load 删掉

第 ② 步的”支配边界”是关键概念(上一节解释了)——它告诉你”一个 block 的定义会影响哪些汇合 block”,从而知道哪里需要 phi。完整算法细节不展开,经典编译器教科书(Cooper & Torczon 的 Engineering a Compiler、Appel 的 Modern Compiler Implementation)都有讲。

6. 在 cuda-oxide 里的位置

cuda-oxide 的 mir-importer 直接用 pliron 提供的 mem2reg pass:

// 在 run_pipeline 里,statement / terminator 翻译完之后
pliron::opts::mem2reg::mem2reg(module_op_ptr, &mut ctx)?;

它跑在 dialect-mir 上,识别的是 mir.alloca / mir.load / mir.store 三个 op。

6.1 跑之前 vs 跑之后

hello_constant 的 entry block 在 docs/run.log 里跑之前:

^block2v1(v0: mir.ptr<si32, ...>):
    v1 = mir.alloca () : ... -> *(*mut i32)      ← 10 个 local 的栈槽
    v2 = mir.alloca () ...
    v3 = mir.alloca () ...
    ... (v4 .. v10 共 10 条 alloca) ...
    mir.store (v1, v0)                            ← 参数 copy 到 _1 槽
    v11 = nvvm.read_ptx_xxx ()                    ← thread::xxx()
    mir.store (v2, v11)                           ← 写 _2 槽
    mir.goto ^block3v1

跑 mem2reg 之后,绝大多数 alloca 都被消掉——v1 这个 out 参数槽只在 entry 被参数 copy 一次,之后所有读都换成直接读 v0。后续 block 里所有 mir.load (v1) 直接换成 v0,v1 这个 alloca 跟着删除。最终保留的可能只有取地址用的几个(printf args 那种)。

6.2 为什么是必需而不是可选

对 GPU 这点尤其重要。没跑 mem2reg 的 IR 翻译成 PTX,所有 alloca 会落进 .local 状态空间——PTX 的 .local 是外置 DRAM,访问比寄存器慢几十倍。

举例:不跑 mem2reg 的版本,*out = 42 这一句会变成:

ld.local.b64  %rd1, [out_slot];     ; 从 .local 读出参数槽
cvta.to.global.u64  %rd2, %rd1;
st.global.b32 [%rd2], 42;

跑了 mem2reg 之后:

ld.param.b64  %rd2, [hello_constant_param_0];   ; 直接从 .param 读
cvta.to.global.u64  %rd1, %rd2;
st.global.b32 [%rd1], 42;

省掉了一次 .local 访问——这一条指令的差距,放到 matmul / reduce kernel 的热循环里就是几十倍性能差距。所以 cuda-oxide 不是”可选优化”地用 mem2reg,而是 lower 之前的必经一步——pipeline 顺序就把它强制塞在 verify 和 lowering 之间。

7. 哪些情况 mem2reg 帮不上

如果你的 kernel 里写了:

let mut x = 0;
let r = &mut x;
unsafe_helper(r);    // ← x 的地址逃逸

x 的 alloca 不会被 promote——逃逸的地址让 mem2reg 不安全。最终生成的 PTX 里还是会有 .local 访问,这一段就是性能损失点。

这是为什么 GPU kernel 风格强调避免取栈变量地址——能用值传递就用值传递,能内联就内联。除了语义考虑,本质就是为了让 mem2reg 跑成功。

8. 历史背景

mem2reg 的核心算法来自 Ron Cytron 等人(Cytron et al.,et al. 是”和其他人”的拉丁缩写)1991 年那篇 “Efficiently Computing Static Single Assignment Form and the Control Dependence Graph”(ACM 期刊 TOPLAS = Transactions on Programming Languages and Systems,第 13 卷第 4 期)——支配边界 + phi 节点放置的高效算法。

LLVM 从最早期就实现了这个 pass,叫 mem2reg(全名 PromoteMemoryToRegister,中文”把内存提升为寄存器”)。所有 LLVM-backed 编译器(Clang、rustc、Swift 等)都默认在 -O0 之上就跑它——因为前端为了写得简单,几乎都依赖它。

pliron / MLIR 把同一个算法重新实现一遍,API 不同但语义一致。cuda-oxide 用 pliron 的版本,效果跟 LLVM 的等价。

9. 一句话总结

几乎所有编译器前端都用 alloca-load-store 模型生成中间 IR——简单但 IR 里塞满冗余。mem2reg 把”地址没逃逸”的 alloca(称为 promotable)消掉,直接用 SSA value 串联,需要的地方靠支配边界算法插 phi 节点。它是 LLVM 标准 pipeline 的固定一环,cuda-oxide 在 dialect-mir 上跑 pliron 的同名 pass——对 GPU 尤其关键,因为没跑的话冗余 alloca 会落进 PTX .local 状态空间(外置 DRAM),性能直接崩。

系列上一篇: Rust 单态化:Instance 类型与惰性实例化