数据流图与依赖关系

机器学习模型本质上是一系列应用于数据的数学变换。尽管Python代码可能逐行顺序执行这些变换，但其底层逻辑通常允许更灵活的执行策略。为了善用这种弹性，机器学习编译器将线性源代码转换为数据流图。此图明确描绘了运算及其间的依赖关系，为编译器提供了数据在系统中流动的完整视图。

数据流图的构成

在编译器中间表示（IR）的背景下，数据流图通常构造为有向无环图（DAG）。这种结构包含两个主要元素：

1. 节点（顶点）： 它们代表计算运算或“操作”。在深度学习场景中，一个节点可以是一个低级原语，如加法或乘法，也可以是一个高级运算符，如二维卷积或Softmax。
2. 边（链接）： 它们代表数据的流动。连接节点 A 到节点 B 的边表示节点 A 的输出是节点 B 所需的输入。这些边承载张量数据，包括形状和类型信息。

与通用编译器中使用的标准控制流图（侧重于指令和跳转的顺序）不同，数据流图侧重于数据的可用性。图中的一个运算在其所有输入边都拥有有效数据的那一刻即可执行。

考虑一个简单的线性层计算：$Y = (X \times W) + b$。

在Python中，这看起来像一系列变量赋值。在IR中，它是一个图，其中 $X$ 和 $W$ 流入一个矩阵乘法节点，结果流入一个加法节点，而 $b$ 作为第二个参数流入同一个加法节点。

一个表示线性层运算的数据流图。该结构明确显示偏置加法必须在矩阵乘法完成后才能进行。

依赖关系与执行顺序

数据流图的主要目的是通过依赖关系确保正确性。当一个运算依赖于另一个运算的结果时，就存在依赖关系。这些是严格的限制：编译器不能以违反这些边的方式重新排序运算。

然而，该图也显示了哪些地方不存在依赖关系。这是优化最宝贵的发现。如果两个节点之间没有有向路径连接，它们就是独立的。

串行与并行执行

在命令式脚本中，独立的运算通常顺序编写，仅仅因为语法要求如此。

# Python执行是串行的
x = load_image()
feat_a = extract_features_method_A(x)  # 步骤 1
feat_b = extract_features_method_B(x)  # 步骤 2
result = concatenate(feat_a, feat_b)   # 步骤 3

对于解释器而言，步骤1必须在步骤2开始前完成。然而，数据流图显示 feat_a 和 feat_b 都只依赖于 x。extract_features_method_A 和 extract_features_method_B 之间没有依赖关系。

这种缺少依赖的情况使得编译器能够将这些运算安排在支持并行的硬件（如具有多流的GPU或多核CPU）上并行运行。该图将程序从严格的序列转变为一系列可利用的契机。

这种分支结构展示了独立的路径。这两个卷积运算共享一个输入，但独立运行，使得编译器可以同时调度它们。

单一静态赋值（SSA）

大多数现代机器学习编译器（包括TVM、XLA和MLIR）都使用一种称为单一静态赋值（SSA）的中间表示形式。在SSA中，每个张量只被赋值一次。在图定义中，你不能就地修改张量；相反，任何看起来修改数据的运算实际上都会生成一个新的张量版本。

这种不变性简化了依赖关系的分析。编译器在图分析阶段无需追踪复杂的状态变化或担心内存别名（即两个变量指向同一内存位置）。如果节点 B 使用张量 $T_1$，并且我们想通过改变 $T_1$ 的生成方式来优化图，我们可以在确切了解哪些下游节点会受影响的情况下进行。

从图到调度

依赖关系映射完成后，编译器执行拓扑排序以确定一个有效的线性执行顺序。拓扑排序是一种节点排序，使得对于从节点 A 到节点 B 的每条有向边，节点 A 在排序中都出现在节点 B 之前。

对于复杂的图，通常存在多个有效的拓扑排序。

1. 深度优先： 优先完成一个计算分支，然后再开始另一个。这可以通过迅速释放中间张量来减少峰值内存使用。
2. 广度优先： 优先执行当前深度所有可用的节点。这可以最大化并行度，但可能增加内存占用，因为更多的中间结果同时保持活动状态。

编译器优化通常涉及修改图结构、融合节点或重写子图，同时保持这些数据依赖关系的有效性。例如，如果一个节点被移除（死代码消除），编译器必须确保没有其他节点依赖其输出。如果节点被融合，新融合节点的依赖关系是原始节点依赖关系的并集。

了解数据流图是看待代码的第一步，它将代码视为一个结构化的运算网络，而非纯文本。这种视角促成了我们将在后续关于图级优化的章节中讨论的强大变换。