计算图常被视为有向无环图 (DAG) 进行分析。尽管 DAG 能有效表示简单的前馈网络,但它们难以描绘现代架构(如循环神经网络 (RNN))或具有数据依赖分支的模型中存在的动态控制流的复杂情况。为有效处理这些结构并实现积极的优化,现代深度学习编译器采用了静态单赋值 (SSA) 形式。SSA 是一种中间表示特性,要求每个变量只能被赋值一次,并且在使用前必须定义。此约束将复杂的依赖链转化为明确的图结构,使数据流清晰无疑。在机器学习环境中,SSA 搭建起了神经网络数学定义与硬件所需命令式指令之间的桥梁。不变性约束在标准命令式编程中,变量充当可变容器。变量 $x$ 在某个时刻可存储张量,在另一个时刻可存储标量。这种可变性使编译器分析变得更难,因为要确定 $x$ 的值需要追踪整个执行路径。SSA 强制执行不变性。编译器不会修改变量,而是创建一个新版本。考虑一个简单的累加操作,这常出现在优化循环或状态更新中。命令式风格:x = 1 x = 2 y = x + 3在此片段中,$y$ 的值取决于对 $x$ 的哪次赋值是最后发生的。编译器必须分析控制流以解决此问题。SSA 形式:x_1 = 1 x_2 = 2 y_1 = x_2 + 3这里,$x_1$ 和 $x_2$ 是不同的符号。 $y_1$ 对 $x_2$ 的依赖是显式的,并直接编码在变量名中。这使得“定义-使用”链变得简单(定义-使用链)。编译器可以立即确定 $x_1$ 在 $y_1$ 的计算中未使用,并将其标记为死代码消除 (DCE) 的候选,而无需复杂的连通性分析。使用 Phi 函数处理控制流当控制流分叉和汇合时,严格的单赋值规则面临挑战。如果一个变量在条件语句的 then 分支和 else 分支中被赋予不同的值,那么在合并点处该变量的值会变得模糊不清。考虑一个使用显式分支实现的修正线性单元 (ReLU):$$ y = \begin{cases} x & \text{若 } x > 0 \ 0 & \text{否则} \end{cases} $$在 SSA 中,我们不能简单地在两个分支中都给 $y$ 赋值,因为这违反了单赋值规则。为解决此问题,SSA 引入了 $\phi$ (Phi) 函数。一个 $\phi$ 函数存在于控制流图的合并点。它根据执行到达该点的路径来选择正确的值。带 Phi 节点的 SSA:// 基本块 1 x_0 = input() condition_1 = x_0 > 0 if condition_1 goto Block 2 else goto Block 3 // 块 2 (真路径) y_1 = x_0 goto Block 4 // 块 3 (假路径) y_2 = 0 goto Block 4 // 块 4 (合并) y_3 = phi(y_1, y_2) return y_3$\phi$ 节点是一个特殊指令,它告诉编译器:“如果我们从块 2 到达,则值为 $y_1$。如果我们从块 3 到达,则值为 $y_2$。”这种结构对 MLIR (多级中间表示) 和 TVM Relay 等框架非常重要。它使编译器能够表示控制流图 (CFG),同时保持函数式数据依赖的好处。digraph G { rankdir=TB; node [shape=box, style=filled, fontname="Helvetica", fontsize=10, color="#dee2e6"]; edge [fontname="Helvetica", fontsize=9, color="#868e96"]; subgraph cluster_0 { label = "SSA 中的控制流"; style = filled; color = "#f8f9fa"; fontname = "Helvetica"; start [label="入口\nx_0 = 输入", fillcolor="#e7f5ff", color="#74c0fc"]; cond [label="x_0 > 0?", shape=diamond, fillcolor="#fff3bf", color="#fcc419"]; true_br [label="真分支\ny_1 = x_0", fillcolor="#d3f9d8", color="#40c057"]; false_br [label="假分支\ny_2 = 0", fillcolor="#ffe3e3", color="#fa5252"]; merge [label="合并块\ny_3 = phi(y_1, y_2)", fillcolor="#eebefa", color="#be4bdb"]; } start -> cond; cond -> true_br [label="真"]; cond -> false_br [label="假"]; true_br -> merge; false_br -> merge; }逻辑流展示了 Phi 节点如何从发散的控制路径中解析变量定义。块参数与 Phi 节点理论上的 SSA 模型使用 $\phi$ 节点,但现代编译器架构常通过“块参数”来实现这一点。这在 MLIR 中尤为明显。块内没有特殊的 $\phi$ 指令,而是将值作为参数传递给块本身,类似于函数参数。以上面的例子来说,块 4 将定义为 Block4(y_arg)。块 2 通过 br Block4(y_1) 跳转到它,块 3 通过 br Block4(y_2) 跳转到它。这种方法结构上更清晰,并简化了图转换的实现,因为数据流模仿函数调用,而非需要对虚指令进行特殊处理。用于张量内存优化的 SSA在深度学习编译器中,SSA 不仅用于逻辑优化;它对内存管理也很关键。张量占用大量内存,在 VRAM 有限的 GPU 上,为每个中间结果 ($x_1, x_2, \dots$) 分配新内存是不可行的。但是,SSA 形式为每个张量提供了精确的“活跃”区间。定义: 张量创建的点。最后使用: SSA 图中引用此版本张量的最终指令。由于 SSA 保证 $x_1$ 永远不会被重新定义,编译器可以执行就地内存重用。一旦执行通过 $x_1$ 的最后使用点,其物理内存缓冲区可以立即被回收并分配给后续变量,例如 $z_5$。这种优化完全依赖于 SSA 提供的显式依赖链。没有 SSA,要证明内存缓冲区可以安全覆盖,就需要昂贵的别名分析。高级 IR 中的函数式 SSATVM Relay 等框架采用 SSA 的函数式变体。在这种方法中,整个图是一个表达式。函数是第一类公民,闭包可以捕获变量。这与 LLVM 中指令级的 SSA 不同。函数式 SSA 将神经网络建模为一组嵌套表达式,而非指令序列。$$ \text{let } %1 = \text{conv2d}(%data, %weight) $$ $$ \text{let } %2 = \text{bias_add}(%1, %bias) $$ $$ \text{relu}(%2) $$这里,% 符号通常表示 IR 中的虚拟寄存器或唯一标识符。函数式 SSA 的作用域规则使编译器能够执行 Lambda 提升和闭包转换,这些技术借鉴自函数式语言编译器(如 Haskell 或 OCaml),以优化递归 RNN 或具有动态序列长度的 Transformer 等动态模型的执行。SSA 启用的优化过程采用 SSA 形式可以实现特定的编译器过程,否则这些过程难以实施:稀疏条件常量传播 (SCCP): SSA 允许编译器在图和分支中传播常量值。如果张量形状是常量,它可以通过通用操作符传播,从而针对该特定形状进行代码优化。循环不变代码外提 (LICM): 由于依赖关系是明确的,识别循环内部不依赖于循环变量的操作就变成了图遍历问题。这些操作可以从循环中提升出来,以减少计算开销。公共子表达式消除 (CSE): 在 SSA 中,如果两条指令对相同的输入变量(版本号匹配)执行相同的操作,它们必定会产生相同的结果。编译器可以计算一次,并用第一个结果替换所有后续使用。通过将计算图转换为 SSA 形式,深度学习编译器从“层”的粗略视图转向数据和控制流的细粒度视图。这种细粒度是高级内核优化和特定硬件代码生成技术的前提,我们将在后续章节中介绍这些技术。