计算图与中间表示：从概念到实践

各位同学，今天我们来聊聊编译器里最核心的一个概念——计算图与中间表示。说实话，我刚入行那会儿，觉得这东西就是画几个圈圈连几条线，没什么大不了的。直到有一次，我接手一个NPU项目，发现模型跑出来的结果总是对不上……嗯，从那以后，我才真正理解了计算图的重要性。

计算图：AI模型的骨架

计算图是什么？说白了，它就是AI模型的骨架。你把神经网络里的每个算子（比如卷积、池化、全连接）看作一个节点，把数据流动的方向看作边，整个图就出来了。

我个人习惯把计算图分成两类：静态图和动态图。这两者的区别，我当年可是踩过坑的。

静态图 vs 动态图

静态图，就是你先搭好整个网络结构，再往里灌数据。TensorFlow 1.x 就是典型代表。动态图呢，边搭边算，PyTorch 就是这种风格。

为什么会有这种区别？你想想看，静态图的好处是编译器能提前看到整个计算流程，可以做很多优化。比如算子融合、内存复用，这些都需要全局信息。但缺点也很明显——调试起来特别痛苦。我记得有一次，我在TensorFlow里写了一个复杂的条件分支，结果图死活构建不对，查了半天才发现是某个张量的shape没对齐。

动态图就灵活多了。你可以像写普通Python代码一样写网络，随时打印中间结果。但代价是，编译器没法做全局优化。每次前向传播，图结构都可能变化。

现在的主流做法是什么？先动态后静态。PyTorch 2.0 引入了 torch.compile，就是先把动态图 trace 成静态图，再做优化。这个思路，我个人觉得非常聪明。

核心要点：静态图适合部署，动态图适合研发。两者不是非此即彼，而是可以互补的。

中间表示（IR）：编译器的通用语言

有了计算图，接下来就是怎么表示它。这就引出了中间表示（IR）。IR 是编译器的内部语言，它把上层框架的多样性统一起来。

我见过很多团队，一开始觉得 IR 无所谓，随便搞搞就行。结果项目做到一半，发现要加新算子，要支持新硬件，IR 根本不够用……最后只能推倒重来。所以，IR 的设计原则，我建议你从一开始就重视。

IR 的设计原则

我总结了三条核心原则：

  • 完备性：IR 必须能表达所有你需要的计算模式。比如，你支持卷积，那 IR 里就得有 Conv2D 这个节点。别想着用几个基础算子拼凑，那样效率会很低。
  • 可优化性：IR 要方便做各种变换。比如算子融合、常量折叠、死代码消除。如果 IR 结构太复杂，优化 pass 写起来会非常痛苦。
  • 可扩展性：新算子、新硬件、新优化策略，都能方便地加进去。我曾经见过一个 IR，为了支持一个新硬件，改了三个月……那叫一个酸爽。

避坑指南：我曾经在项目初期，为了省事，直接用 Python 对象做 IR。结果后期性能优化时，发现 Python 对象太慢了，根本没法做大规模图变换。后来全部重写为 C++ 结构体，性能提升了 10 倍以上。

多级IR的层次结构

一个成熟的编译器，不会只有一种 IR。为什么？因为不同阶段，你需要不同的抽象层次。

我习惯把 IR 分成三级：

层级 名称 特点 典型操作
高层IR Graph IR 接近原始计算图，保留算子语义 算子融合、布局转换
中层IR Tensor IR 张量运算级别，引入循环结构 循环分块、向量化
底层IR 指令IR 接近硬件指令，处理寄存器、内存 指令调度、寄存器分配

举个例子，你有一个 Conv + ReLU + Conv 的序列。在高层IR里，它们还是三个独立的算子。到了中层IR，编译器可能会把 ReLU 融合进第一个 Conv，变成 ConvReLU。到了底层IR，就要考虑怎么把数据从全局内存搬到寄存器，怎么利用向量指令并行计算。

为什么要分这么多层？你想想看，如果一开始就在底层IR做算子融合，那代码会复杂到没法看。每一层只关注自己该做的事，这就是分治思想。

高层IR：Graph IR

Graph IR 是最接近原始模型的。它保留了算子的语义信息，比如 Conv2D、BatchNorm、Relu。这个阶段的优化，主要是图级别的变换。

我举个例子，常量折叠。如果你的模型里有一个 Add 节点，两个输入都是常量，那编译器可以直接算出结果，把这个节点替换成一个常量节点。这种优化，在高层IR做最合适。

// 高层IR示例（伪代码）
Graph {
  Node: Conv2D(input=A, weight=W, stride=1)
  Node: ReLU(input=Conv2D)
  Node: Add(input=ReLU, bias=B)
}

中层IR：Tensor IR

到了中层IR，我们开始关注张量的具体计算过程。这里会引入循环结构，因为卷积、矩阵乘法这些操作，本质上都是多层循环。

中层IR的优化，主要是循环变换。比如循环分块（tiling），把大循环拆成小循环，提高缓存命中率。还有向量化，把标量运算变成向量运算。

// 中层IR示例（伪代码）
for i in range(0, N, 16):
  for j in range(0, M, 16):
    C[i:i+16, j:j+16] = matmul(A[i:i+16, :], B[:, j:j+16])

注意：中层IR的优化，对NPU性能影响最大。我见过一个项目，同样的模型，优化前后性能差了5倍。关键就在于循环分块的大小和向量化的策略。

底层IR：指令IR

底层IR是最接近硬件的。这里要考虑寄存器分配、内存访问模式、指令流水线等问题。

对于NPU来说，底层IR尤其重要。因为NPU通常有专用的向量单元、矩阵单元、以及复杂的内存层次结构。底层IR要能精确描述这些硬件特性。

// 底层IR示例（伪代码）
LOAD R1, [A_addr]
LOAD R2, [B_addr]
VECTOR_MUL R3, R1, R2
STORE [C_addr], R3

嗯，这里要注意，底层IR的优化，往往需要硬件工程师的参与。因为只有他们最清楚，某个指令序列在硬件上到底是怎么执行的。

总结

计算图是AI模型的骨架，IR是编译器的通用语言。静态图和动态图各有优劣，现在的主流做法是先动态后静态。IR的设计要遵循完备性、可优化性、可扩展性。多级IR的层次结构，让编译器能分阶段做优化，每一层只关注自己该做的事。

我个人觉得，理解这些概念最好的方式，就是动手写一个简单的编译器。哪怕只是把几个算子融合一下，你都会对IR的设计有更深的理解。好了，今天就到这里，下次我们聊聊算子融合的具体实现。

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