第2章：计算图与IR——让编译器看懂你的模型

说实话，我刚入行做AI推理优化那会儿，最头疼的就是模型格式。PyTorch训练好的模型，转成ONNX报错，转成TensorRT又报错，折腾半天发现是某个算子不支持。那时候我就想，要是有一套统一的中间表示，大家都能看懂，该多好。

后来我明白了，这就是IR要做的事。IR，全称是Intermediate Representation，翻译过来就是中间表示。它就像编译器世界的通用语言。你写Python也好，C++也罢，最终都要翻译成机器码。但直接翻译太复杂，所以中间加一层IR，分步优化。

计算图，就是IR最常见的一种形式。说白了，就是把神经网络的计算过程，画成一张有向图。节点是算子，边是张量流动。你想想看，这多直观。

2.1 计算图的基本概念

计算图，我习惯叫它“数据流图”。每个节点代表一个操作，比如卷积、ReLU、矩阵乘。每条边代表数据的流向，也就是张量。

举个例子，一个简单的表达式：z = (x + y) * w。画成计算图就是：

    x ---\
          + ---\
    y ---/      * --- z
                /
           w --/

节点有3个：加法节点、乘法节点。边有4条：x、y、w、中间结果。就这么简单。

我在项目中遇到过一个问题：模型跑着跑着显存爆了。后来一查，是计算图里有个巨大的中间张量没及时释放。这就是计算图分析的价值——你能看到每个节点的内存占用。

计算图有两种模式：

• 静态图：先定义图，再执行。TensorFlow 1.x就是典型。优点是优化空间大，缺点是不够灵活。
• 动态图：边定义边执行。PyTorch就是典型。优点是灵活，缺点是优化空间小。

嗯，这里要注意：现在很多框架都在做“动静结合”。比如PyTorch的torch.compile，就是先把动态图trace成静态图，再做优化。说白了，静态图才是编译器真正喜欢的东西。

2.2 中间表示（IR）的作用

IR到底有什么用？我总结了三句话：

1. 解耦前端和后端：前端只管把模型转成IR，后端只管把IR转成目标代码。各管各的，互不干扰。
2. 提供优化入口：在IR层面做算子融合、常量折叠、死代码消除，比在源码层面容易得多。
3. 跨平台支持：一套IR，可以生成CPU、GPU、NPU等多种后端的代码。

我曾经踩过一个坑：自己写了个算子融合pass，结果把IR的依赖关系搞乱了，模型推理结果全错。后来花了整整两天才定位到问题。所以啊，IR的优化一定要保证语义等价，这是底线。

2.3 常见的IR格式

现在业界主流的IR格式，我接触过的有这三种：ONNX、MLIR、TVM Relay。它们各有各的脾气。

IR格式	设计哲学	适用场景	我的评价
ONNX	模型交换标准	跨框架模型转换	通用性强，但优化能力弱
MLIR	多级IR，可扩展	编译器基础设施	灵活，但学习曲线陡
TVM Relay	面向张量计算	推理优化	算子融合做得好

2.3.1 ONNX

ONNX，全称Open Neural Network Exchange。名字就说明了它的定位——开放、交换。说白了，它就是想让PyTorch、TensorFlow、Caffe这些框架能互相串门。

ONNX的IR是protobuf格式的。一个ONNX模型，就是一个protobuf文件。里面定义了计算图、算子、张量类型等信息。

我刚开始用ONNX时，最常遇到的问题是算子不支持。比如PyTorch里的某个自定义算子，ONNX标准里没有。这时候就得自己写自定义算子，或者用组合算子替代。嗯，挺麻烦的。

// ONNX的算子定义示例（protobuf格式）
node {
  input: "input"
  input: "weight"
  output: "output"
  op_type: "Conv"
  attribute {
    name: "kernel_shape"
    ints: 3
    ints: 3
  }
  attribute {
    name: "strides"
    ints: 1
    ints: 1
  }
}

2.3.2 MLIR

MLIR，全称Multi-Level Intermediate Representation。这是Google搞出来的东西，设计理念很前卫。它不只有一层IR，而是有多层。从高层到底层，逐层lowering。

举个例子，你可以用TOSA（Tensor Operator Set Architecture）这种高层IR表示卷积，然后lowering到LLVM IR这种底层IR。每一层都有自己的优化pass。

我个人觉得MLIR最大的优点是可扩展性。你可以自定义dialect（方言），定义自己的算子和类型。这对于做硬件编译器的人来说，简直是神器。

但说实话，MLIR的学习曲线确实陡。我当初看MLIR的文档，看了三遍才搞懂它的核心概念。如果你刚开始接触，建议先从Toy Tutorial入手。

2.3.3 TVM Relay

TVM Relay是Apache TVM项目里的IR。它专门为张量计算设计，所以对神经网络的支持特别好。

Relay IR的特点是：

• 函数式风格，每个表达式都是不可变的
• 支持自动微分，可以求梯度
• 内置了丰富的算子融合pass

我记得有一次，我用TVM优化一个BERT模型。默认情况下，推理延迟是50ms。我手动调了几个Relay的融合策略，把延迟降到了35ms。这就是IR优化的威力。

// TVM Relay IR示例
def @main(%x: Tensor[(1, 3, 224, 224), float32]) {
  %1 = nn.conv2d(%x, meta[relay.Constant][0], padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3])
  %2 = nn.relu(%1)
  %3 = nn.max_pool2d(%2, pool_size=[2, 2], strides=[2, 2])
  %4 = nn.conv2d(%3, meta[relay.Constant][1], padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3])
  %5 = nn.relu(%4)
  %6 = nn.max_pool2d(%5, pool_size=[2, 2], strides=[2, 2])
  %7 = nn.batch_flatten(%6)
  %8 = nn.dense(%7, meta[relay.Constant][2])
  %9 = nn.softmax(%8)
  %9
}

你看，Relay的IR可读性很强。每个算子都一目了然。而且它支持let binding，可以复用中间结果。

2.4 知识体系总览

为了让你更直观地理解本章的知识结构，我画了一张图：

这张图把本章的核心内容串起来了。左边是计算图，中间是IR的作用，右边是常见的IR格式。底部是它们之间的关系。你保存下来，以后复习时看一眼就能想起来。

2.5 小结

这一章我们聊了计算图和IR。说白了，计算图就是模型的骨架，IR就是编译器理解这个骨架的语言。没有IR，编译器就是个瞎子，什么都优化不了。

ONNX、MLIR、TVM Relay，这三者各有千秋。ONNX适合做模型交换，MLIR适合做编译器基础设施，TVM Relay适合做推理优化。选哪个，取决于你的具体场景。

我个人建议，如果你是做模型部署的，先从ONNX入手。如果你是做编译器开发的，MLIR是必学的。如果你是做推理引擎的，TVM Relay值得深入研究。

嗯，这一章就到这里。下一章我们会深入算子融合的具体技术，到时候我会分享一些我在实际项目中用到的融合策略和踩过的坑。