深度学习编译器基础：从计算图到中间表示（IR）

深度学习编译器是连接高层神经网络模型与底层硬件执行效率的关键桥梁。在AI芯片（如NPU、GPU、TPU）的软件栈中，编译器承担着将用户用PyTorch、TensorFlow等框架书写的模型，翻译成可在特定硬件上高效运行的机器指令的重任。这一过程并非一步到位，而是遵循经典编译器理论中的“三段式”结构：前端（Frontend）、中端（Middle-end）与后端（Backend）。理解这一分层架构，是掌握NPU指令调度优化的前提。本章将深入剖析从计算图（Graph）到中间表示（IR）的转换机制，并详细拆解编译器三个阶段的职责边界与核心算法。

上图展示了AI编译器典型的“三段式”流水线。前端负责将不同深度学习框架（PyTorch、TensorFlow、PaddlePaddle等）产出的计算图统一转换为编译器内部的图表示（Graph IR）。中端在此图表示上执行与硬件无关的优化，如算子融合、常量折叠、死代码消除等。后端则将优化后的IR逐步lowering为硬件相关的指令，并完成寄存器分配、指令调度、内存分配等关键步骤，最终生成可执行的二进制码或NPU微码。下面我们逐一深入各阶段的核心原理与实战要点。

2.1 前端：从框架计算图到统一图IR

核心原理：计算图的本质与转换挑战

深度学习模型本质上是一个有向无环图（DAG），节点代表算子（如Conv2D、MatMul、ReLU），边代表张量（Tensor）的数据流向。前端的主要任务是将框架的动态图（PyTorch的Eager模式）或静态图（TensorFlow的Graph模式）捕获为可静态分析的图结构。对于动态图，编译器通常采用追踪（Tracing）或源码转换（Scripting）两种机制。追踪执行一次模型并记录算子调用序列，但无法处理依赖运行时数据的控制流（如if tensor.sum() > 0）；源码转换则通过解析Python AST将函数体编译为静态图，能保留控制流但实现复杂度高。

转换过程中的核心挑战是算子语义对齐。不同框架对同一数学运算的命名、属性定义、数据格式（NHWC vs NCHW）存在差异。例如，PyTorch的nn.Conv2d默认权重格式为[out_channels, in_channels/groups, kH, kW]，而TensorFlow的tf.nn.conv2d的filter格式为[kH, kW, in_channels, out_channels]。前端必须将这些差异统一为编译器内部的规范表示。

操作流程：前端转换的标准步骤

1. 图捕获：通过框架提供的API（如torch.jit.trace）获取原始计算图。
2. 规范化：将算子属性统一（如将padding的字符串形式"SAME"转为具体数值）。
3. 类型与形状推断：沿图传播张量的数据类型（dtype）和形状（shape），为后续优化提供信息。
4. 转换为Graph IR：将框架算子映射到编译器内置的算子集（OpSet），生成框架无关的中间表示。

案例拆解：PyTorch模型转Relay IR

假设我们有一个简单的两层卷积网络。使用TVM的from_pytorch转换时，前端会执行以下操作：

import torch
import tvm
from tvm import relay

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, 3, padding=1)
        self.relu = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(16, 32, 3, padding=1)
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        return x

model = Model()
input_shape = (1, 3, 224, 224)
trace = torch.jit.trace(model, torch.randn(input_shape))
mod, params = relay.frontend.from_pytorch(trace, {"input": input_shape})
print(mod)

转换后，mod中的IR会显示为nn.conv2d、nn.relu等标准算子，且权重参数被提取到params字典中。此时IR仍是高层图表示，尚未涉及任何NPU相关的细节。

2.2 中端：硬件无关的IR优化

核心原理：优化Pass的组织与执行

中端是编译器体现“智能”的核心区域。它接收前端产出的Graph IR，执行一系列Pass（优化遍）来降低计算量、减少内存访问、改善数据局部性。这些Pass与具体硬件无关，因此可以被所有后端复用。常见的Pass包括：

• 算子融合（Operator Fusion）：将Conv+BN+ReLU等连续算子融合为单个算子，减少中间张量的内存读写。
• 常量折叠（Constant Folding）：在编译期计算所有输入均为常量的子图。
• 代数化简（Algebraic Simplification）：如x*1替换为x，x+0替换为x。
• 内存布局转换（Layout Transformation）：将NCHW转为NHWC或NC1HW0等硬件友好的格式。

判断方法：如何选择优化Pass顺序

Pass的执行顺序至关重要。例如，先做算子融合再做布局转换，可以避免融合后的算子因布局不匹配而无法匹配硬件指令。通常遵循以下原则：

1. 先做“简化类”Pass（常量折叠、死代码消除），缩小图规模。
2. 再做“融合类”Pass，减少算子数量。
3. 最后做“数据布局”Pass，为后端做准备。

实战场景：NPU上的Conv+BN融合

在NPU上，卷积运算通常由专用硬件单元完成，而BN（Batch Normalization）在推理阶段可以简化为逐元素的缩放和偏移。将BN融合进卷积的权重和偏置中，可以避免一次额外的全张量遍历。具体操作如下：

