ML 中的静态单赋值

计算图常被视为有向无环图 (DAG) 进行分析。尽管 DAG 能有效表示简单的前馈网络，但它们难以描绘现代架构（如循环神经网络 (neural network) (RNN)）或具有数据依赖分支的模型中存在的动态控制流的复杂情况。为有效处理这些结构并实现积极的优化，现代深度学习 (deep learning)编译器采用了静态单赋值 (SSA) 形式。

SSA 是一种中间表示特性，要求每个变量只能被赋值一次，并且在使用前必须定义。此约束将复杂的依赖链转化为明确的图结构，使数据流清晰无疑。在机器学习 (machine learning)环境中，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$ 的计算中未使用，并将其标记 (token)为死代码消除 (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)，同时保持函数式数据依赖的好处。

逻辑流展示了 Phi 节点如何从发散的控制路径中解析变量定义。

块参数 (parameter)与 Phi 节点

理论上的 SSA 模型使用 $\phi$ 节点，但现代编译器架构常通过“块参数”来实现这一点。这在 MLIR 中尤为明显。块内没有特殊的 $\phi$ 指令，而是将值作为参数传递给块本身，类似于函数参数。

以上面的例子来说，块 4 将定义为 Block4(y_arg)。块 2 通过 br Block4(y_1) 跳转到它，块 3 通过 br Block4(y_2) 跳转到它。这种方法结构上更清晰，并简化了图转换的实现，因为数据流模仿函数调用，而非需要对虚指令进行特殊处理。

用于张量内存优化的 SSA

在深度学习 (deep learning)编译器中，SSA 不仅用于逻辑优化；它对内存管理也很关键。张量占用大量内存，在 VRAM 有限的 GPU 上，为每个中间结果 ( $x_1, x_2, \dots$ ) 分配新内存是不可行的。

但是，SSA 形式为每个张量提供了精确的“活跃”区间。

定义： 张量创建的点。
最后使用： SSA 图中引用此版本张量的最终指令。

由于 SSA 保证 $x_1$ 永远不会被重新定义，编译器可以执行就地内存重用。一旦执行通过 $x_1$ 的最后使用点，其物理内存缓冲区可以立即被回收并分配给后续变量，例如 $z_5$ 。这种优化完全依赖于 SSA 提供的显式依赖链。没有 SSA，要证明内存缓冲区可以安全覆盖，就需要昂贵的别名分析。

高级 IR 中的函数式 SSA

TVM Relay 等框架采用 SSA 的函数式变体。在这种方法中，整个图是一个表达式。函数是第一类公民，闭包可以捕获变量。这与 LLVM 中指令级的 SSA 不同。

函数式 SSA 将神经网络 (neural network)建模为一组嵌套表达式，而非指令序列。

\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 形式，深度学习 (deep learning)编译器从“层”的粗略视图转向数据和控制流的细粒度视图。这种细粒度是高级内核优化和特定硬件代码生成技术的前提，我们将在后续章节中介绍这些技术。

这部分内容有帮助吗？