【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 一个最简单的例子
importtvmfromtvmimportrelay# 定义输入和参数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和Modulefunc=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层：我知道这里要做矩阵乘# 具体怎么算？交给TIRout=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_1for(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")# 外层绑定到blockIdxsch.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++编译好的函数，就需要一个统一的调用约定。
importtvm# 加载编译产物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的四层架构：
- 前端（FrontEnd）：把不同框架的模型解析成TVM IR
- Relay：静态图函数式IR，强类型、Pass丰富，适合CNN/BERT
- Relax：动态图IR，支持动态shape、控制流、call_tir，适合LLM/MoE
- TIR：底层张量IR，描述循环、buffer、线程，是调度和Codegen的基础
- 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
- 更多内容持续更新中……
