【AI编译器入门②】TVM架构拆解：Relay、Relax、TIR与Runtime四层结构

【AI编译器入门②】TVM架构拆解：Relay、Relax、TIR与Runtime四层结构

本文是「AI编译器」系列的第二篇。TVM是学术界和工业界最成熟的AI编译器之一。很多人学TVM时容易被一堆名词搞晕：Relay、Relax、TIR、TE、Schedule、Runtime......它们都是什么？有什么关系？本文把TVM的四层架构拆开讲清楚，帮你建立整体认知。

---

一、写在前面：TVM为什么要有多层IR？

学TVM的第一个困惑就是：为什么要有这么多层IR？一层不行吗？

答案是：不同的抽象层次，解决不同的问题。

想象一下造房子：

• 设计图纸（高层抽象）→ 表达"房子长什么样"
• 施工图（中层抽象）→ 表达"墙怎么砌、梁怎么搭"
• 施工指令（底层抽象）→ 表达"先搬砖、再和灰、再砌墙"

每一层都有存在的意义------高层改起来方便，底层精确可控。

TVM的多层IR也是一个道理：

• Relay/Relax（高层IR）：描述"算什么"（图结构、算子语义）
• TIR（底层IR）：描述"怎么算"（循环、内存、线程）
• Runtime（执行层）：描述"怎么跑"（加载、调度、执行）

分层的好处是：各层可以独立优化，不用一层承担所有职责。

先上一张架构图：

---

二、TVM四层架构总览

TVM架构可以粗略分为四层，从高到低：

┌─────────────────────────────────────────┐
│         TVM Runtime / RPC               │  执行引擎
├─────────────────────────────────────────┤
│   TIR (Tensor IR) + TE/Schedule         │  底层IR + 调度原语
├─────────────────────────────────────────┤
│      Relax（动态图）/ Relay（静态图）     │  高层IR
├─────────────────────────────────────────┤
│  FrontEnd：ONNX / PyTorch / TensorFlow  │  前端解析
└─────────────────────────────────────────┘

数据流转的方向是从上到下：

ONNX / PyTorch 模型
  → FrontEnd解析为Relay或Relax（高层IR）
  → MiddleEnd Pass优化高层IR
  → Lower到TIR（底层IR）
  → Schedule优化 / Codegen生成CUDA/LLVM代码
  → Runtime加载执行

我们逐层往下看。

---

三、Relay：静态图IR

3.1 什么是Relay？

Relay是TVM的函数式静态图IR ，适合CNN、BERT等静态shape模型。

两个核心概念：

• 函数式表达式：计算图是嵌套的表达式，无副作用（就像数学函数，给定输入就有确定输出）
• 强类型系统：每个表达式都有明确的类型和shape，便于编译期优化

3.2 一个最简单的例子

import tvm
from tvm import relay

# 定义输入和参数
x = relay.var("x", shape=(1, 784), dtype="float32")
w = relay.var("w", shape=(784, 128), dtype="float32")
b = relay.var("b", shape=(128,), dtype="float32")

# 构建图：y = relu(x @ w + b)
fc = relay.nn.dense(x, w)       # 全连接（矩阵乘）
out = relay.nn.bias_add(fc, b)   # 加偏置
out = relay.nn.relu(out)         # ReLU激活

# 打包成Function和Module
func = relay.Function([x, w, b], out)
mod = tvm.IRModule.from_expr(func)
mod = relay.transform.InferType()(mod)
print(mod)

输出大概长这样：

def @main(%x: Tensor[(1, 784), float32],
          %w: Tensor[(784, 128), float32],
          %b: Tensor[(128), float32]) {
  %0 = nn.dense(%x, %w, units=128)
  %1 = nn.bias_add(%0, %b)
  nn.relu(%1)
}

是不是和ONNX很像？都是节点+边的图结构。

3.3 Relay的特点

特点	说明
静态shape	所有张量维度在编译期已知
函数式	无变量赋值，用表达式嵌套表达计算
Pass丰富	TVM早期积累的优化Pass大多基于Relay
动态shape弱	需要用VM或bucket编译来workaround

