4、中间表示层（IR）：计算图的中间表示形式，从Graph IR到Linear IR的演进

说到AI芯片的软件栈，中间表示层（IR）是个绕不开的话题。很多人觉得IR就是个“中间过渡”，没什么技术含量。我刚开始做编译器的时候也这么想，直到被一个Graph IR的优化问题折磨了整整两周……嗯，从那以后我再也不敢小看这层了。

说白了，IR就是计算图的“中间语言”。它把前端框架（PyTorch、TensorFlow）的模型描述，翻译成后端芯片能理解的指令。这个翻译过程不是一步到位的，而是分阶段演进：从Graph IR到Linear IR。

4.1 为什么需要多层IR？

你想想看，前端框架的模型五花八门。PyTorch的动态图、TensorFlow的静态图、ONNX的中间格式……每个都有自己的表达方式。如果直接把这些映射到芯片指令，那编译器得疯掉。

所以我们需要一个“中间人”。这个中间人分两层：

• Graph IR：保留计算图的结构信息，适合做算子融合、内存规划等高层次优化。
• Linear IR：把图展开成线性指令序列，适合做寄存器分配、指令调度等低层次优化。

我个人习惯把Graph IR比作“建筑设计图”，Linear IR比作“施工图纸”。设计图告诉你房间怎么布局，施工图纸告诉你砖头怎么砌。两者缺一不可。

4.2 Graph IR：计算图的“骨架”

Graph IR的核心是“节点+边”。节点代表算子（Conv、ReLU、Add），边代表数据依赖（张量流动）。这个结构天然适合做全局优化。

我在项目中遇到过一个问题：一个BERT模型在Graph IR层面做算子融合时，发现多个连续的MatMul和Add可以合并成一个FusedMatMulAdd。这个优化让推理速度提升了15%。如果等到Linear IR再做，就来不及了——因为那时候图结构已经丢失了。

// 伪代码示例：Graph IR表示
Graph {
  Node: Conv2d (input: tensor_A, weight: tensor_B)
  Node: ReLU (input: tensor_C)
  Node: Add (input: tensor_D, tensor_E)
  Edge: tensor_A -> Conv2d -> tensor_C
  Edge: tensor_C -> ReLU -> tensor_D
  Edge: tensor_D -> Add -> tensor_F
}

Graph IR的优化主要围绕“图变换”展开：

• 算子融合：把多个小算子合并成一个大算子，减少内存访问。
• 常量折叠：把编译时就能算出来的常量提前计算好。
• 死代码消除：去掉那些输出不被使用的节点。
• 内存规划：分析张量的生命周期，决定内存复用策略。

4.3 Linear IR：从“图”到“线”

Graph IR虽然好，但它有个问题：图结构太“松散”了。芯片执行指令时，是一条一条顺序执行的。所以我们需要把图“拍平”成线性序列——这就是Linear IR。

Linear IR的核心是“基本块+控制流”。基本块是一段顺序执行的指令序列，控制流负责处理分支、循环等跳转逻辑。

// 伪代码示例：Linear IR表示
BasicBlock_0:
  LOAD tensor_A, addr_0
  LOAD tensor_B, addr_1
  CONV2D tensor_C, tensor_A, tensor_B
  STORE tensor_C, addr_2
  JUMP BasicBlock_1

BasicBlock_1:
  LOAD tensor_D, addr_2
  RELU tensor_E, tensor_D
  STORE tensor_E, addr_3
  JUMP BasicBlock_2

从Graph IR到Linear IR的转换，我称之为“图展开”。这个过程有几个关键步骤：

1. 拓扑排序：把图中的节点按依赖关系排成线性顺序。
2. 基本块划分：根据控制流（if-else、loop）把指令分组。
3. 虚拟寄存器分配：为每个张量分配一个虚拟寄存器编号。
4. 控制流图构建：记录基本块之间的跳转关系。

为什么需要虚拟寄存器？因为这时候我们还没决定哪些数据放寄存器、哪些放内存。虚拟寄存器就是个“占位符”，等后面做寄存器分配时再决定。

4.4 Graph IR vs Linear IR：一张表说清楚

说白了，Graph IR负责“做什么”，Linear IR负责“怎么做”。前者是战略层，后者是战术层。

4.5 从Graph IR到Linear IR的演进：一个具体例子

我们拿一个简单的计算图来演示：Conv2d -> ReLU -> Add。

Step 1：Graph IR表示

Graph {
  Node_0: Conv2d (input: X, weight: W) -> Y
  Node_1: ReLU (input: Y) -> Z
  Node_2: Add (input: Z, bias: B) -> Out
}

Step 2：图展开（拓扑排序）

顺序：Node_0 -> Node_1 -> Node_2

Step 3：Linear IR表示（带虚拟寄存器）

BasicBlock_0:
  %v0 = LOAD X
  %v1 = LOAD W
  %v2 = CONV2D %v0, %v1
  %v3 = RELU %v2
  %v4 = LOAD B
  %v5 = ADD %v3, %v4
  STORE %v5, Out

你看，从Graph IR到Linear IR，信息没有丢失，但结构变了。图变成了线，节点变成了指令，边变成了数据依赖（通过虚拟寄存器体现）。

4.6 业界主流IR方案

不同的AI编译器，IR的设计也不一样。我整理了几个主流方案：

• TVM的Relay IR：Graph IR的代表，支持函数式编程风格，适合做高阶优化。
• MLIR的Dialect机制：可以自定义多层IR，从高层的TOSA Dialect到底层的LLVM Dialect。
• XLA的HLO IR：Google的Graph IR，支持跨设备优化（CPU/GPU/TPU）。
• Glow的IR：Facebook的Linear IR风格，直接面向硬件指令。

我个人比较喜欢MLIR的方案。它允许你在同一个框架里定义多个IR层次，从Graph IR到Linear IR的转换就是Dialect之间的转换。这样既保持了灵活性，又降低了维护成本。

4.7 本章小结

中间表示层是AI编译器的心脏。从Graph IR到Linear IR的演进，本质上是“从抽象到具体”的过程。Graph IR保留计算图的结构，方便做全局优化；Linear IR把图展开成指令序列，方便做底层调度。

嗯，这里要注意：不要试图跳过Graph IR直接做Linear IR。我见过一些团队为了省事，直接从框架模型映射到芯片指令，结果优化空间小得可怜。Graph IR的全局优化价值，是Linear IR无法替代的。

最后留个思考题：如果你的芯片有专用的矩阵计算单元（比如NVIDIA的Tensor Core），Graph IR和Linear IR的设计会有什么不同？想清楚这个问题，你对IR的理解就更深一层了。