图优化：公共子表达式消除CSE

可视化演示

三步展示XLA CSE优化：HloCSE哈希分桶+精确比较的等价检测、重复子表达式只算一次并共享、以及非确定性/副作用指令的排除约束。

概念定义

回到后厨场景。厨师长翻看午市菜单，发现三台灶都在做“剁蒜泥”这同一个动作。每台灶对应计算图中的一个算子节点，而“剁蒜泥”这个动作就是“公共子表达式”——它在图的不同位置重复出现，且操作相同、输入相同。在AI编译器的术语里，计算图是一张有向无环图（DAG）：节点是算子（如矩阵乘法、加法、激活函数），边是张量数据流。当某个子图结构在图中出现多次，且对应的输入张量来源一致时，它就是一个公共子表达式，编译器可以只保留一份计算。

CSE的优化目标很明确：消除冗余计算，减少算子执行次数。假设图中两处分别计算 a×b，第一次用到结果后，第二次完全可以复用同一份结果。正如备菜间一次性剁好一大碗蒜泥，两台灶随取随用，不再各自重复劳动。

这就是CSE的“备菜哲学”：一次准备，全店共享。理解了概念，接下来我们看后厨如何快速认出“哪两碗菜是同一份活儿”。

识别原理

帮厨每天要面对几十种备菜需求，如何一眼认出“蒜蓉”和“姜末”不是同一件事？靠的是标签。CSE识别公共子表达式的核心手段是哈希表与值编号（Value Numbering）。每个表达式根据“操作符 + 左右操作数的编号”计算出一个唯一的哈希值，这个哈希值就是表达式的“编号”。当两个表达式的哈希值相同，编译器就判定它们等价，可以合并。

具体做法是：遍历计算图时，为每个节点分配一个值编号。对于二元运算 e = op(l, r)，其值编号 v(e) 由操作符 op 与左右操作数的值编号组合后哈希得到。由于 a×b 的两个操作数编号相同，两处 a×b 得到的哈希值必然相同——不管它们在图中相隔多远。为了识别更隐蔽的重复，编译器还会利用代数恒等式做“标准化”，比如 a×b 与 b×a 视为同一表达式（交换律）， a+0 规范化为 a（单位元），从而让哈希判断更聪明。

哈希值匹配只是“认出重复”，但真正要安全地合并，还要回答一个问题：第二次使用时，第一次的结果是否一定还“活着”？这就引出了算法的全局视野——支配关系。

算法流程

备菜间准备的蒜蓉，必须确保在它被使用之前不会被打扫卫生的阿姨倒掉。CSE同样面临“生命周期”问题：合并两个相同表达式的前提是，第一个计算点必须能安全地“主宰”第二个使用点。这就是支配关系（Dominance）的由来。我们称节点 d 支配节点 n，如果从入口到 n 的所有路径都必经 d。只有第一个表达式所在节点支配第二个表达式所在节点时，第二个节点才能确定地复用第一个节点的结果——否则控制流可能绕开了第一个计算点，导致使用时变量未定义。

编译器按拓扑序遍历基本块，在每个基本块内维护一个局部哈希表，记录“当前可见的表达式编号”。遇到新表达式先查表：命中则用旧编号重写当前节点；未命中则插入表。对于跨越基本块的情况，借助支配树提前将表达式提升到支配者节点，从而让CSE从局部扩展到全局。整个流程是：遍历图 → 查哈希匹配 → 用支配关系验证安全性 → 重写计算图。

有了哈希匹配和支配关系双保险，CSE就能放心地“一份备菜，多处共享”。但这个优化并非没有代价。

性能影响

备菜共享节省了剁蒜的时间，但代价是那一大碗蒜蓉得一直放在台面上占地方。CSE同样存在权衡：它减少了重复计算（省算力、省延迟），但提取出的公共子表达式会产生中间张量，这些张量可能长时间驻留在寄存器或显存里，增加内存压力。尤其在现代AI芯片上，有时“重新算一遍”反而比“从远处读取中间结果”更快——因为访存开销可能大于浮点运算开销。

因此，CSE并非总是净收益，编译器需要启发式地评估：如果公共表达式出现的次数很多、单次计算代价很高，合并绝对划算；如果只出现两次且代价低廉（如一次加法），而中间结果又需要跨多个算子保持活跃，那么不如不合并。性能工程师常用的分析指标是“净节省操作数”：

理解了代价模型，再看AI编译器里的真实场景：现代深度学习模型动辄几百个算子，CSE在其中扮演着什么角色？