3.4 Relay模块结构

IRModule（整个模块）
  ├── GlobalVar @main（主函数名）
  ├── Function（函数：参数、函数体、返回类型）
  ├── Expression（表达式）
  │     ├── Var / Constant / Call / Tuple
  │     └── Let / Function / If
  └── Type（类型）
        ├── TensorType (shape, dtype)
        └── FuncType (arg_types, ret_type)

---

四、Relax：新一代动态图IR

4.1 为什么需要Relax？

Relay虽然好，但面对大模型推理时，有几个硬伤：

• 动态shape支持弱（batch size、sequence length都是变的）
• 控制流支持有限（if/while不好表达）
• 大模型场景下笨重

Relax就是为了解决这些问题而生的------它是TVM的新一代IR，专门为LLM、MoE等动态场景设计。

4.2 Relay vs Relax 对比表

维度	Relay	Relax
动态shape	弱，需bucket/VM	原生支持
控制流	有限	原生支持if/while
设计目标	CNN/BERT等静态模型	LLM/MoE等大模型
算子语义	函数式表达式	绑定式 + call_tir
Pass写法	ExprMutator	BlockBuilder + 模式匹配

4.3 Relax的关键抽象

BindingBlock（绑定块）

Relax不用函数式嵌套，而是用顺序绑定的方式------就像写普通代码一样，一行一行赋值：

with R.dataflow():
    lv0 = R.matmul(x, w)       # 第一行：矩阵乘
    lv1 = R.add(lv0, b)        # 第二行：加偏置
    out = R.relu(lv1)          # 第三行：激活
    R.output(out)

这种写法更符合人类的直觉，也更容易处理动态shape和控制流。

call_tir：Relax与TIR的桥梁

Relax里有一个非常重要的概念：call_tir。

什么意思呢？Relax只管高层的图结构，具体的计算（比如矩阵乘怎么算），它不关心------直接调用底层的TIR PrimFunc。

# Relax层：我知道这里要做矩阵乘
# 具体怎么算？交给TIR
out = relax.call_tir("matmul_tir", [x, w], out_sinfo=...)

好处：高层图和底层计算解耦了。

• Relax层专心做图优化（融合、布局、量化）
• TIR层专心做调度优化（分块、向量化、线程绑定）
• 互不干扰，各干各的

4.4 StructInfo：动态shape的类型系统

Relay的类型是固定shape的，比如Tensor[(1, 784), float32]。

Relax的类型可以有符号变量，比如Tensor[("batch", 784), float32]------batch就是个符号变量，运行时才知道具体值。

这就是StructInfo------不仅描述dtype和shape，还描述符号变量之间的关系。

---

五、TIR：底层张量IR

5.1 什么是TIR？

TIR（Tensor IR）是TVM的底层IR，描述循环、buffer、线程、存储等细节。后端Codegen就是基于TIR生成CUDA/LLVM代码的。

如果说Relay/Relax是"设计图纸"，那TIR就是"施工方案"------精确到每一步怎么做。

5.2 核心抽象

抽象	含义	类比
Buffer	多维张量存储	全局内存/共享内存/寄存器
Loop	循环嵌套	for循环
Block	TensorIR调度单元	计算块（可以被分块、重排）
Thread	线程绑定	threadIdx.x, blockIdx.x
PrimFunc	TIR函数	可以被Relax通过call_tir调用

5.3 TIR长什么样？

一个简单的向量加TIR大概长这样：

primfn(A_1: handle, B_1: handle) -> () {
    A = A_1
    B = B_1
    for (i: int32, 0, 1024) {
        B[i] = (A[i] + 1f)
    }
}

看起来像C代码对吧？但它不是普通的C------它是可以被调度和优化的IR。你可以对它做split（分块）、reorder（重排）、vectorize（向量化）、bind（绑定线程）等操作。

5.4 Schedule：调度原语

Schedule就是对TIR做变换的操作。比如：

