XLA图优化：代数简化让计算图更高效

可视化演示

三步交互演示XLA编译器代数简化核心机制：常量折叠编译时预计算、表达式化简消除冗余算子、优化Pass流水线协同执行。

计算图基础：读懂菜谱才能改菜谱

主厨拿到一份新菜谱，第一件事一定是通读全篇：总共有多少道工序？哪道工序必须在哪道之前？切西红柿必须在打蛋前面吗？不一定，但“先备料再开火”和“一边炒一边找盐”的出菜效率截然不同。读懂工序和依赖，是一切改动的起点。计算图之于 AI 编译器，正是这张菜谱。

深度学习模型的计算过程，就是一张巨大的“菜谱”：计算图（Computation Graph）。在 XLA 里，这张图有一个专门的名字——HLO（High Level Optimizer 的中间表示）。图上的节点是算子，比如 Add、Mul、MatMul、ReLU；边是张量的流动方向，数据从输入节点流入，经过一个个算子，最后到达输出节点。XLA 后面的所有优化，都在这张图上展开。

图优化的第一个前提是“读图”。编译器需要知道三件事：每个节点的输入来自哪里（数据依赖），哪些节点可以并行（节点间没有依赖链），每个中间张量的形状是什么（Shape）。有了这些信息，后续的代数简化、常量折叠、算子融合规则才有用武之地——就像主厨只有看清每道工序的先后关系，才知道哪些步骤可以删、哪些可以提前准备。

就像主厨只有看清每道工序的依赖关系，才能决定先做哪一步、哪些步骤可以提前准备，编译器也只有读懂计算图的结构，才知道哪些节点可以删、可以提前算、可以合并。读图，是后面所有“偷懒技巧”的前提。

代数简化：划掉“加水又倒掉”的步骤

菜谱里最气人的步骤是“加一勺盐，觉得咸，再加水稀释”。这个操作不改变最终味道，却白白占了灶台和厨师的时间。后厨高手看到这种步骤会毫不犹豫划掉。计算图里也大量存在这种“做了等于没做”的算子：乘以 1、加 0、乘以 0。

代数简化（Algebraic Simplification）就是 XLA 里专门做这件事的优化 pass（可以理解为一个“专门挑冗余步骤的审核员”）。它的工作方式非常朴素：在图上寻找符合特定恒等模式的子图，然后替换成更简单的等价子图。单个规则看上去像是“小聪明”，但在动辄成千上万个节点的模型里，冗余算子往往是以“万”为单位出现的。

除了恒等消去，代数简化还包括两类常用技巧。一类是强度削减（strength reduction）：把昂贵的算子换成便宜的等价算子，比如把 x 的平方改成 x×x，把除以常数改成乘以它的倒数；另一类是公共子表达式消除（CSE）：如果图里两处计算的是同一个子表达式，只算一次，结果复用——就像熬一锅高汤，多道菜共用，而不是每道菜各熬一锅。

划掉“加水又倒掉”，菜品味道一字不差，出菜速度却明显变快。代数简化的精髓正是如此：语义严格保持不变，计算成本悄悄下降。对编译器而言，这种“零风险”的优化永远排在性能优化金字塔的最底层——它不依赖硬件，却能让后面的每一层优化受益。

常量折叠：提前把配料称好

高峰期出餐，主厨一定会在开市前把盐、糖、淀粉按小份装好。顾客点单时直接倒进锅里，而不是现场称量——现场称量不仅慢，还容易忙中出错。计算图里也经常出现这种“现场称量”的浪费：3+2、0.5×2 这类子表达式，输入全是常量，与任何动态数据无关，却要在每次前向传播里重复计算。

常量折叠（Constant Folding）的做法非常直接：在编译期把常量子表达式算出来，替换成新的常量节点。比如 (3+2)×x，编译器在编译阶段算出 3+2=5，把整棵子图压成一个常量节点 5。要知道，模型训练一个 epoch 可能要跑几千步，推理服务每秒钟要跑几百次——把"3+2"从每次运行里抹掉，累积收益相当可观。

更有意思的是常量传播（Constant Propagation）：把已知常量的值沿着计算图的边向前推，让下游原本依赖变量的算子也被“感染”成常量子表达式，从而连锁触发更多折叠。常量折叠还有一个隐藏收益——形状推断：当所有维度都在编译期确定后，编译器可以提前分配显存、规划 kernel 调度，而不是运行时手忙脚乱地分配。

提前称好的配料，让出餐高峰期从容不迫。常量折叠省下的，正是模型每一次前向传播里的“现场称量”。这也是为什么编译器工程师常说：能在编译期做完的事，绝不留到运行时。

算子融合：一锅出的智慧

