1. 课程导入
深度学习编译器是连接高层神经网络模型描述与底层硬件执行之间的关键桥梁。随着模型规模从百万参数膨胀至千亿甚至万亿级别，单纯依赖框架运行时（如PyTorch eager模式）已无法满足训练和推理场景对性能、显存效率和延迟的极致要求。编译器通过静态分析与动态优化，将模型计算图转换为高度优化的机器码，并在这一过程中完成两项核心任务：内存规划（Memory Planning）与图重写（Graph Rewriting）。本课程将系统拆解这两大技术支柱，帮助学员建立从计算图到物理内存的完整认知链路。
上图展示了深度学习编译器从模型输入到硬件执行的核心链路。编译器接收高层模型描述后，首先构建计算图，随后进入核心优化阶段——图重写负责改变计算图的结构与算子组合方式，内存规划则负责为图中每个张量分配物理内存并最大化复用效率。两者相互依赖：图重写改变了张量的生命周期，内存规划则需根据重写后的图重新计算分配策略。最终生成的可执行代码在目标硬件上运行，其性能直接取决于这两项优化的质量。
1.1 深度学习编译器概述
深度学习编译器并非传统意义上的"编译器"（如GCC、LLVM），而是一个面向深度学习领域的专用编译栈。它接收的输入是高层神经网络模型（如PyTorch的nn.Module或ONNX图），输出是可在特定硬件上高效执行的机器码或底层指令序列。其核心价值在于自动化硬件适配与跨平台性能迁移——同一份模型代码，通过不同编译器后端，可以运行在NVIDIA GPU、AMD GPU、华为昇腾NPU、寒武纪MLU甚至手机端NPU上，且性能接近手写优化内核。
从架构角度看，深度学习编译器通常分为三层：前端负责将框架模型转换为统一的中间表示（IR）；中端执行与硬件无关的优化（图重写、内存规划、算子规范化）；后端将优化后的IR映射到目标硬件的指令集或库调用。本课程聚焦中端优化中最关键的两个环节——图重写与内存规划，因为它们是决定最终性能上限的"幕后推手"。
核心认知：深度学习编译器的性能天花板由中端优化决定，而非后端代码生成。即使后端拥有完美的指令调度，如果图重写没有消除冗余计算、内存规划没有实现高效复用，最终性能仍会大打折扣。以BERT-base模型为例，未经优化的PyTorch eager模式推理延迟约8ms，经过TVM编译优化后可降至2ms以下，其中约60%的性能提升来自图重写（算子融合、常量折叠），约25%来自内存规划（减少分配开销、提升缓存命中率），其余来自后端代码生成优化。
1.2 内存规划：显存效率的基石
内存规划解决的核心问题是：计算图中每个张量应该存放在哪里？何时分配？何时释放？如何复用？在深度学习训练中，显存往往比算力更稀缺——一块A100 80GB GPU在训练 175B模型时，仅参数和优化器状态就需要约1.2TB显存，必须通过张量并行、流水线并行和内存优化技术才能勉强容纳。内存规划的目标就是在满足计算依赖的前提下，最小化峰值内存占用，同时避免碎片化和不必要的拷贝。
内存规划的核心原理包括三个层次：
- 生命周期分析（Liveness Analysis）：对计算图中的每个张量，确定其从定义到最后一次被使用的区间。生命周期不重叠的张量可以共享同一块内存。例如，两个连续的卷积层，第一层的输出在第二层计算完成后即不再使用，因此第二层的输出可以复用第一层输出的内存空间。
- 内存池化（Memory Pooling）：将多个生命周期不重叠的张量分配到同一块物理内存区域。池化策略包括最佳适配（Best-fit）、首次适配（First-fit）和伙伴系统（Buddy System）。TVM中默认使用最佳适配策略，通过维护空闲块列表，为每个张量寻找最接近其大小的空闲块，以减少内部碎片。
- In-place操作优化：某些算子（如ReLU、Dropout）允许输入输出共享同一内存地址。通过识别这些算子，编译器可以在图重写阶段标记in-place属性，从而减少额外内存分配。例如，PyTorch中的
torch.relu_()就是in-place版本，但编译器需要静态分析确认安全性。
| 内存规划策略 | 适用场景 | 优势 | 风险/代价 |
|---|---|---|---|
| 静态规划（编译期确定） | 推理场景、固定batch size | 零运行时开销、可预测性强 | 无法应对动态shape |
| 动态规划（运行时决策） | 训练场景、可变batch size | 灵活适应输入变化 | 运行时分配开销、碎片化 |
| 混合策略 | 训练+推理混合部署 | 兼顾灵活性与效率 | 实现复杂度高 |
以TVM的内存规划器为例，其工作流程为：首先遍历计算图，为每个张量节点分配唯一的ID；然后执行逆拓扑排序，计算每个张量的生命周期区间；接着将所有张量按生命周期排序，使用贪心算法尝试将新张量放入已释放的空闲块中；最后输出内存分配表，指导代码生成阶段插入分配/释放指令。整个过程在编译期完成，运行时仅需执行预定的分配方案。
1.3 图重写：计算图的结构优化
图重写是对计算图进行等价变换，以降低计算量、减少内存访问或改善数据局部性。它不同于传统编译器中的"优化Pass"，而是针对深度学习算子的特性设计了一系列模式匹配与替换规则。图重写的核心价值在于：将高层语义的算子组合转换为硬件友好的底层实现。
常见的图重写模式包括：
- 算子融合（Operator Fusion）：将多个连续算子合并为一个复合算子，减少内核启动开销和中间张量的内存读写。最经典的案例是Conv-BN-ReLU融合——将卷积、批归一化和ReLU激活合并为单个内核，在GPU上可减少约30%的延迟。融合的数学基础是BN在推理阶段可以折叠为卷积的缩放因子和偏置项。
- 代数化简（Algebraic Simplification）：包括常量折叠（Constant Folding）、冗余消除（Dead Code Elimination）、零乘消除（Zero Multiplication Elimination）等。例如，当某个张量乘以0时，可以直接将结果替换为全零张量，跳过计算。
- 数据布局转换（Layout Transformation）：将张量的内存布局从NCHW转换为NHWC或NC1HW0等硬件友好的格式。GPU上的cuDNN通常对NHWC布局有更好的向量化支持，而NPU可能偏好NC1HW0的分块布局。布局转换本身有开销，因此编译器需要权衡转换成本与后续计算的收益。
- 算子替换（Operator Substitution）：将计算复杂度高的算子替换为等价的低复杂度实现。例如，将大卷积核分解为多个小卷积核（如5×5分解为两个3×3），或将全连接层转换为1×1卷积。
经验提醒：图重写并非总是"越多越好"。过度融合可能导致寄存器溢出或缓存抖动，反而降低性能。例如，将整个ResNet-50的卷积层全部融合为一个巨型内核，虽然消除了所有中间张量的内存读写，但内核的寄存器压力过大，导致GPU占用率下降。实践中，融合的粒度需要根据目标硬件的寄存器数量和缓存大小进行调优。
1.4 案例拆解：BERT推理的内存规划与图重写实战
为了直观理解两项技术的协同作用，我们以BERT-base推理场景为例进行拆解。模型包含12层Transformer编码器，每层包含多头注意力（MHA）和前馈网络（FFN）。输入序列长度为128，batch size为32，FP16精度。
初始状态与约束：模型参数约110M（FP16下220MB），中间激活张量峰值约1.2GB（未优化时）。目标是将峰值显存控制在800MB以内，同时将推理延迟从8ms降至3ms以下。
图重写阶段的操作：
- 将LayerNorm的均值/方差计算与后续缩放合并为单内核，减少两次全局归约。
- 将MHA中的QKV线性变换合并为单个GEMM（3×768→768），减少内核启动次数。
- 将FFN中的两个线性层与GELU激活融合，消除中间张量的内存写入。
- 将残差连接与Add操作合并到后续算子中，避免额外的内存读写。
内存规划阶段的操作：
- 分析重写后图中每个张量的生命周期，发现MHA的中间注意力矩阵（batch×head×seq×seq）生命周期极短，仅存在于softmax计算前后。
- 将注意力矩阵的内存与FFN第一层的输出复用——因为FFN的输入在注意力计算完成后才产生，两者生命周期不重叠。
- 对LayerNorm的中间统计量（均值、方差）使用in-place优化，直接覆盖输入张量的部分区域。
结果对比：
| 指标 | 未优化 | 图重写后 | 图重写+内存规划 |
|---|---|---|---|
| 峰值显存 | 1.2GB | 950MB | 720MB |
| 推理延迟 | 8.2ms | 4.5ms | 3.1ms |
| 内核启动次数 | 1,842 | 1,156 | 1,156 |
| 内存分配次数 | 2,300 | 1,800 | 42 |
风险与代价：内存复用策略使得张量之间的依赖关系更加复杂，调试时难以直接观察中间张量的值。为此，编译器提供了"内存规划验证模式"，在调试阶段关闭复用，确保计算正确后再启用优化。此外，图重写后的融合内核在CPU上可能无法获得同等收益，因为CPU的缓存层次与GPU不同，需要针对不同后端分别调优。
风险提示：内存规划中的in-place优化存在安全隐患。如果两个算子之间存在隐藏的数据依赖（例如通过别名或视图操作），in-place覆盖可能导致计算错误。编译器必须进行严格的依赖分析，包括检查张量的存储偏移、步长和视图关系。TVM中通过StorageRewritePass的保守策略处理——只有确认无别名时才启用in-place。
1.5 常见误区与总结复盘
在学习内存规划与图重写时，初学者常陷入以下误区：
- 误区一：认为内存规划只是"减少显存占用"。实际上，内存规划更重要的目标是提升访存局部性——将生命周期相近的张量放在相邻内存区域，可以提高缓存命中率，减少DRAM带宽压力。在内存带宽受限的算子（如Elementwise）中，这一优化对性能的影响远大于显存容量的节省。
- 误区二：认为图重写是"万能加速器"。图重写只能消除结构性冗余，无法优化单个算子的计算效率。如果模型本身包含大量低效算子（如逐元素操作），图重写能做的只是将它们融合，但无法改变算法复杂度。此时需要结合算子级优化（如tiling、vectorization）才能获得显著加速。
- 误区三：忽略动态shape场景。许多编译器优化（尤其是内存规划）假设张量形状在编译期已知。但在自然语言处理中，输入序列长度可能动态变化。处理动态shape需要引入运行时分支或使用"桶化"（bucketing）策略——将相似长度分组，每组使用静态规划。这增加了编译器的复杂度，但也是工业级部署的必备能力。
总结复盘：本章建立了深度学习编译器的整体认知框架。核心要点包括：编译器通过图重写改变计算图结构，通过内存规划优化张量存储，两者协同决定了最终性能。学习路径上，建议先掌握IR设计（第2-4章），再深入图重写Pass框架（第5-8章），随后学习内存生命周期分析与池化算法（第9-12章），最后通过端到端实战（第13-15章）巩固知识。每章将结合TVM、XLA等开源编译器的真实实现，帮助学员将理论转化为可落地的工程能力。