3. AI编译器概览

在深入MLIR中间表示之前，必须理解AI编译器所处的生态位。AI编译器并非凭空产生，它上承深度学习框架（TensorFlow、PyTorch）的图表示，下接各类硬件后端（GPU、NPU、CPU）。本章将剖析主流框架的图表示差异，并系统对比TVM、XLA、Glow三大代表性编译器的架构设计，为后续理解MLIR如何统一这些分散的中间层奠定基础。

上图展示了AI编译器生态的整体架构。左侧是框架层与编译器层，右侧是MLIR作为统一中间表示层的位置。理解这一生态，需要先厘清两个关键问题：框架的图表示如何影响编译器的输入形态？以及现有编译器在IR设计上有何根本差异？下面逐一展开。

3.1 主流AI框架的图表示：静态图与动态图的博弈

3.1.1 核心原理：图表示的本质

深度学习框架的核心抽象是计算图（Computation Graph）——将神经网络表达为有向无环图（DAG），节点表示算子（Op），边表示张量（Tensor）流动。但不同框架对图的构建时机、可变性、优化空间的处理截然不同，这直接决定了编译器需要面对什么样的输入。

• TensorFlow（静态图）：采用"定义-执行"分离模式。用户先构建GraphDef（协议缓冲区格式），再通过Session执行。图在运行前已完全固定，便于全局优化（如内存复用、算子融合），但调试困难、控制流表达笨拙。
• PyTorch（动态图）：采用"命令式"执行模式。每次前向传播都会重新构建图（通过Autograd记录反向图），灵活性极高，但编译器难以在运行前获得完整图信息。PyTorch通过TorchScript（追踪+脚本）和FX（符号追踪）尝试弥补这一缺陷。

3.1.2 判断方法：如何识别框架图的优化潜力

评估一个框架图表示是否适合编译器优化，可以从以下维度判断：

维度	TensorFlow GraphDef	PyTorch TorchScript	PyTorch FX Graph
图构建时机	运行前（静态）	追踪或脚本编译时	符号追踪时
控制流支持	强（tf.cond/while）	中（脚本模式支持）	弱（需手动展开）
动态形状	受限（需placeholder）	支持（但优化受限）	支持（但需约束）
优化友好度	高（全局可见）	中（混合语义）	中高（纯Python语义）
调试难度	低（难以断点）	高（可逐步调试）	中（可打印图）

实际工程中，编译器导入框架图时，需要做规范化（Canonicalization）：将框架特有的算子映射到编译器内部IR。例如，将TensorFlow的tf.nn.conv2d和PyTorch的torch.nn.functional.conv2d统一映射为linalg.conv_2d方言操作。

3.1.3 案例拆解：PyTorch模型导出到编译器的完整链路

以一个ResNet-50模型为例，展示PyTorch图如何被编译器消费：

1. 初始状态：PyTorch模型处于Eager模式，每次调用forward()动态构建图。
2. 第一步：TorchScript追踪。使用torch.jit.trace(model, example_input)，记录实际执行路径，生成ScriptFunction。此时图是"扁平化"的，包含所有张量操作。
3. 第二步：FX符号追踪。使用torch.fx.symbolic_trace(model)，生成GraphModule。FX图保留了Python语义（如get_attr、call_function），更易于编译器解析。
4. 第三步：导入编译器。TVM的from_pytorch或XLA的PyTorch/XLA桥接层将FX图转换为编译器IR。关键转换包括：

   // FX图中的节点
   %conv1 = call_function[target=torch.nn.functional.conv2d](%input, %weight, %bias, stride=[2,2])
   // 转换为TVM Relay
   %1 = nn.conv2d(%input, %weight, strides=[2,2], padding=[3,3])

5. 结果对比：Eager模式延迟约15ms/step，经TorchScript+TVM优化后降至8ms/step，再经TVM auto-tuning后降至5ms/step。代价是编译时间增加约2分钟（含自动调优）。

3.2 编译器架构对比：TVM、XLA、Glow的设计哲学

3.2.1 核心原理：三种IR设计路线

三大编译器分别代表了不同的设计思路，理解它们的差异有助于把握MLIR的定位。

• XLA（Accelerated Linear Algebra）：Google出品，深度绑定TensorFlow。采用两段式IR：高层HLO（High-Level Optimizer）负责代数化简、算子融合；低层LLO（Low-Level Optimizer）负责内存分配、流水线调度。XLA的强项在于激进融合——将多个小算子融合为单个kernel，减少内存往返。
• TVM：Apache开源，框架无关。采用三段式IR：Relay（高层图IR）负责图优化；TIR（Tensor IR）负责循环变换、向量化；后端代码生成。TVM的核心竞争力是AutoTVM/Ansor自动调优，通过搜索空间探索最优kernel配置。
• Glow：Facebook开源，主打低延迟推理。采用单层IR（Glow IR），直接面向硬件指令。Glow的独特之处在于图调度器——显式管理内存生命周期，避免动态分配。

3.2.2 对比分析：关键架构维度

维度	TVM	XLA	Glow
IR层级	Relay + TIR（多级）	HLO + LLO（两级）	单级Glow IR
框架支持	TF/PyTorch/ONNX等	TF/JAX（深度绑定）	PyTorch/Caffe2（ONNX）
优化策略	算子融合 + 自动调优	代数化简 + 激进融合	内存规划 + 指令级优化
硬件后端	CPU/GPU/FPGA/自定义	CPU/GPU/TPU	CPU/GPU/FPGA
动态形状	支持（Relay支持动态shape）	有限支持（需bounded）	不支持（静态shape）
可扩展性	高（Pass机制）	中（HLO扩展需改C++）	低（IR固定）

从表中可见，TVM的多级IR设计最接近MLIR的理念，但TVM的IR是封闭的——自定义算子需要侵入式修改Relay/TIR。XLA的HLO虽然清晰，但缺乏通用性。Glow则过于特化，难以扩展到新硬件。

3.2.3 实战场景：选择编译器的决策路径

在实际项目中，选择编译器需考虑以下因素：

1. 框架绑定度：若已深度使用TensorFlow且无迁移计划，XLA是低风险选择；若使用PyTorch且需要多后端，TVM更合适。
2. 部署形态：边缘设备推理（低延迟、小内存）优先考虑Glow；云端训练+推理优先XLA；异构硬件（NPU/FPGA）优先TVM。
3. 算子覆盖：检查目标编译器是否支持模型中的自定义算子。TVM通过TE（Tensor Expression）可快速添加新算子；XLA需要编写HLO的C++实现；Glow则需要修改IR。

3.2.4 常见误区与总结复盘