后厨最懂“一锅出”的妙处：西红柿炒蛋，先炒蛋盛出来、再炒西红柿、最后倒回去混合？高手直接一锅搞定。省掉的不仅是来回折腾的时间，还有洗盘子的水——对应到 GPU 上，就是中间张量写回显存再读出来的开销。这个开销往往比计算本身还要昂贵。

算子融合（Operator Fusion）就是把计算图中相邻的、逐元素计算的算子合并成一个自定义 kernel（GPU 上执行的一段程序）。比如 conv + bias + relu 三个算子，原本要启动三次 kernel、产生两个中间张量；融合之后变成一个“融合算子”，中间结果根本不离开片上缓存，直接在寄存器里传递。

这里要澄清一个常见误解：融合不会减少浮点运算次数（FLOPs 一点没少），它省的是两笔工程开销。第一笔是访存——中间张量不用写回显存再读出来，省掉了最贵的全局内存往返；第二笔是 kernel 启动次数——GPU 上每次启动 kernel 都有固定开销，融合把三次启动合并成一次。但融合不是无脑合并：涉及 shape 改变的算子（如 reshape、pad）必须谨慎，强行融合可能让访存模式变差，反而更慢。融合规则的精髓，正是判断哪些工序值得一锅出，哪些必须分锅。

一锅出不代表所有菜都该一锅出——炖汤和油炸必须分开，否则互相破坏口感。算子融合也是如此，它是一门“看菜下锅”的手艺：选对能融合的算子组合，性能起飞；选错组合，反而得不偿失。

优化Pass调度：后厨动线设计

后厨动线是门学问：备料台放哪、灶台怎么排、出菜口在哪，都有讲究。更重要的是顺序——先切配还是先煲汤？先备料再开火，还是边炒边备料？顺序错了，再好的厨师也会手忙脚乱。编译器的优化过程同样面临这个问题，而它的解决方案是“优化 Pass 调度”。

编译器里的优化不是一口气完成的，而是按一个个“优化 pass”逐步执行的。每个 pass 是一个独立的小优化器，按顺序作用在计算图上：前一个 pass 的输出图，是后一个 pass 的输入图。XLA 的优化管线里，代数简化、常量折叠、算子融合、死代码消除等 pass 被编排成一个序列，各司其职、依次执行。

Pass 的顺序大有讲究。比如常量折叠通常要在恒等消去之前：折叠产生的新常量，可能正好触发下游的“乘 1 消去”；而代数简化把图精简之后，算子融合的模式匹配又变得更简单。更复杂的是“机会触发”——某些 pass 只在检测到特定图模式时才开工，需要前面的 pass 先把模式“造”出来。因此编译器通常不会只跑一遍，而是循环迭代整个 pass 集合，直到图不再变化——这个稳定状态叫“不动点”。

精妙的动线设计让整条出菜流水线顺滑流畅。Pass 调度也是同样的道理：顺序对了，各个优化规则互相成就；顺序乱了，某个 pass 可能把图改得让后面所有 pass“无从下手”——这就是为什么编译器工程师把 pass 管线的编排看得和 pass 本身一样重要。

性能评估：出菜计时器说了算

主厨改菜谱从来不是凭感觉，而是靠计时——改前一桌菜 3 分钟出齐，改后 2 分钟，数据说话。如果计时器显示更慢了，再“有道理”的改动也得回滚。图优化同样如此，任何优化规则都不能只停留在“看起来更快”，必须用性能数据证明自己。

性能评估常用的指标有四类。端到端延迟：处理一个样本的耗时，是最直观的用户体验指标；吞吐：每秒能处理的样本数，对推理服务至关重要；显存峰值：优化前后占用的显存，融合经常能带来惊喜；还有 kernel 启动次数和浮点运算量，用来定位瓶颈到底在调度层还是计算层。

XLA 提供了专业的剖析工具，最常用的是导出 pass 前后的 HLO 图（通过 XLA_FLAGS 里的 --xla_dump_to 参数），工程师可以逐 pass 对比图的变化，精确看到每个优化规则到底删了什么、合并了什么。最终的综合指标是加速比（Speedup）：优化前耗时除以优化后耗时，大于 1 说明优化有效，越大越好。

计时器不会说谎。优化是“看起来更快”还是“真的更快”，最终都要看这个比值。这也是为什么成熟的编译器团队会把性能回退测试（perf regression test）写进 CI——每次改动一个 pass，都要用一组标准模型跑一遍加速比，防止优化 A 给优化 B“帮了倒忙”。

XLA与PyTorch：两位大厨的同款理念

不同餐厅的菜系天差地别，但主厨的简化理念是相通的：删冗余、预配料、并工序。AI 编译器世界也是如此——XLA 和 PyTorch 编译器就像两位风格不同的大厨，用的却是同一套“后厨心法”。理解了代数简化这套理念，你在任何编译器里都能快速找到对应物。

