【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 一个最简单的例子

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的特点

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 对比表

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 核心抽象

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的四层架构：

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
• 更多内容持续更新中……