想象一个大型建筑事务所。业主拿着餐巾纸上的涂鸦、手机里的参考照片、Excel里的房间清单来表达需求——这些“草稿”形式五花八门,但建筑师不能直接拿着它们上工地。他必须先画出一张标准化的设计蓝图:标注承重墙、门窗位置、水电走向,却不用操心用红砖还是空心砖。这张蓝图统一了所有表达,让结构工程师能审核优化,让施工队能照图施工。 AI 编译器 TVM 的世界里也有同样的故事。PyTorch、TensorFlow、ONNX 就像不同的业主,各自用不同的方式描述神经网络;而 GPU、CPU、NPU 就是施工队。TVM 需要一张“设计蓝图”来衔接两端——RELAY IR(Relay 中间表示) 就是这张蓝图。它用函数式编程语言的形式,把神经网络描述成结构严谨的程序,所有优化都在蓝图上进行,最后才翻译成“施工图”TIR,交给后端生成机器码。
理解 RELAY IR,你就握住了 TVM 编译流水线的总开关——所有自动优化、代码生成都是从读懂这张蓝图开始的。
可视化演示
展示RELAY IR三维度:Let绑定(let x=e1 in e2每个中间结果显式命名数据依赖和生命周期一目了然支持精确优化分析)、编译流水线(relay.build从IRModule经ANF规范化FoldConstant常量折叠FuseOps算子融合Lowering降级到Execute共6阶段PassManager管理优化pass)、执行后端对比(Graph Executor静态图一次性执行吞吐高95%不支持动态VM支持运行时分支和动态形状灵活性90%按场景选择),三模式切换理解TVM端到端编译流程。
本节配有交互动效(进入视口后自动加载,约 3KB)。
'+C1+'
IR体系定位
在一家建筑事务所里,图纸按成熟度分成三种:业主的初步设想——可能只是一张照片加几个箭头;设计蓝图——标清楚承重墙、门窗、楼梯和水电走向;施工图——细到每一颗螺丝的规格。RELAY IR 恰好对应中间那张设计蓝图:它既不是业主模糊的想法,也不是施工队手里繁琐的细节。
TVM 编译器的“图纸链”同样环环相扣。前端框架产出的模型(Model)经过 TVM 前端解析器(frontend)翻译成 RELAY IR;RELAY IR 经过一系列优化后被降级为 TIR(Tensor Intermediate Representation,张量中间表示);TIR 最后通过代码生成(codegen)变成可在 GPU、CPU 或 NPU 上执行的机器码。Relay 就卡在“用户描述”与“硬件执行”的正中间,承担桥梁角色。
从 IR 层级的视角看,RELAY IR 是高级中间表示:它保持函数式编程风格,程序由不可变的表达式组成;它完全不关心寄存器、线程、内存布局等硬件细节;它的数据流视图等价于一张计算图,却比计算图多出一层表达力——能用高阶函数和递归描述控制流。这张图“承上启下”,上层模型要落地必须先经过它。
这张蓝图之所以能“通吃”上下游,正是因为它站在一个恰到好处的抽象高度:足够具体,让优化器能分析数据依赖;又足够抽象,让所有硬件后端都能消费。理解了 Relay 的位置,你就理解了 TVM 为什么能把“一处编程、多处部署”变成现实。
- Model(模型)前端框架(如 PyTorch、TensorFlow)训练出来的神经网络模型。
- frontend(前端解析)TVM 前端解析器把模型翻译成 Relay 程序。
- Relay IR(设计蓝图)统一、规范、与硬件无关的高级中间表示。
- optimize(优化)在 Relay 上执行的多趟优化 pass。
- TIR(施工图)张量中间表示,携带循环、内存等低级信息。
- codegen(代码生成)后端代码生成器把 TIR 变成目标机器码。
Model → Relay IR → TIR → Machine Code设计目标
业主的草稿为什么不能直接施工?因为草图缺尺寸、比例乱、不同业主的画法千差万别。建筑师的职责是“重画一张”,回答三个问题:画得全(表达力)、看得懂(可分析)、改得动(可优化)。RELAY IR 的设计目标也围绕这三个词展开。
早期编译器常用计算图(Computation Graph)作为中间表示。计算图由节点(算子)和边(张量数据流)组成,直观易懂,但表达能力有限:它无法表达条件分支、循环、递归这类控制流。现代模型——BERT 的动态序列长度、强化学习策略网络的条件决策——用静态计算图往往力不从心。Relay 的诞生正是为了突破这个天花板。
Relay 的三个核心设计目标:第一,表达现代网络,借助函数式语言的抽象能力,支持控制流、递归和共享子图;第二,简化程序分析,程序由递归定义的表达式组成,每个子表达式都有明确类型和语法结构,pass 可以像结构工程师审图一样逐层检查;第三,统一前端入口,无论模型来自 PyTorch 还是 ONNX,最终都翻译成同一种 Relay 程序,编译器只需写一套优化逻辑。
这个公式揭示了一个重要的设计智慧:Relay 不是抛弃计算图,而是把计算图吸收为自身的数据流视图。当程序没有控制流时,Relay 就是一张图;当程序需要控制流时,图做不到的事,Relay 依然做得到。这正是它作为“通用蓝图”的底气——不牺牲表达能力换取规范性,而是二者兼得。
- 任意计算图 G传统静态计算图,表示张量的数据流关系。
- 存在 Relay 程序 e总能找到一个 Relay 表达式与之对应。
- 数据流视图把 Relay 程序按数据依赖展开得到的形态。
- 结构等价Relay 的数据流形态与计算图在结构上互通。
程序结构
拿到一张设计蓝图,你会先看图签栏里的工程名称、轴线编号和标高说明——规范让每个读者都能快速定位信息。Relay 程序也有固定骨架,浏览起来同样“处处有章法”。
Relay 程序的最外层容器是模块(Module),一个模块里装载一组具名函数。最常用的入口函数叫 @main,相当于图纸的“总图”。函数体内部采用数据流形式:一行一行绑定中间张量,行与行之间靠数据依赖串联。每个中间结果都用 % 开头命名,方便后续表达式引用。用一个极简例子感受一下:
这段程序描述了一个“全连接 → 加偏置 → ReLU 激活”的小网络。%0、%1、%2 是三个中间张量,后面的表达式可以直接引用前面的名字——这种命名绑定机制让程序的结构天然支持中间张量复用:同一个结果可以被多处使用,等价于计算图里的共享子图。Let 表达式就是这种绑定关系的抽象形态:
把 let 链条展开,你得到的就是一张数据流图;把数据流图压缩回嵌套表达式,你又得到了一段 Relay 程序。这种“图与程序”的等价性,是后面所有优化 pass 能够放心改写程序的根基。
- let 关键字声明一个局部绑定,类似“把这张图命名为 x”。
- x(绑定名)新引入的变量名,代表一个中间结果。
- e₁(绑定体)被绑定的子表达式,其计算结果就是 x 的值。
- in 分隔符标记绑定作用域的开始。
- e₂(使用体)在绑定生效的范围内使用 x 的后续表达式。
def @main(%x: Tensor[(10, 10), float32], %w: Tensor[(10, 10), float32], %b: Tensor[(10,), float32]) { %0 = nn.dense(%x, %w); %1 = nn.bias_add(%0, %b); %2 = nn.relu(%1); %2}
核心语法
蓝图的右下角总有一块“图例”:实线代表承重墙,虚线代表轻质隔墙,双圈代表排水立管。读懂图例才能读懂图纸。Relay 的“图例”是表达式(expression)——整个语言由五种基本表达式构建而成,任何一种网络结构都能用这五种积木拼出来。
五种表达式中,调用(Call) 是网络计算的发动机:一个算子在 Relay 里就是一个函数,把输入张量和参数喂给它,返回输出张量。匿名函数则让 Relay 能表达“函数作为参数传递”——比如把一段自定义的计算逻辑作为回调传给控制流原语。let 负责把纯函数式的组合风格拉回贴近直觉的“先算这个、再算那个”的数据流叙事。灵活组合这五种积木,线性网络、残差结构、动态分支都能被忠实描画。
- e(表达式)Relay 程序的任意表达式,它是语法的最小递归单元。
- c(常量)张量字面量,直接内嵌数据或引用元数据中的常量。
- x(变量引用)引用一个已绑定的张量名字,如
%x、%0。 - λx.e(匿名函数)定义接收参数 x 的函数体 e,对应函数式编程的“高阶能力”。
- e₁(e₂)(函数调用)把参数 e₂ 传给函数 e₁,
nn.dense(%x, %w)本质就是一次调用。 - let 绑定为表达式命名并限定作用域,是中间张量复用的语法载体。
类型系统
蓝图会标注每堵墙的厚度和材料——“外墙 240 厚,混凝土”,这样结构工程师不用猜也能审核。Relay 的类型系统就是蓝图上这种“标注规则”的化身:每个张量表达式都必须带类型,编译器据此做检查。
Relay 中最常见的类型是张量类型,写作 Tensor[(10, 10), float32]:前半部分是形状(shape),后半部分是元素数据类型(dtype)。有了这样的标注,一个 dense 算子的输入输出是否匹配,通过类型推导就能在编译期发现,而不是跑到 GPU 上炸出运行时错误。
类型系统还支持多态函子:函数可以声明类型参数,让同一个函数适配不同形状的张量,就像蓝图里同一型号的窗框可以装在不同尺寸的墙上。有了这套机制,Relay 既保证了“图纸”的严谨性,又没有牺牲“图纸”的通用性。
- Γ(类型环境)记录变量到类型的映射集合,类似“图纸的标注清单”。
- 前提一在环境 Γ 下,表达式 e₁ 的类型是 T₁。
- 环境扩展把 x 绑定到类型 T₁,得到新环境 Γ[x↦T₁]。
- 前提二在扩展环境中,表达式 e₂ 的类型是 T₂。
- 横线(推导规则)如果横线上方的前提都成立,则下方的结论成立。
- 结论整个 let 表达式的类型与 e₂ 一致,为 T₂。
编译流水线
设计蓝图定稿后,还要经过结构专业审查、机电专业会签,最后翻译成施工图交到工地。Relay 在 TVM 中的旅程完全复刻了这条流水线:构建 → 优化 → 降级。
构建阶段,TVM 前端把 PyTorch、TensorFlow 或 ONNX 模型解析成 Relay 表达式,包装成 IRModule。优化阶段,一串 pass 依次登场:常量折叠(ConstFold)把可以预先算好的节点算掉,推理简化(SimplifyInference)去掉训练专用的操作,算子融合(FuseOps)把相邻算子打包成一个 kernel。降级阶段,Relay 表达式被逐层翻译成 TIR——这时蓝图正式变成施工图,循环、内存布局、线程绑定全部显式出现,交给后端做进一步优化和代码生成。
算子融合是提升性能的关键一步。考虑两个相邻算子 f(dense)和 g(bias_add),传统方式先执行 f 把中间结果写进内存,再执行 g 读回来;融合后,两个计算合并在同一个 kernel 里,中间张量根本不用写进内存。
在 GPU 上,每次 kernel 启动都有不容忽视的固定开销,中间张量的读写更是耗电大户。融合让整个“dense → bias_add → relu“链条凝聚成一个 kernel,访存次数大幅下降。这就是为什么同一份 Relay 程序,经过优化后性能可以天差地别——蓝图还是那张蓝图,但”施工工艺”已经完全不同。
- f、g(相邻算子)数据流中紧挨着的两个计算,如 dense 和 bias_add。
- x(输入张量)进入算子链的原始数据。
- g(f(x))(未融合)先算 f 结果写内存,再读回算 g,多一次内存往返。
- (g∘f)(x)(融合后)合成一个函数,一次遍历完成全部计算。
AI 应用场景
多前端统一部署
阶段:模型转换与统一表达(前端编译)
解决的问题:算法工程师用 PyTorch 训练模型,生产环境却要部署到不同硬件。传统方式要为每个硬件重复开发推理脚本,且框架自带推理引擎很难跨硬件移植。用 Relay 作为统一中转格式,一次转换、多处部署。
具体案例:某公司算法组用 PyTorch 训练图像分类模型后,通过 TVM Relay 前端把模型翻译成 Relay IR,再编写一套优化流水线,同时部署到 NVIDIA T4 服务器和国产 NPU 加速卡,省去两套推理代码的维护成本。
动态 NLP 模型优化
阶段:图优化与编译(Relay pass)
解决的问题:BERT 等 NLP 模型的输入序列长度在推理时动态变化,静态计算图 IR 无法应对。Relay 的函数式控制流和类型系统支持动态形状,让优化器在序列长度未知时依然能完成算子融合与内存规划。
具体案例:某搜索团队用 TVM 编译 BERT 推理服务,借助 Relay 的动态形状能力,在 batch 和序列长度实时变化的情况下,融合后的 kernel 仍然保持稳定的低延迟,相比 TensorRT 的静态图方案更灵活。
自定义算子融合调优
阶段:性能调优(FuseOps)
解决的问题:LayerNorm 类算子由多个基础算子组合而成,逐算子执行会产生多次 kernel 启动和中间张量访存。通过用 Relay 表达算子组合并开启 FuseOps,可以自动融合成单个 kernel。
具体案例:某大模型推理团队将 RMSNorm 用 Relay 的算子组合表达,经过算子融合后端到端延迟降低约 30%,避免了手写 CUDA kernel 的高昂开发成本。
Python 演示
import numpy as npimport tvmfrom tvm import relay def build_simple_net(): # 输入变量声明, shape=(10,10), dtype=float32 x = relay.var("x", shape=(10, 10), dtype="float32") w = relay.var("w", shape=(10, 10), dtype="float32") b = relay.var("b", shape=(10,), dtype="float32") # 构建数据流: dense -> bias_add -> relu y = relay.nn.dense(x, w) y = relay.nn.bias_add(y, b) z = relay.nn.relu(y) # 包装成函数和模块 func = relay.Function([x, w, b], z) mod = tvm.IRModule.from_expr(func) return mod def main(): mod = build_simple_net() print("=== 原始 Relay IR ===") print(mod["main"]) # 执行算子融合优化 with tvm.transform.PassContext(opt_level=3): fused_mod = relay.transform.FuseOps(3)(mod) print("=== 融合后 Relay IR ===") print(fused_mod["main"]) # 降级到 TIR with tvm.transform.PassContext(opt_level=3): tir_mod = relay.transform.ToTIR()(fused_mod) print("=== 降级后模块类型 ===") print(type(tir_mod)) if __name__ == "__main__": main()
x = relay.var("x", shape=(10, 10), dtype="float32")声明一个名为 x 的输入张量变量,带形状和数据类型信息。y = relay.nn.dense(x, w)调用 dense 算子,把输入 x 与权重 w 做矩阵乘,生成中间结果 y。z = relay.nn.relu(y)对中间结果施加 ReLU 激活,得到最终输出。func = relay.Function([x, w, b], z)把输入变量和输出表达式包装成一个具名函数。tvm.IRModule.from_expr(func)将函数放入 IRModule,这是 Relay 程序的模块容器。relay.transform.FuseOps(3)(mod)以融合等级 3 执行算子融合 pass,合并可融合的算子以减少 kernel 启动。relay.transform.ToTIR()(fused_mod)把优化后的 Relay 模块降级为 TIR 模块,进入底层优化与代码生成阶段。
常见误区
误区我一直觉得 Relay IR 就是计算图,只不过把节点和边改成文本形式存储,换个序列化格式罢了,没什么本质区别。
正解Relay 是一门完整的函数式语言,能表达递归、高阶函数、let 作用域和控制流,而传统计算图只能表达无环数据流。Relay 程序的数据流视图确实等价于计算图,但控制流是图 IR 表达不了的。就像蓝图里可以标注“如果下雨就关门”,照片式的草图却画不出这种逻辑。
为何常见:入门教程多从“计算图优化”切入,读者看到 Relay 展开后的数据流形态与图相似,便误以为二者只是存储格式不同。
误区网络模型的输入输出不都是张量吗?反正类型只有浮点数几种,类型系统听着像个摆设,直接翻译成代码不就行了。
正解Relay 的类型系统能在编译期发现形状不匹配。比如矩阵乘要求内维相等、卷积要求通道对齐,这些错误如果留到运行时才暴露,排错成本极高。类型推导规则像“图纸审查”:结构工程师在开工前就发现承重墙位置不对,而不是等楼塌了再补救。
为何常见:Python 动态语言的开发体验让初学者习惯了“跑起来才知道错”,低估了静态检查在大规模编译中的价值。
误区FuseOps 就是把几个算子合并写进一个函数里,代码看起来紧凑了一些,但实际效果应该差不多吧?
正解融合的核心收益是减少 kernel 启动次数和中间张量的内存读写。GPU 上每次启动 kernel 都有固定开销,中间结果在显存里写一遍再读一遍更是消耗带宽。融合后多个操作在一个 kernel 内完成,访存次数大幅下降,这是量级层面的性能差异,而不是代码美观问题。
为何常见:融合前后打印出的 Relay 源码行数变化不大,初学者从文本层面看不出性能差异,忽略了“运行方式”才是关键。
误区TVM 和 Relay 都是给 GPU 做加速的,CPU 上部署模型直接用 PyTorch 就行,根本不需要 Relay 这套东西。
正解Relay 是硬件无关的高级 IR,CPU、GPU、NPU、FPGA 后端都可以消费 Relay 做相同的优化,再降级到各自的 TIR 代码生成路径。计算平台团队恰恰最需要 Relay 的抽象能力——一套优化逻辑适配所有硬件。
为何常见:大多数教程以 GPU 为例展示性能提升,让读者误以为 Relay 绑定了某个具体硬件。
面试题
Q1. TVM 中的 RELAY IR 属于哪一层?
- A.前端模型表达层
- B.高级中间表示层← 正确答案
- C.后端汇编代码层
- D.操作系统内核层
答案:B。Relay 是 TVM 的高级中间表示,位于前端模型与后端 TIR 之间。A 是框架本身做的事,C 是代码生成后的产物,D 与编译链路无关。用蓝图比喻来说,Relay 是设计蓝图而非施工图。
Q2. Relay 程序中一组具名函数的集合被组织成什么?
- A.模块容器 Module← 正确答案
- B.张量容器 Tensor
- C.算子容器 Operator
- D.优化器容器 Pass
答案:A。Module 是 Relay 程序的最外层容器,装载一组具名函数,相当于蓝图的图册。Tensor 是数据类型,Operator 是算子对象,Pass 是优化流程,它们都不是函数集合的容器。
Q3. let 绑定 x = e1 in e2 中,e1 的结果在哪里可以使用?
- A.只能在 e1 内部使用
- B.只能在 e2 内部使用← 正确答案
- C.全程序任意位置可用
- D.只能在函数末尾使用
答案:B。let 的语义是:先计算 e1 并把结果绑定到 x,然后 x 只在 e2 的作用域内可用。这类似蓝图中的“局部标注”,只在本层有效,不会泄漏到整个程序。
Q4. 下列哪些是 Relay 表达式语法支持的构造?
- A.常量张量构造← 正确答案
- B.匿名函数定义← 正确答案
- C.递归函数调用← 正确答案
- D.指针算术运算
答案:ABC。Relay 的表达式文法包含常量、变量、匿名函数、函数调用和 let 绑定,借助匿名函数与调用,可以构造递归逻辑。Relay 是高级函数式 IR,不提供指针及指针运算。
Q5. Tensor[(10, 10), float32] 中 float32 表示什么?
- A.张量的形状信息
- B.元素的数据类型← 正确答案
- C.张量的变量名字
- D.算子的融合等级
答案:B。Tensor 类型前半部分是形状 (10,10),后半部分 float32 是元素数据类型,两者共同构成张量类型标注。形状和数据类型分离,正是类型系统能够检查 shape 是否匹配的前提。
Q6. 类型推导规则 Γ[x↦T1] ⊢ e2 : T2 中,x 的新类型是什么?
- A.与 e1 类型无关
- B.就是 e1 的类型 T1← 正确答案
- C.继承 e2 的类型 T2
- D.由运行时动态确定
答案:B。环境扩展 x↦T1 表示把 x 绑定到 e1 的类型 T1,然后在扩展环境中检查 e2。这就像图纸审查时,构件 x 的类型由它的设计值 T1 决定,而非由后续使用方式决定。
Q7. 相比传统计算图 IR,Relay 的哪些能力让它能表达 BERT 这类动态模型?
- A.支持动态形状的输入← 正确答案
- B.支持条件分支与循环← 正确答案
- C.支持递归与共享子图← 正确答案
- D.支持直接操作寄存器
答案:ABC。Relay 的类型系统支持动态形状,函数式控制流支持条件分支和循环,匿名函数与递归支持共享子图与复杂结构。直接操作寄存器是低级 IR 的特征,与 Relay 的高级抽象无关。
Q8. Relay 优化中 FuseOps 把多个算子合并的主要收益是什么?
- A.仅仅减少代码行数
- B.减少核启动与访存← 正确答案
- C.让解释器运行更快
- D.增加冗余的算子数
答案:B。算子融合将多个计算合并为一个 kernel,减少 GPU 上的 kernel 启动次数,并避免中间张量写回显存再读回的开销。它改的是“施工工艺”,而不是图纸上的行数。
Q9. Relay 程序在 TVM 流水线中降级后的目标 IR 是什么?
- A.先翻译为 LLVM IR
- B.降级为目标 TIR← 正确答案
- C.直接生成 CUDA 源码
- D.编译成 Python 字节码
答案:B。Relay 降级后得到 TIR(张量中间表示),携带循环、内存等低级信息。LLVM IR 是 TIR 之后代码生成阶段可能再使用的后端 IR,CUDA 源码和 Python 字节码都不是 Relay 的直接降级目标。
Q10. 下列关于 Relay 类型系统的说法正确的有?
- A.编译期可发现形状错误← 正确答案
- B.支持多态函子机制← 正确答案
- C.类型信息推动优化← 正确答案
- D.运行时才确定类型
答案:ABC。Relay 在编译期通过类型推导检查形状是否匹配,支持多态函子让函数适配不同形状,类型信息还用于布局推导等优化。Relay 是静态类型语言,类型在编译期确定而非运行时。