3. 计算图基础

计算图（Computational Graph）是深度学习编译器一切优化工作的起点与核心载体。无论是静态图（如 TensorFlow 1.x、TorchScript）还是动态图（如 PyTorch Eager 模式），其底层执行语义最终都会被映射为一张由算子节点（Operator Node）和张量边（Tensor Edge）构成的有向无环图（DAG）或带控制依赖的图。编译器对内存规划、算子融合、布局转换、并行调度等优化，本质上都是在计算图上进行模式匹配与重写。因此，理解计算图的表示方式、节点语义、数据依赖关系以及控制流与动态形状的处理策略，是掌握深度学习编译器内存规划的前提。

上图展示了计算图的基本构成要素：算子节点（如 MatMul、BiasAdd、ReLU）通过数据依赖边连接，形成数据流；控制流（if/while/switch）则通过控制边约束执行顺序；动态形状则是在静态图骨架上引入符号化维度，使同一张图能够处理不同 batch size 或序列长度的输入。下面我们将从四个维度深入剖析。

3.1 计算图表示：从 AST 到 Graph 的编译前端

深度学习框架的前端负责将用户编写的 Python 代码转换为计算图。这一过程通常分为两个阶段：AST（抽象语法树）解析与图构建。在 PyTorch 中，TorchScript 通过 torch.jit.trace 或 torch.jit.script 将 Python 函数转换为 Graph 对象；在 TensorFlow 中，tf.function 使用 AutoGraph 将 Python 控制流转换为图内控制流节点。

计算图的存储结构通常采用有向图，每个节点（Node）包含以下关键字段：

• op_type：算子类型（如 "conv2d"、"add"、"reshape"）。
• inputs：输入张量列表，每个输入是一个 Value 对象，指向产生它的节点。
• outputs：输出张量列表，每个输出有唯一的 Value 标识。
• attributes：算子属性（如 stride、padding、axis），这些是编译期常量。
• control_deps：控制依赖列表，用于强制节点执行顺序（如资源初始化、同步操作）。

在内存规划视角下，计算图的拓扑序决定了张量生命周期。编译器通常对图进行拓扑排序，然后依次为每个中间张量分配内存。例如，一个典型的卷积网络，其计算图拓扑序为：Conv → BN → ReLU → Pool → FC。编译器可以据此确定每个张量的活跃区间（从第一次被写入到最后一次被读取），从而复用内存。

3.2 算子节点与数据依赖：内存规划的基石

算子节点是计算图的基本执行单元。每个算子定义了从输入张量到输出张量的映射关系。数据依赖分为两种：读依赖（算子读取输入张量）和写依赖（算子产生输出张量）。编译器必须保证：一个算子的所有输入在其执行前已就绪，且其输出在其执行后才可被后续算子使用。

从内存规划角度，数据依赖直接决定了张量别名分析的可能性。例如，y = x + 1 中，如果 x 在后续不再被使用，编译器可以将 y 直接覆盖到 x 的内存地址上，实现原地更新（in-place）。但若 x 被多个算子共享，则必须分配新内存。下表总结了常见算子的内存行为特征：

算子类型	内存行为	是否可原地	典型内存峰值
Elementwise (Add/ReLU)	输入输出形状相同	是（若输入无其他引用）	1×输出大小
MatMul / Conv	输出形状由输入形状与属性决定	否（输入输出形状不同）	输出 + 中间累加缓冲
Reshape / Transpose	仅改变元数据，不复制数据	是（视图操作）	0（共享底层存储）
Concat / Split	沿指定维度拼接或切分	否（通常需要新内存）	输出大小 + 输入引用
Gather / Scatter	按索引访问内存	Scatter 可部分原地	输出大小 + 索引缓冲

编译器在内存规划时，会遍历计算图，构建张量活跃性分析（Liveness Analysis）。对于每个张量，记录其定义点（Def）和最后使用点（Last Use）。活跃区间为 [Def, LastUse]。内存分配器（如 Arena Allocator）利用这些区间进行区间着色，将不重叠的活跃区间映射到同一物理内存块。

3.3 控制流与动态形状：编译器的挑战与策略