AI编译实践

如果把深度学习模型比作一场满汉全席，几十道菜共用高汤、酱汁这些“半成品”，那么XLA、PyTorch编译器等AI编译器就是统筹全局的总厨。XLA在将TensorFlow或JAX模型编译为HLO（High Level Optimizer）中间表示后，会执行多轮Pass优化，其中就包含HLO级别的CSE——把重复的卷积、矩阵乘、归一化子图合并成公共算子。在残差网络（ResNet）中，x + f(x) 的分支结构里，x 被反复引用，CSE确保这个激活值只计算一次，大幅减少冗余内核启动。

PyTorch 2.x中的TorchInductor编译器同样通过CSE优化算子融合后的计算图，减少中间张量的分配与拷贝。不过动态图（Eager模式）下，模型边运行边构建图，CSE难以执行——因为你永远不知道下一步会算什么。这也是编译器工程师推动“图模式”（torch.compile）的原因：把动态执行收敛为静态图，CSE等优化才有用武之地。在手机端侧推理引擎（如TFLite、NCNN）中，CSE还被用于减少内存带宽消耗，让低功耗设备跑起大模型。

从框架到端侧，CSE贯穿了AI编译的各个环节。那么在实际工程项目中，工程师如何落地一个CSE Pass？

工程实战

备菜流程不是写一次就完事，主厨会根据翻台率、菜品种类持续调整SOP。CSE在编译器工程中同样是一个可插拔的Pass，通常运行在目标无关的IR优化管线中。实现一个CSE Pass，基本步骤包括：遍历IR中的基本块 → 维护表达式哈希表 → 匹配成功后用第一个结果的SSA值替换后续使用 → 最后清理无效的冗余指令。但在真实工程里，CSE通常与死代码消除（DCE）配合：CSE合并后会产生一些“没人再用的旧计算”，DCE负责把这些死代码删除，两兄弟协同作战，才能最大化收益。

调试CSE是件细致活。工程师常用 print-ir-after-all 之类开关逐Pass dump IR，对比优化前后的计算图；出现浮点精度差异时，要检查是否错误合并了本不该合并的表达式（例如涉及非确定性算法的算子）。基准测试方面，通常用GPU内核启动次数、FLOPs、端到端延迟三个指标量化CSE的收益。整个优化管线可以形式化表示为：

至此，从后厨到编译器，CSE的全貌已经展开。它不只是“省一次乘法”的小技巧，更是AI编译器优化管线中承上启下的基石。

AI 应用场景

Python 演示

class CSEPass:    def __init__(self):        self.expr_map = {}        self.next_id = 0     def canonicalize(self, op, lhs, rhs):        # 利用交换律标准化：a*b 与 b*a 视为同一个表达式        if op == 'mul' and lhs > rhs:            lhs, rhs = rhs, lhs        return (op, lhs, rhs)     def run(self, instructions):        rewrite = {}        new_ins = []        for ins in instructions:            op, lhs, rhs, dst = ins            key = self.canonicalize(op, lhs, rhs)            if key in self.expr_map:                # 哈希命中：用已有的值编号重写此指令的目标                rewrite[dst] = self.expr_map[key]            else:                # 新表达式：分配值编号并保留指令                new_id = 'v' + str(self.next_id)                self.next_id += 1                self.expr_map[key] = new_id                new_ins.append((op, lhs, rhs, new_id))                rewrite[dst] = new_id        return new_ins, rewrite  def main():    prog = [        ('mul', 'a', 'b', 't1'),        ('mul', 'a', 'b', 't2'),        ('add', 't1', 'c', 'x'),        ('add', 't2', 'd', 'y'),    ]    cse = CSEPass()    new_prog, rewrite = cse.run(prog)    for ins in new_prog:        print(ins)    print('rewrite:', rewrite)  if __name__ == '__main__':    main()

1. if op == 'mul' and lhs > rhs:利用交换律对乘法操作数排序，让 a*b 与 b*a 哈希到同一槽位。
2. key in self.expr_map查询哈希表，判断当前表达式是否已出现过。
3. rewrite[dst] = self.expr_map[key]将重复表达式的目标变量重写为首次计算的结果编号。
4. new_ins.append((op, lhs, rhs, new_id))仅保留首次出现的表达式指令，其余重复指令被丢弃。
5. ('mul', 'a', 'b', 't1')输入程序用四元组表示：算子、左操作数、右操作数、目标变量。

常见误区

面试题