机器学习模型通常包含在最终执行中不起作用的操作。这些多余的节点可能源于代码中遗留的调试语句、从不执行的条件分支,或训练逻辑中在推理时无关紧要的残留。在编译器设计中,移除这些不必要操作被称为死代码消除(DCE)。虽然标准软件编译器移除无法访问的代码以减小二进制文件大小,但机器学习编译器执行 DCE 主要为了节省计算资源和内存带宽。计算图中的每个操作都会耗费时间和能量。如果张量操作的结果从未被模型的输出或任何后续节点读取,那么计算它就是资源浪费。机器学习图中的死代码来源神经网络中的死代码通常源于三个方面。首先是用户定义。在试验模型架构时,您可能会定义用于可视化或特征提取的辅助分支,它们并非主要预测路径的一部分。如果这些分支被捕获到中间表示(IR)中,但其输出未被请求,它们就会成为冗余部分。第二个来源是不同优化阶段的关联。像常量折叠这样的技术常常会作为副产品产生死代码。如果一个条件语句依赖于编译器确定实际为常量的值(例如 if use_dropout:,其中 use_dropout 为 False),那么与真条件关联的整个分支将变得不可达。编译器必须运行 DCE 来清除折叠操作遗留的无用部分。第三个来源是从训练到推理的转变。训练图包含反向传播操作、梯度计算和损失函数。当编译模型以进行部署时,反向图是不必要的。虽然大多数框架会明确分离这些,但复杂的自定义层可能会意外保留对训练专用张量的引用。可达性分析为了识别哪些节点可以安全移除,编译器会进行活跃性或可达性分析。此过程将计算图视为一个有向无环图(DAG),其中边表示数据依赖关系。该算法通常按以下步骤运作:确定根节点: 编译器识别图的“输出”。这些是用户请求的外部张量。反向遍历: 从根节点开始,编译器沿输入边反向遍历图。在此遍历期间访问的每个节点都标记为“活跃”。清理: 图中在此遍历期间未被标记为活跃的任何节点都被视为“死”节点。这些节点不对最终输出做贡献,可以移除。考虑一个简化的计算图,其中一个旁支计算用于记录的标准差,但该值从未被函数返回。digraph G { rankdir=TB; node [style="filled", fontname="Helvetica", shape="box", margin=0.2]; edge [fontname="Helvetica", fontsize=10, color="#495057"]; subgraph cluster_0 { style=invis; node [style=filled]; Input [label="输入张量", fillcolor="#a5d8ff", color="#a5d8ff"]; Weight [label="权重", fillcolor="#a5d8ff", color="#a5d8ff"]; Conv [label="Conv2D", fillcolor="#69db7c", color="#69db7c"]; Relu [label="ReLU", fillcolor="#69db7c", color="#69db7c"]; Output [label="输出", fillcolor="#ffec99", color="#ffec99"]; Input -> Conv; Weight -> Conv; Conv -> Relu; Relu -> Output; } subgraph cluster_1 { style=invis; Stats [label="均值计算\n(未使用)", fillcolor="#e9ecef", color="#adb5bd", fontcolor="#868e96"]; Log [label="日志操作\n(未使用)", fillcolor="#e9ecef", color="#adb5bd", fontcolor="#868e96"]; Relu -> Stats [color="#adb5bd", style="dashed"]; Stats -> Log [color="#adb5bd", style="dashed"]; } }一个数据流图,显示了活跃节点(绿色)与死节点(灰色)。计算均值的分支源自 ReLU 输出,但未连接到最终图输出,使其成为可消除的对象。处理副作用机器学习编译器中的一个情况是对具有副作用的操作的处理。在纯数据流图中,一个操作只有在其数据被使用时才有意义。然而,某些操作,例如 print() 语句或自定义日志操作符,不会产生被数学管道使用的张量输出,但会执行对用户或系统可见的操作。如果编译器激进地移除所有没有数据后继的节点,它可能会移除这些日志操作。为了避免这种情况,IR 通常支持将特定节点标记为“非纯”或具有“副作用”。活跃性分析算法将这些节点视为辅助根。如果一个节点有副作用,它会被标记为活跃,而且重要的是,它的所有输入依赖也都会被标记为活跃。递归优化循环死代码消除很少作为单次遍历运行。它通常是迭代优化循环的一部分。这是因为移除一个死节点可能会使其输入没有剩余消费者,从而导致它们也成为死节点。考虑以下操作序列:$$ \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$ 被消除,编译器就不需要为这些张量分配中间缓冲区。对于大型神经网络,中间激活图可以消耗数千兆字节的显存。通过仔细裁剪图,编译器减少了峰值内存占用,可能允许更大的批次大小或更大的模型适应设备。当为内存限制严格的边缘设备优化时,这尤其有意义。在编译器栈中的实现在诸如 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 是函数返回值。尽管理论上很简单,死代码消除是一个基础性阶段,它确保后续复杂转换的效率。它充当图优化阶段的垃圾回收器,确保硬件只执行为产生预期预测而绝对必需的计算。