# 假设conv权重为W，偏置为b，BN参数为gamma, beta, mean, var, eps
# 融合后的权重 W' = W * gamma / sqrt(var + eps)
# 融合后的偏置 b' = (b - mean) * gamma / sqrt(var + eps) + beta
# 在Relay中，可以通过FoldScaleAxis或CombineParallelConv2D等Pass自动完成。

在TVM中，relay.transform.FoldScaleAxis会识别Conv2D后接Mul/Add的模式，并将其吸收进卷积权重。这一优化在NPU上尤其重要，因为NPU的片上存储有限，减少中间张量的写回DDR能显著降低功耗和延迟。

2.3 后端：面向NPU的指令生成与调度

核心原理：从优化IR到硬件指令的Lowering

后端是编译器中最贴近硬件的部分。它接收中端优化后的IR，将其逐步lowering为硬件抽象层（HAL）或直接生成NPU指令。这一过程涉及多个关键步骤：

• 算子选择（Op Selection）：将IR中的算子映射到NPU支持的硬件指令（如矩阵乘指令、向量指令、标量指令）。
• 循环优化（Loop Optimization）：包括循环分块（Tiling）、循环展开（Unrolling）、循环重排序（Reordering），以匹配NPU的脉动阵列或SIMD单元。
• 内存分配（Memory Allocation）：将张量分配到片上SRAM、全局内存或寄存器文件中，并处理DMA传输。
• 指令调度（Instruction Scheduling）：对生成的指令序列进行排序，以最大化流水线并行度，隐藏内存访问延迟。

操作流程：后端编译的典型流水线

1. IR Lowering：将Relay IR转换为Tensor Expression（TE）或直接转换为MLIR的Linalg方言。
2. Tiling与Binding：根据NPU的硬件参数（如SRAM大小、计算单元数量）确定循环分块大小，并将循环轴绑定到硬件线程或流水线阶段。
3. 指令发射：将优化后的循环结构转换为NPU汇编指令或LLVM IR。
4. 寄存器分配与指令调度：使用图着色算法分配寄存器，并采用列表调度（List Scheduling）算法对指令排序。
风险说明：NPU的指令集通常具有VLIW（超长指令字）特性，一条指令可以同时控制多个计算单元。如果调度不当，会导致大量NOP（空操作）填充，浪费计算资源。因此，后端调度器必须精确建模NPU的流水线深度和资源冲突。

案例拆解：NPU上的矩阵乘指令调度

假设NPU有一个16x16的脉动阵列，SRAM大小为256KB。我们要计算一个[1024, 1024] x [1024, 1024]的矩阵乘。后端会执行以下操作：

# 1. Tiling: 将输出矩阵分为 (64, 64) 个 tile，每个 tile 大小为 16x16
# 2. 数据搬运: 将 A 的 tile 和 B 的 tile 从 DDR 加载到 SRAM
# 3. 计算: 在脉动阵列上执行 16x16x16 的矩阵乘
# 4. 累加: 将部分和累加到输出 tile
# 5. 写回: 将输出 tile 从 SRAM 写回 DDR
# 调度器需要交错执行上述步骤，使得 DMA 搬运与计算重叠。

初始调度可能如下：加载A_tile(0,0) → 加载B_tile(0,0) → 计算 → 写回。但这样DMA和计算是串行的。优化后的调度为：预加载A_tile(0,1)和B_tile(1,0)，在计算当前tile的同时进行下一次DMA传输，从而隐藏内存延迟。通过软件流水线（Software Pipelining）技术，可以将端到端延迟降低约30%。

调度策略	计算时间 (cycles)	DMA时间 (cycles)	总时间 (cycles)
串行调度	4096	2048	6144
软件流水线	4096	2048 (隐藏)	4300

总结复盘：后端优化的核心是“算”与“搬”的重叠。NPU的算力通常远大于DDR带宽，因此瓶颈往往在数据搬运。通过增大Tiling粒度、使用双缓冲（Double Buffering）机制，可以有效提升计算访存比（Arithmetic Intensity）。

2.4 常见误区与工程实践要点

误区一：忽视IR的语义保持

在Pass编写过程中，如果对IR的改写破坏了原有语义（如改变了浮点运算顺序），会导致数值精度下降。编译器开发者应使用等价变换，并在关键Pass后添加验证器（Validator）检查IR的合法性。

误区二：过早进行硬件相关优化

在中端就引入NPU特定的Tiling策略，会导致IR无法复用于其他硬件（如GPU）。正确的做法是保持中端IR的硬件无关性，将硬件特定优化推迟到后端。

误区三：忽略编译时间

深度学习模型动辄上百个算子，如果每个Pass都遍历全图，编译时间可能长达数分钟。工程上应使用Pass管理器（Pass Manager）的缓存机制，只对变更的子图重新优化。

实践建议：在开发NPU编译器时，建议先使用MLIR框架搭建中后端，利用其方言（Dialect）机制逐步lowering。MLIR提供了丰富的Pass基础设施和验证工具，能显著降低开发成本。

本章从计算图出发，详细阐述了AI编译器前端、中端、后端的基本流程。前端解决“如何看懂模型”，中端解决“如何优化模型”，后端解决“如何高效执行”。三者的协作关系如同翻译、编辑和排版：翻译保证语义准确，编辑提升文采，排版适配纸张。掌握这一分层思想，是后续深入NPU指令调度优化的基石。

← 上一章 📖 返回目录 下一章 →