真实模型（如 NLP 中的 BERT、推荐系统中的 Embedding 层）往往包含控制流（if/else、while 循环）和动态形状（如变长序列、动态 batch）。这给计算图的内存规划带来了巨大挑战。

控制流处理策略：

• 图内控制流：将 if/while 转换为 Switch、Merge、Enter、Exit 等控制流节点。编译器需要为每个分支单独进行内存规划，并在分支汇合点处理张量版本冲突。
• 子图划分：将控制流边界作为图切分点，每个子图独立编译和内存分配。子图间通过宿主张量（Host Tensor）传递数据。
• 循环优化：对于 while 循环，编译器可进行循环不变量代码提升（将循环内不变的计算移到循环外），减少重复内存分配。

动态形状处理策略：

动态形状意味着张量的维度在编译期未知，通常用符号变量（如 batch_size）表示。编译器采用符号形状推断（Symbolic Shape Inference）传播维度约束。内存规划时，有两种主流方案：

1. 保守分配：为动态维度分配最大可能大小（如 batch=64），浪费内存但实现简单。
2. 动态内存池：运行时根据实际形状动态扩展或收缩内存池。TVM 的 VM 和 XLA 的 DynamicDim 机制即采用此方案。

3.4 案例拆解：BERT 模型中的控制流与动态形状内存规划

我们以一个简化的 BERT 推理场景为例，说明编译器如何处理控制流与动态形状。假设输入为变长序列，batch size 固定为 8，序列长度在 32~128 之间变化。

初始现象与目标：模型包含 while 循环（用于多层 Transformer 堆叠）和 if 分支（用于 attention mask 处理）。目标是使内存峰值不超过 4GB，并保证不同序列长度下均能正确执行。

关键输入参数：

• 隐藏层维度：768
• 层数：12
• 序列长度：动态（32~128）
• batch size：8

分析判断过程：

编译器首先将 Python 代码转换为计算图，识别出 while 循环体（包含 12 个 Transformer 层）和 if 分支（mask 处理）。对于循环体，编译器采用循环展开（Unrolling）策略，将 12 层展开为 12 个独立的子图，每层的内存规划独立进行。对于动态序列长度，编译器引入符号变量 S，并生成形状多态内核。

实际操作步骤：

# 伪代码：编译器生成的形状多态内核调度
def attention_kernel(q, k, v, mask, seq_len):
    # seq_len 为运行时参数
    if seq_len <= 64:
        return attention_impl_small(q, k, v, mask)  # 专用小形状内核
    else:
        return attention_impl_large(q, k, v, mask)  # 通用内核

内存规划时，编译器为每个 Transformer 层分配一个内存池。由于循环展开后各层独立，层间张量（如 hidden_states）可以复用同一块内存。对于动态形状，编译器在运行时根据实际 seq_len 计算所需内存大小，并从全局内存池中动态申请。

结果对比与风险：

策略	内存峰值	执行时间	风险
静态形状（固定 seq=128）	4.8 GB	12 ms	短序列时浪费内存
动态形状 + 内存池	3.2 GB	10 ms	运行时分配开销
动态形状 + 内核特化	3.2 GB	8 ms	内核编译时间增加

复盘经验：动态形状下，内存规划必须与内核调度协同设计。单纯的内存池无法解决内核执行效率问题，而内核特化虽然提升了速度，但增加了编译开销。实际工程中，通常采用分层策略：对常见形状（如 seq=64、128）预编译特化内核，对罕见形状回退到通用内核。

3.5 常见误区与总结复盘

在计算图基础的学习与实践中，以下误区值得警惕：

总结：计算图是深度学习编译器的中间表示（IR），其质量直接决定了后续内存规划与算子优化的上限。理解数据依赖是进行张量生命周期分析的基础；控制流要求编译器具备跨分支的内存隔离能力；动态形状则要求编译器在运行时具备形状自适应与内存动态管理能力。掌握这些基础概念，是深入理解内存池设计、算子融合、图重写等高级主题的前提。在实际工程中，建议使用 torch.fx 或 MLIR 等工具对计算图进行可视化与静态分析，以验证图结构的正确性。