sch = tir.Schedule(func)
i = sch.get_loops("B")[0]
io, ii = sch.split(i, factors=[32, 32])  # 把循环拆成两层
sch.bind(io, "blockIdx.x")                  # 外层绑定到blockIdx
sch.bind(ii, "threadIdx.x")                 # 内层绑定到threadIdx

这就是CUDA编程的"Grid/Block/Thread"在TIR中的表达。

---

六、TVM Runtime：执行引擎

6.1 Runtime做什么？

Runtime负责加载编译产物、管理内存、发射kernel。它是连接"编译"和"运行"的桥梁。

核心模块：

Module：加载.so / .tar / .aom文件
PackedFunc：统一的函数调用接口
DeviceAPI：CPU / CUDA / ROCm / OpenCL
Object：NDArray、Module、String等对象系统

6.2 PackedFunc：TVM的"万能函数指针"

PackedFunc是TVM Runtime的核心抽象------它是一个统一的函数调用接口，允许C++、Python、RPC之间无缝互相调用函数。

为什么需要这个？因为TVM的编译器在C++层，但用户接口在Python层------Python要调用C++编译好的函数，就需要一个统一的调用约定。

import tvm

# 加载编译产物
mod = tvm.runtime.load_module("model.so")
# 获取函数
ex = mod.get_function("main")
# 调用执行
out = ex(input_ndarray)

看起来很简单，但背后PackedFunc做了很多事情------参数类型检查、内存管理、设备切换......

---

七、四层协作：一个模型的完整旅程

最后，我们把四层串起来，看一个PyTorch模型在TVM里的完整旅程：

PyTorch模型
  ↓ from_pytorch（前端）
Relay计算图
  ↓ Pass优化（中端）
    FoldConstant → 常量折叠
    FuseOps → 算子融合
    ...
优化后的Relay图
  ↓ Lower到TIR（降低）
TIR PrimFunc（带循环和buffer）
  ↓ Schedule优化
    split → 分块
    vectorize → 向量化
    bind → 线程绑定
优化后的TIR
  ↓ Codegen
CUDA C++代码
  ↓ nvcc编译
.so文件
  ↓ Runtime加载执行
GPU上跑推理

---

八、常见误区

误区一："Relay和Relax只是写法不同"

错。Relay是函数式静态图IR，Relax是支持动态shape和命令式绑定的IR。Relax的算子通过call_tir显式连接TIR，而Relay是隐式lower。完全是两种设计哲学。

误区二："TIR就是打印出来的C代码"

TIR是可被调度和分析的IR。Codegen前可以对TIR做各种优化（split、reorder、vectorize、bind thread），这些优化直接影响最终性能。TIR是"活"的，C代码是"死"的。

误区三："Runtime只负责加载.so"

TVM Runtime还提供跨语言调用（PackedFunc）、跨设备内存管理（NDArray）、RPC等功能。理解Runtime是部署和调试的基础。

---

九、总结一下

今天这篇文章，我们拆解了TVM的四层架构：

1. 前端（FrontEnd）：把不同框架的模型解析成TVM IR
2. Relay：静态图函数式IR，强类型、Pass丰富，适合CNN/BERT
3. Relax：动态图IR，支持动态shape、控制流、call_tir，适合LLM/MoE
4. TIR：底层张量IR，描述循环、buffer、线程，是调度和Codegen的基础
5. Runtime：执行引擎，提供Module、PackedFunc、DeviceAPI

下一篇预告：我们深入编译前端------ONNX和TorchScript是怎么解析的？动态shape怎么处理？自定义算子怎么注册？

---

📚 本系列导航：

• 第1篇：AI编译器基础：三段式Pipeline与全链路协同
• 第2篇：TVM架构拆解：Relay、Relax、TIR与Runtime（本文）
• 第3篇：编译前端：ONNX、TorchScript与动态shape处理
• 第4篇：编译中端：图优化Pass全解析
• 第5篇：编译后端：CUDA Codegen、CUTLASS与TensorRT
• 更多内容持续更新中......