计算图和有向无环图

深度学习 (deep learning)编译器中的基本分析单元不是单个指令或内存地址，而是计算图。虽然高级框架通常向开发者提供命令式接口，允许类似标准Python的逐行执行，但编译器将模型视为静态数据流网络。这种观念的转变使得编译器能够在硬件上启动任何内核之前，对全局数据依赖性、内存生命周期和操作粒度进行考量。

计算图正式表示为有向无环图（DAG）。在这种结构中，节点代表计算操作（例如矩阵乘法、卷积或逐元素加法），有向边代表这些操作之间的数据（张量）流动。这种表示方法将模型的数学定义与其执行策略分离。

考量一个由以下方程表示的标准全连接层：

$$Y = \text{ReLU}(X W^T + b)$$

在通用编程语言中，这可能作为三个不同函数调用的序列执行。在计算图中，它是一个连接结构，其中矩阵乘法节点的输出直接流入加法节点，随后流入激活节点。这种结构明确编码了依赖链：加法不能在乘法完成之前进行，但它也表明$X$和$W$是独立的，理论上可以同时载入。

图的构成

图由编译器必须管理的三种主要组件组成：

1. 节点（操作符）： 它们是功能单元。在TVM Relay或XLA HLO等高级中间表示（IR）中，节点带有语义属性。例如，一个Conv2D节点将内核大小、步长、填充和膨胀率等静态信息作为属性存储，而不是作为输入张量。
2. 边（数据依赖）： 边代表流经系统中的张量。重要的是，编译器图中的边带有类型信息，即形状和数据类型（dtype）。在编译过程中，“形状推断”阶段通过图传递此信息，以确定内存分配需求，而无需执行代码。
3. 源节点和汇点节点： 图具有明确的入口点（参数 (parameter)或输入数据）和出口点（模型预测）。

DAG的无环特性在大多数张量编译器中都严格遵守。数据流中的循环将意味着一个时间悖论，即一个操作依赖于其自身的未来输出。循环神经网络 (neural network)（RNN）本身包含循环，通常通过将循环展开成固定时间步长的长而无环的操作序列来处理，或者使用专门的控制流操作符将循环体封装为子图来处理。

数据流的可视化

为了了解编译器如何看待模型，我们可以可视化一个简化后的残差块子图。这展现了数据如何拆分和合并，创建出调度器必须解决的复杂依赖模式。

一个残差块的表示。虚线表示跳跃连接，它在数据路径中创建了一个分支，并在“加法”节点处合并。编译器使用这种拓扑结构来确定输入张量必须在内存中保留，直到加法操作完成。

数据流与控制流

标准编译器（如GCC或LLVM）擅长处理控制流图（CFG），其中节点代表指令的基本块，边代表跳转或分支。然而，AI编译器侧重于数据流图（DFG）。在DFG中，数据的可用性触发执行，而不是程序计数器。

当模型需要条件逻辑时，例如递归网络中的动态执行或Transformer中的条件生成，这种区别会使问题变得复杂。深度学习 (deep learning)编译器通过两种主要方法来处理这种情况：

1. 追踪/展开： 如果控制流依赖于编译时常量（例如，一个循环精确运行4次），编译器会将图展开成一个扁平的序列。if/else逻辑会消失，取而代之的是追踪期间所采取的特定分支。
2. 函数式控制流： 对于依赖于数据的控制流（例如，依赖于运行时计算的张量值的逻辑），编译器引入了像If或While这样的专用节点。与标准CPU分支不同，这些节点通常将整个子图作为参数 (parameter)。一个If节点可能接受一个条件张量，一个“then”子图和一个“else”子图。这在支持动态行为的同时，保持了图的属性。

拓扑排序和调度

DAG定义了依赖关系，但没有固定的执行顺序。对于上面显示的残差块，编译器知道conv1必须在bn1之前发生，但它需要将此图线性化为硬件指令序列。这通过拓扑排序实现。

拓扑排序算法生成顶点的线性排序，使得对于从$u$到$v$的每条有向边，顶点$u$在排序中位于$v$之前。在复杂的图中，可能存在多个有效的拓扑排序。编译器的调度器会选择一个针对特定指标进行优化的排序，例如：

• 峰值内存使用： 以允许中间张量尽早释放的顺序执行分支。
• 缓存局部性： 在生产者之后立即调度消费者，以将数据保留在L1或L2缓存中。

以下可视化显示了DAG中执行深度的观念。相同深度的节点仅依赖于先前深度的节点，表明在有足够硬件资源的情况下，这些操作理论上可以并行执行。

操作按其依赖深度分组。具有多个块的步骤表示并行执行或流并发的可能性。任何步骤中条形的宽度代表该点图的“宽度”。

图的粒度和降低

DAG的复杂程度很大程度上取决于抽象级别。在最高层（前端IR），一个节点可能代表一个“多头注意力 (multi-head attention)”块。这是一种宏观视角，有助于图级优化，例如检查硬件是否具有特定的注意力机制 (attention mechanism)加速器。

随着编译的进行，此图被“降低”。单个注意力节点会分解成几十个较小的节点：矩阵乘法、转置、softmax函数和逐元素除法。编译器继续将其作为DAG进行管理，但图的大小大幅增加。最终，这些节点被转换为循环嵌套，最后转换为硬件指令。

理解DAG是第二章讨论的优化技术的先决条件。例如，操作符融合就是识别DAG中特定子模式（如Conv2D后跟ReLU），并用一个在一次内核启动中执行两个操作的“融合”节点替换它们。如果没有图结构提供的全局可见性，此类优化将无法安全应用。