XLA 的编译管线清晰而经典：前端把 TensorFlow 或 JAX 模型转换成 HLO 中间表示，然后在 HLO 上依次运行一串优化 pass——代数简化、常量折叠、算子融合都在其中；优化完的 HLO 再交给 LLVM，生成针对 GPU/TPU 的机器码。PyTorch 2.0 的 torch.compile 走的也是同样的路：Dynamo 负责捕获 Python 代码里的计算图，Inductor 在图上做算子融合，并生成高效的 Triton kernel。

两位大厨有一个共同的“头疼事”：动态 Shape。如果张量形状在运行时才确定，常量折叠和形状推断在编译期就难以发挥作用，内存预分配也无从谈起——相当于菜谱上写着“盐适量、水少许”，主厨只能临场发挥。这也是为什么实际生产中，工程师常常用 padding 把输入形状固定下来，牺牲一点显存换来编译优化的全面激活。

不管哪家大厨，拿手的都是“把菜谱打磨到极致”的功夫。XLA 和 PyTorch 编译器各有所长：XLA 适合静态图、追求极致性能，torch.compile 兼顾动态 Python 生态和编译优化。但它们的核心心法完全一致——代数简化，删冗余、预配料、并工序，让计算图更高效。

AI 应用场景

Python 演示

class Expr:    """计算图表达式节点: 仅支持常量, 变量, 加法和乘法"""    def __init__(self, op, *args, value=None):        self.op = op          # const / var / add / mul        self.args = args      # 子节点        self.value = value    # 常量节点的数值     def __repr__(self):        if self.op == 'const':            return str(self.value)        if self.op == 'var':            return self.args[0]        sym = {'add': '+', 'mul': '*'}        return '(' + str(self.args[0]) + ' ' + sym[self.op] + ' ' + str(self.args[1]) + ')'     def count_nodes(self):        """统计子图节点总数, 用于观察优化效果"""        if self.op in ('const', 'var'):            return 1        return 1 + self.args[0].count_nodes() + self.args[1].count_nodes()  def simplify(expr):    """代数简化主函数: 自底向上折叠常量, 消去恒等操作"""    if expr.op not in ('add', 'mul'):        return expr    left = simplify(expr.args[0])    right = simplify(expr.args[1])    # 常量折叠: 两个子节点都是常量, 直接算出结果    if left.op == 'const' and right.op == 'const':        if expr.op == 'add':            return Expr('const', value=left.value + right.value)        return Expr('const', value=left.value * right.value)    # 恒等消去: x + 0 = x, x * 1 = x, x * 0 = 0    if expr.op == 'add' and right.op == 'const' and right.value == 0:        return left    if expr.op == 'mul' and right.op == 'const':        if right.value == 1:            return left        if right.value == 0:            return Expr('const', value=0)    # 构造一个新的, 左右子树均已简化的节点    return Expr(expr.op, left, right)  def main():    # 构造表达式: (x * 1) + (3 + 2)    x = Expr('var', 'x')    raw1 = Expr('add', Expr('mul', x, Expr('const', value=1)),                Expr('add', Expr('const', value=3), Expr('const', value=2)))    print('原始表达式:', raw1, '节点数:', raw1.count_nodes())    opt1 = simplify(raw1)    print('简化后表达式:', opt1, '节点数:', opt1.count_nodes())     # 再试: (x + 0) * (5 * 0), 期望直接折叠成常量 0    raw2 = Expr('mul', Expr('add', x, Expr('const', value=0)),                Expr('mul', Expr('const', value=5), Expr('const', value=0)))    print('原始表达式2:', raw2, '节点数:', raw2.count_nodes())    opt2 = simplify(raw2)    print('简化后表达式2:', opt2, '节点数:', opt2.count_nodes())  if __name__ == '__main__':    main()

1. def simplify(expr):简化入口，自底向上递归处理每个节点，对应 XLA 中“先化简子树再处理父节点”的遍历方式。
2. left = simplify(expr.args[0])先简化左子树，保证子问题解决后再处理当前节点，这是“自底向上”的关键。
3. if left.op == 'const' and right.op == 'const':两个子节点都是常量，触发常量折叠，对应 XLA 的 Constant Folding pass。
4. if expr.op == 'add' and right.op == 'const' and right.value == 0:加零恒等消去，x+0 直接返回 x，对应“划掉加水又倒掉的步骤”。
5. if expr.op == 'mul' and right.op == 'const':乘一消去和乘零吸收两条规则汇聚在这里：乘以 1 返回原节点，乘以 0 返回常量 0。
6. print('简化后表达式:', opt1, '节点数:', opt1.count_nodes())打印优化前后的节点数，用量化数据验证“图变简洁了”，对应性能评估中的指标对比。

常见误区

面试题