死代码消除

机器学习 (machine learning)模型通常包含在最终执行中不起作用的操作。这些多余的节点可能源于代码中遗留的调试语句、从不执行的条件分支，或训练逻辑中在推理 (inference)时无关紧要的残留。在编译器设计中，移除这些不必要操作被称为死代码消除（DCE）。

虽然标准软件编译器移除无法访问的代码以减小二进制文件大小，但机器学习编译器执行 DCE 主要为了节省计算资源和内存带宽。计算图中的每个操作都会耗费时间和能量。如果张量操作的结果从未被模型的输出或任何后续节点读取，那么计算它就是资源浪费。

机器学习 (machine learning)图中的死代码来源

神经网络 (neural network)中的死代码通常源于三个方面。首先是用户定义。在试验模型架构时，您可能会定义用于可视化或特征提取的辅助分支，它们并非主要预测路径的一部分。如果这些分支被捕获到中间表示（IR）中，但其输出未被请求，它们就会成为冗余部分。

第二个来源是不同优化阶段的关联。像常量折叠这样的技术常常会作为副产品产生死代码。如果一个条件语句依赖于编译器确定实际为常量的值（例如 if use_dropout:，其中 use_dropout 为 False），那么与真条件关联的整个分支将变得不可达。编译器必须运行 DCE 来清除折叠操作遗留的无用部分。

第三个来源是从训练到推理 (inference)的转变。训练图包含反向传播 (backpropagation)操作、梯度计算和损失函数 (loss function)。当编译模型以进行部署时，反向图是不必要的。虽然大多数框架会明确分离这些，但复杂的自定义层可能会意外保留对训练专用张量的引用。

可达性分析

为了识别哪些节点可以安全移除，编译器会进行活跃性或可达性分析。此过程将计算图视为一个有向无环图（DAG），其中边表示数据依赖关系。

该算法通常按以下步骤运作：

确定根节点： 编译器识别图的“输出”。这些是用户请求的外部张量。
反向遍历： 从根节点开始，编译器沿输入边反向遍历图。在此遍历期间访问的每个节点都标记 (token)为“活跃”。
清理： 图中在此遍历期间未被标记为活跃的任何节点都被视为“死”节点。这些节点不对最终输出做贡献，可以移除。

考虑一个简化的计算图，其中一个旁支计算用于记录的标准差，但该值从未被函数返回。

一个数据流图，显示了活跃节点（绿色）与死节点（灰色）。计算均值的分支源自 ReLU 输出，但未连接到最终图输出，使其成为可消除的对象。

处理副作用

机器学习 (machine learning)编译器中的一个情况是对具有副作用的操作的处理。在纯数据流图中，一个操作只有在其数据被使用时才有意义。然而，某些操作，例如 print() 语句或自定义日志操作符，不会产生被数学管道使用的张量输出，但会执行对用户或系统可见的操作。

如果编译器激进地移除所有没有数据后继的节点，它可能会移除这些日志操作。为了避免这种情况，IR 通常支持将特定节点标记 (token)为“非纯”或具有“副作用”。活跃性分析算法将这些节点视为辅助根。如果一个节点有副作用，它会被标记为活跃，而且重要的是，它的所有输入依赖也都会被标记为活跃。

递归优化循环

死代码消除很少作为单次遍历运行。它通常是迭代优化循环的一部分。这是因为移除一个死节点可能会使其输入没有剩余消费者，从而导致它们也成为死节点。

考虑以下操作序列：

\begin{align*} t_1 &= x + y \\ t_2 &= t_1 \times 2 \\ t_3 &= \text{ReLU}(t_2) \\ \text{返回 } & t_1 \end{align*}

在初步分析中，根是返回值 $t_1$ 。

$t_1$ 是活跃的，因为它就是输出。
$x$ 和 $y$ 是活跃的，因为 $t_1$ 依赖它们。
$t_3$ 是死的，因为它不是输出，并且没有活跃节点使用它。
$t_2$ 是死的，因为它的唯一消费者是 $t_3$ ，而 $t_3$ 现在是死的。

根据实现情况，编译器可能通过引用计数在一次反向遍历中识别 $t_3$ 和 $t_2$ ，或者可能需要多次遍历，其中消除 $t_3$ 会使 $t_2$ 成为新的消除对象。

这种循环关系在与其他优化结合使用时特别有意义。例如：

图简化 运行并将一个复杂的子图转换为一个更简单的恒等操作。
DCE 运行以移除现在已断开连接的复杂节点。
常量折叠 在简化后的图上运行，可能将一个分支解析为一个静态值。
DCE 再次运行以移除未选择的分支。

对内存分配的影响

DCE 的好处不仅限于跳过算术指令。最显著的收益常常来自内存管理。在上面的例子中，如果 $t_3$ 和 $t_2$ 被消除，编译器就不需要为这些张量分配中间缓冲区。

对于大型神经网络 (neural network)，中间激活图可以消耗数千兆字节的显存 (VRAM)。通过仔细裁剪图，编译器减少了峰值内存占用，可能允许更大的批次大小或更大的模型适应设备。当为内存限制严格的边缘设备优化时，这尤其有意义。

在编译器栈中的实现

在诸如 TVM 或 XLA（加速线性代数）等框架中，DCE 作用于该栈的特定中间表示（例如，Relay IR 或 HLO）。

检查此阶段之前和之后的 IR 时，您会看到更清晰的结构。

DCE 之前:

fn (%x: Tensor[(10, 10), float32]) {
  %0 = add(%x, 1.0f);
  %1 = multiply(%0, 2.0f);  /* 死代码: %1 从未被使用 */
  %2 = relu(%0);
  %2
}

DCE 之后:

fn (%x: Tensor[(10, 10), float32]) {
  %0 = add(%x, 1.0f);
  %2 = relu(%0);
  %2
}

请注意，%1 及其相关计算被完全移除。依赖 %0 保留，因为它为 %2 所需，而 %2 是函数返回值。

尽管理论上很简单，死代码消除是一个基础性阶段，它确保后续复杂转换的效率。它充当图优化阶段的垃圾回收器，确保硬件只执行为产生预期预测而绝对必需的计算。

这部分内容有帮助吗？