深度学习编译器中的优化流程通常表现为扩展与收缩的循环。循环分块或向量化等阶段会增加代码的复杂程度以适应硬件限制,而图层级的清理阶段则负责收缩过程。死代码消除(DCE)和公共子表达式消除(CSE)在编译过程中扮演着“垃圾回收器”的角色。它们清除先前转换或高级框架引入的多余内容,确保硬件只运行生成模型输出所必需的逻辑。公共子表达式消除(CSE)在张量图的背景下,CSE 识别生成相同值的相同操作,并将其替换为单个计算。这种冗余常见于神经网络架构搜索(NAS)生成的复杂模型结构中,或者当高级框架将 Softmax 或 GELU 等复合运算符展开为其基本组成部分时。如果一个图包含两个节点 $N_1$ 和 $N_2$,它们满足以下条件时可被消除:它们具有相同的操作符类型(opcode)。它们的输入相同(结构相等)。它们的属性(步长、填充、扩张)一致。这些操作是纯粹的(无副作用)且确定性的。设想一个场景,模型的一个分支计算一个方差项,该项用于归一化和后续的缩放因子。如果框架将其降级为两个不同的子图,编译器将看到:$$ \begin{aligned} t_1 &= \text{均值}(X) \ t_2 &= (X - t_1)^2 \ v_1 &= \text{均值}(t_2) \ \dots \ v_2 &= \text{均值}((X - \text{均值}(X))^2) \end{aligned} $$没有 CSE,硬件会计算两次方差。编译器的任务是检测生成 $v_2$ 的子图在拓扑上与生成 $v_1$ 的子图相同,重写对 $v_2$ 的依赖使其指向 $v_1$,并移除冗余节点。通过哈希实现高效的 CSE 实现依赖于值编号或结构哈希。编译器按拓扑顺序遍历图,维护一个哈希表,其中键是运算符及其输入操作数的唯一标识符的复合哈希。当访问一个节点时:规范化: 确保属性采用标准格式(例如,重新排列 $A + B$ 与 $B + A$ 等可交换输入)。哈希: 计算 $H = \text{hash}(\text{opcode}, \text{inputs}, \text{attrs})$。查找: 检查哈希表中的已访问节点是否存在 $H$。命中: 将当前节点的所有下游用法替换为表中找到的节点。标记当前节点待移除。未命中: 将当前节点插入表中并继续。ML 编译器中的一个特殊难题是随机操作。像 dropout 或 random_uniform 这样的运算符并非纯粹;执行两次会产生不同的结果。CSE 阶段必须明确地将这些非确定性运算符列入黑名单,以保持模型随机行为的正确性。digraph CSE_Example { rankdir=TB; node [shape=rect, style=filled, fontname="Helvetica", fontsize=10, color="#dee2e6", penwidth=0]; subgraph cluster_before { label="CSE 前"; style=dashed; color="#adb5bd"; in1 [label="输入 X", fillcolor="#e9ecef"]; op1 [label="Conv2D", fillcolor="#a5d8ff"]; op2 [label="ReLU", fillcolor="#ffc9c9"]; op3 [label="Conv2D\n(相同)", fillcolor="#a5d8ff"]; op4 [label="ReLU\n(相同)", fillcolor="#ffc9c9"]; add [label="Add", fillcolor="#b2f2bb"]; in1 -> op1; op1 -> op2; in1 -> op3; op3 -> op4; op2 -> add; op4 -> add; } subgraph cluster_after { label="CSE 后"; style=dashed; color="#adb5bd"; in2 [label="输入 X", fillcolor="#e9ecef"]; opA [label="Conv2D", fillcolor="#a5d8ff"]; opB [label="ReLU", fillcolor="#ffc9c9"]; add2 [label="Add", fillcolor="#b2f2bb"]; in2 -> opA; opA -> opB; opB -> add2 [label="x2"]; } }应用公共子表达式消除前后拓扑结构的比较。冗余的 Conv2D 和 ReLU 分支被合并,有效将该部分的计算需求减半。死代码消除(DCE)死代码消除移除对图的最终输出没有贡献的操作。在深度学习中,死节点通常是其他优化的次要效果。例如,如果代数简化阶段将 $x * 0$ 替换为零张量,则计算 $x$ 的子图就变成了死代码,前提是 $x$ 没有在其他地方使用。当部署带有辅助头(例如 Inception 或 R-CNN 架构)训练的模型时,DCE 也很有用,因为用于训练监督的中间输出在推理期间不相关。标记-清除方法计算图中的标准 DCE 类似于内存管理中的垃圾回收。它通常采用标记-清除算法:识别根节点: 编译器识别图的输出节点(“汇点”)。这些节点本质上是活跃的。标记: 从根节点开始,算法向后(上游)遍历图,将每个输入节点标记为“活跃”。清除: 图中任何未标记为活跃的节点都将从计算中解除连接并移除。尽管简单,但这个过程必须是迭代的。移除一个死节点会减少其输入的引用计数。如果一个输入的引用计数降至零,它也变为死代码。因此,DCE 通常以循环方式运行或使用工作列表算法,直到无法移除更多节点。考虑常量折叠阶段后的以下图状态:$$ \begin{aligned} A &= \text{加载}(\text{"权重"}) \ B &= \text{常量}(0) \ C &= \text{矩阵乘}(X, A) \ D &= \text{乘}(C, B) \quad \rightarrow \quad D = \text{类零}(C) \end{aligned} $$一旦编译器用零张量替换乘法(因为乘以零),对 $C$ 的依赖就被切断了。因此,MatMul 操作成为死代码。移除 MatMul 会使 $A$(权重)成为死代码。最终的图只包含生成零张量的逻辑,通过完全避免权重加载来节省内存带宽。优化不动点DCE 和 CSE 很少单独运行。它们通常在循环中与其他阶段(如常量折叠和规范化)交错执行。这是因为转换通常会暴露新的清理机会。常量折叠 计算编译时已知的值。DCE 移除生成这些常量的原始操作。规范化 重新排列图结构(例如 $(A+B)+C \rightarrow A+(B+C)$)。CSE 找到由规范形式暴露的重复项。DCE 清理由 CSE 断开的冗余分支。这个循环重复进行,直到图达到一个不动点,即应用更多阶段也不会改变图拓扑的状态。{ "data": [ { "x": ["初始", "阶段 1 (常量折叠)", "阶段 2 (DCE)", "阶段 3 (CSE)", "阶段 4 (DCE)"], "y": [100, 100, 85, 70, 65], "type": "scatter", "mode": "lines+markers", "name": "节点数量", "line": {"color": "#4dabf7", "width": 3}, "marker": {"size": 10, "color": "#228be6"} }, { "x": ["初始", "阶段 1 (常量折叠)", "阶段 2 (DCE)", "阶段 3 (CSE)", "阶段 4 (DCE)"], "y": [0, 15, 15, 30, 35], "type": "bar", "name": "移除节点数 (累计)", "marker": {"color": "#ff8787"} } ], "layout": { "title": "优化阶段中的图收缩", "font": {"family": "Helvetica"}, "yaxis": {"title": "节点数量"}, "xaxis": {"title": "优化顺序"}, "showlegend": true, "plot_bgcolor": "#f8f9fa", "paper_bgcolor": "white", "margin": {"l": 50, "r": 20, "t": 50, "b": 50} } }图的复杂程度降低通常分步进行。请注意,在 CSE 阶段之后运行的 DCE 阶段(阶段 4)如何移除通过分支逻辑合并而隔离的剩余节点。处理副作用和控制流MLIR 或 Relay 等高级 IR 在控制流和内存副作用方面带来复杂状况。在这些环境下,简单的图连接不足以判断活跃性。例如,一个操作可能写入某个内存缓冲区,而该缓冲区随后会被读取。即使该操作的直接输出张量未使用,其副作用(内存写入)也保持了它的活跃性。处理有状态图(常见于循环网络或复杂训练图)的编译器必须执行内存 SSA(静态单赋值)分析或别名分析,以确保 DCE 不会移除必要的内存操作。同样,条件分支(if/else)中的节点需要仔细处理。分支内的节点仅当分支本身不可达,或者节点的结果在控制流合并点之后的所有可能执行路径中都未使用时,才被认为是死代码。这需要构建支配者树来证明某个特定计算在所有路径中都是完全多余的。