想象你是这家米其林后厨的总策划。晚上八点,客人的订单像雪片一样飞来:有人点红烧肉,有人点清蒸鱼,还有人要两菜一汤。你作为主厨,不会直接冲进灶台炒菜——你要先看完整批订单,把能一起准备的食材归拢、把重复的工序合并、规划出一条“最高效的出菜路线”。这份规划,是图级思维。而真正到了灶台,掌勺师傅盯着的不是整桌宴席,而是眼前这口锅:葱花切多细、火候开多大、翻炒几下……每一步都是张量级别的执行细节。规划与执行,各有各的图纸。AI 编译器 TVM,正是用两层中间表示(IR)把这两张图纸分开管理:Relay IR 是主厨的菜单蓝图,TIR 是师傅的工序卡片。本文就用这个后厨比喻,带你看懂两层图纸如何分工、如何转换、如何各自优化。
TVM 的两层 IR 设计是理解 AI 编译器的主干知识:掌握 Relay 与 TIR 的分工,你就拿到了读懂编译流水线、定位性能瓶颈、应对 AI 编译器岗位面试的通关钥匙。
目录
可视化演示
展示编码器逐步处理输入序列并生成上下文向量,解码器基于上下文向量逐步生成多步预测结果的完整流程。点击模式按钮切换编码、解码或完整流程视图。
本节配有交互动效(进入视口后自动加载,约 3KB)。
编码器阶段 编码器逐步读取输入序列x₁到x₄,每步更新隐藏状态h 最终隐藏状态h₄作为上下文向量C传递给解码器 hₜ = f(hₜ₋₁, xₜ)
Relay IR 计算图:主厨的菜单蓝图
主厨拿到订单后,不会立刻冲到灶台,而是先画一张“统筹图”:哪些菜共用烤箱、哪些菜能同时蒸、哪些汤可以提前吊好。这张图只关心菜与菜之间的依赖和顺序,完全不关心刀工火候。对应到 AI 编译器,这张统筹图就是 Relay IR——它以计算图(data flow graph)描述整个神经网络,每一个节点是一个算子,每一条有向边是张量数据的流动方向。
Relay IR 是 TVM 的“高层中间表示”。它有几个关键设计:数据流节点把卷积、池化、全连接等算子组织成一张有向无环图,节点之间显式传递张量;算子抽象则用统一的形式表达不同算子及其属性(比如卷积的 stride、padding),使上层优化无需关心硬件实现;类型推断为每个节点标注张量的形状(shape)与数据类型(dtype),这是后续所有优化决策的依据。
用数学语言概括,一张 Relay 计算图可以表示为三元组:
主厨的统筹图只画“菜与菜的先后”,不画刀工火候。同理,Relay IR 也不关心循环怎么排、内存怎么放、寄存器怎么用,这些细节全部留给第二张图纸——TIR。
- 节点集合 $V$图中每个节点代表一个算子(如 conv2d、add),好比菜单上的一道菜,有明确的“做法属性”。
- 数据流边 $E$每条有向边表示张量在算子间的流动,好比食材从洗切台递到灶台,方向不能乱。
- 类型标注 $\lambda$为每个节点记录输出张量的 shape 与 dtype,好比菜谱上标注“大份、中份、小份”。
TIR 张量级表示:掌勺师傅的工序卡
灶台师傅手里的工序卡精确到每一步:热锅、下油、入葱姜、倒主料、大火翻炒三十秒。它完全不关心整桌菜的先后顺序,只关心眼前这一道菜怎么做完。TVM 中的 TIR(Tensor Intermediate Representation,张量中间表示) 就是这张工序卡,它以张量表达式、循环结构和存储缓冲来描述单个算子的具体计算过程。
以矩阵乘法为例,TIR 会把“计算 C = A × B”展开成三重嵌套循环:外层遍历输出矩阵的行和列,内层做累加。每个循环变量对应一个索引,循环体内声明了缓冲区和访存操作。与 Relay 只画“节点与边”不同,TIR 已经精确到“哪一步读哪个地址、哪一步写哪个寄存器”。
数学上,矩阵乘法可以写成对每个输出元素的内积求和:
师傅的工序卡写清了每一步的“手和脚”,TIR 同样把计算落实到循环、索引与缓冲上,让硬件真正“有活可干”。但工序卡不会自己从菜单蓝图中长出来——两者之间需要一位“翻译官”。
- 输出张量 $C[i][j]$工序的“出品”,每个位置的值都由累加得到,对应师傅最终盛出的那盘菜。
- 内层累加下标 $k$把 $A$ 的第 $i$ 行与 $B$ 的第 $j$ 列逐元素相乘再相加,对应工序中“反复翻炒让食材入味”。
- 双重遍历 $i$、$j$外层循环遍历输出的每一行每一列,对应“每一盘菜都要按同一道工序做一遍”。
两层转换流程:从菜单到工序的翻译
菜单蓝图不会自动变成工序卡,需要主厨把“红烧肉”翻译成“切三厘米方块、小火炖四十分钟”。TVM 中这个翻译动作叫作 Lowering(降级):从高层抽象的 Relay IR 逐层降低到贴近硬件的 TIR,最终到达机器码。翻译质量直接决定后厨效率——翻译得笨,师傅再勤快也白搭。
Lowering 的完整路径是:先把 Relay 计算图做类型推断,再执行 FuseOps(算子融合)——把 Conv2D、BN、ReLU 这类相邻算子合并成一个 fusion 节点,让中间结果不落回全局内存;随后每个 fusion 节点被翻译成 TIR 的循环与存储结构;最后经过 schedule(调度)和 codegen(代码生成)产出目标机器码。
整体流水线可以浓缩成一条“翻译链”:
翻译链的两端各有一批“优化顾问”:Relay 侧优化菜单结构,TIR 侧优化工序细节。接下来分别看看这两层各自握着哪些优化武器。
- 左侧 Relay 计算图整张菜单蓝图,节点是算子、边是数据流。
- 中间 Lowering(FuseOps 融合算子)主厨决定“哪几道菜可以一锅出”,减少中间环节。
- 右侧 TIR 与代码生成工序卡翻译完成后,由机器照着执行。
Relay 图级优化:优化出菜顺序
主厨可以不改任何菜谱,只调整“统筹图”就显著提升翻台率:高汤提前吊好、重复的配菜一次性切好、能一锅出的菜绝不分开炒。Relay 图级优化的哲学正是如此——不改变每个算子的计算方式,只改变图的结构,让整张图更省算力、更省内存、更少中间张量读写。
最常见的三类图级优化是:常量折叠(FoldConstant),把全常量子表达式在编译期就算完,运行时直接取结果;代数化简(AlgebraicSimplify),把 add(x, 0)、mul(x, 1) 这类恒等计算直接删除;算子融合(FuseOps),把 Conv2D+BN+ReLU 合并成一个节点,中间张量不再落回主存。
常量折叠是最直观的例子。假设图中有一个乘法节点 mult(2, 3),它的两个输入都是编译期常量:
图级优化像“排兵布阵”,不碰任何一口锅。但同样的菜,师傅手法不同,速度也能差好几倍——这就是 TIR 张量级优化要解决的问题。
- 两个操作数都是常量编译期就能确定乘法结果,不需要等运行时。
- 用常量 6 替换原子表达式图中不再出现 mult 节点,等价于主厨提前把高汤熬好,出菜时直接舀。
- 运行时只剩一次加法少一次乘法与一次中间张量访存,整体延迟下降。
TIR 张量级优化:精进翻炒手法
工序卡已经写明“大火翻炒三十秒”,但老师傅还能再精进:先翻炒再颠勺、把配菜提前摆好队、一次炒两锅的量。同一道菜,不同手法,出品速度差好几倍。TIR 级优化就是这套“手法进化”:不改变计算结果,只改变循环与存储的安排,让硬件资源(缓存、SIMD 向量单元、多核)被喂得更饱。
典型手段有三类:循环变换(循环交换、分块、展开),目标是提高缓存命中率与指令级并行;向量化(vectorize),把多个标量操作打包成一条 SIMD 指令一次算完;存储布局重排,把 NCHW 改成 NHWC 或其他布局,让访存地址更连续。
以循环交换为例,原始代码按“外层 i、内层 j”的顺序访问矩阵:
TIR 优化把每一道工序打磨到极致,但单点最优不等于整体最优。一家餐厅能稳定出菜,靠的是从接单到装盘的整套流程——这正是 TVM 统一编译管线做的事。
- 原始嵌套顺序外层 $i$、内层 $j$:内层循环变化最快,对应“一行一行地切配菜”。
- 交换后外层 $j$、内层 $i$:访问顺序改成按列推进,适配列优先存储或分块访存。
- 复杂度不变、常数因子下降总操作数仍是 $O(MN)$,但缓存命中率提升,实际耗时明显缩短。
统一编译管线:一整天的后厨运转
一家米其林餐厅能稳稳出菜,靠的不是某位大师傅灵光一现,而是整套流水线:接单→规划→翻译→炒制→装盘。TVM 的编译管线也是如此,它是 Relay 与 TIR 两层级 IR 协同工作的完整舞台。从前端框架导入模型开始,到最终产出可执行库,每一步都有明确分工。
完整管线包含五个环节:前端解析用 from_pytorch、from_onnx 等接口把 PyTorch / ONNX 模型翻译成 Relay IR;图级优化在 Relay 上做常量折叠、算子融合等变换;Lowering 把优化后的图翻译成 TIR;张量级优化与调度在 TIR 上做循环变换、向量化;后端代码生成基于 LLVM、CUDA 或 NPU 指令集产出目标二进制。现代 TVM 还内置 AutoTVM 与 Ansor 自动搜索最优调度,替代手工调参。
整条流水线用一张图可以浓缩为:
从接单到装盘,两层 IR 各司其职又紧密咬合。理解了这条流水线,你就理解了 TVM 的骨架——下面看看这套骨架在实际工作中能解决哪些真问题。
- Model(框架模型)客人留下的原始订单,PyTorch / ONNX / TensorFlow 格式五花八门。
- 前端解析为 Relay把订单翻译成主厨能看懂的菜单蓝图,统一成一张计算图。
- 图优化与 Lowering 到 TIR先优化蓝图,再把每个算子翻译成工序卡(循环与存储)。
- Schedule 与代码生成工序卡交给师傅(硬件)执行,AutoTVM / Ansor 负责找到最优手法。
AI 应用场景
场景一:边缘设备部署 CV 模型
阶段:模型编译与部署
解决的问题:原始 PyTorch 模型直接跑在手机或嵌入式设备上,速度慢且内存爆炸。传统做法是手工改写成 C++ 并逐算子调优,工作量巨大。通过 TVM 前端解析为 Relay,图级优化(算子融合)减少中间张量,再 Lowering 到 TIR 结合目标架构生成高效代码,一次编译多端部署。
具体案例:将 YOLOv5s 部署到 Jetson Nano。原始 PyTorch 模型 1080p 推理约 800ms,经 Relay 图优化(融合 Conv+BN+ReLU)与 TIR 自动调度后降至约 120ms,内存占用下降约 40%。
场景二:自研推理引擎的融合算子优化
阶段:算子融合与图优化
解决的问题:线上推荐模型中 Conv+BN+ReLU 三段算子反复读写全局内存,访存成为瓶颈。传统优化需要手工在 CUDA 层写融合 kernel,周期长。通过 Relay 的 FuseOps 先在图级决定“谁和谁融合”,再在 TIR 层自动生成融合后的循环实现。
具体案例:某 CTR 预估模型部署到 CPU 推理服务,模型精度不变,单请求时延 p99 从 23ms 降至 14ms,吞吐提升约 60%。
场景三:大模型推理的访存与显存调优
阶段:张量级优化与自动调度
解决的问题:大模型参数量巨大,访存带宽成为首要瓶颈。传统手工调优难以覆盖所有算子。TIR 层的循环分块、向量化、存储布局重排等手段能显著提高带宽利用率,AutoTVM 还能自动搜索最优分块参数。
具体案例:LLaMA-7B 在 A100 上用 TVM Ansor 自动调度,batch=1 推理时 prefill 阶段吞吐相比默认配置提升约 1.6 倍。
Python 演示
"""TVM 双层 IR 演示:构建 Relay 计算图并 Lowering 至 TIR"""import tvmfrom tvm import relay def build_conv_relu_graph(): # 1. 构建 Relay 计算图:Conv2D -> ReLU data = relay.var("data", relay.TensorType((1, 3, 32, 32), "float32")) weight = relay.var("weight", relay.TensorType((8, 3, 3, 3), "float32")) conv = relay.nn.conv2d(data, weight, padding=(1, 1)) out = relay.nn.relu(conv) func = relay.Function([data, weight], out) return tvm.IRModule.from_expr(func) def main(): mod = build_conv_relu_graph() # 2. 图级优化:常量折叠 + 公共子表达式消除 seq = tvm.transform.Sequential( [relay.transform.FoldConstant(), relay.transform.EliminateCommonSubexpr()], opt_level=2) mod_opt = seq(mod) print("=== Relay 图级 IR ===") print(mod_opt["main"]) # 3. Lowering 前置:类型推断 + FuseOps 算子融合 with tvm.transform.PassContext(opt_level=2): mod_fused = relay.transform.InferType()(mod_opt) mod_fused = relay.transform.FuseOps()(mod_fused) print("\n=== 算子融合后 Relay(Conv+ReLU 合并为 fusion 节点)===") print(mod_fused["main"]) # 4. 代码生成:融合后的图 Lowering 到 TIR 并编译 target = tvm.target.Target("llvm") with tvm.target.Target(target): lib = relay.build(mod_fused, target=target) source = lib.get_lib().get_source() print("\n=== TIR 张量级代码生成结果 ===") print("LLVM 源码行数:", len(source.splitlines())) print("编译通过:Relay 规划图 + TIR 执行算,双层分工完成。") if __name__ == "__main__": main()
relay.var("data", relay.TensorType((1, 3, 32, 32), "float32"))声明 Relay 图输入节点,相当于菜单上标注“这道菜用哪种食材”。relay.nn.conv2d(data, weight, padding=(1, 1))创建卷积算子节点,图级表示不关心底层循环怎么实现。tvm.transform.Sequential([...], opt_level=2)按顺序执行图级优化 pass,相当于主厨按既定流程优化上菜顺序。relay.transform.FoldConstant()常量折叠优化,编译期算完常量表达式,运行时少干活。relay.transform.EliminateCommonSubexpr()消除公共子表达式,避免重复计算同一结果。relay.transform.InferType()(mod_opt)类型推断,为每个节点补全 shape 与 dtype,是融合与生成的依据。relay.transform.FuseOps()(mod_fused)算子融合,把 Conv+ReLU 合并成一个 fusion 节点,中间张量不落回内存。relay.build(mod_fused, target=target)编译驱动,串联 Lowering、调度与代码生成,输出可执行库。
常见误区
误区TVM 既然有两种 IR,那我只需要选其中一种来用就行了,比如只用 Relay 做优化就够。
正解两层 IR 不是竞争关系,而是上下级流水线:Relay 负责计算图级规划(主厨的菜单),TIR 负责张量级执行(师傅的工序)。缺了 Relay,模型无法做高层结构优化;缺了 TIR,计算无法落到硬件。二者必须按顺序配合使用。
为何常见:TVM 文档常分开介绍两级 IR,先学 Relay 再学 TIR,容易误以为它们各自独立、可互相替代,就像只看过菜单和只看过工序卡的人都以为厨房只需要自己那一张纸。
误区FuseOps 算子融合最终把循环合并了,所以融合优化肯定发生在 TIR 张量级,跟 Relay 没关系。
正解融合的“决策”发生在 Relay 图级——FuseOps pass 决定哪些算子合并成 fusion 节点;TIR 只是按决策去生成合并后的循环实现。决策与执行是两回事,好比主厨决定“这两道菜一锅出”在前,师傅照着合并工序在后。
为何常见:融合的效果表现为循环合并、中间缓冲消失,初学者看到循环就联想到 TIR,于是误把图级决策归到了执行层。
误区Relay 做图优化,TIR 做循环优化,两边各干各的,互不影响,分别调优就行。
正解两层优化深度耦合。Relay 的算子融合会直接改变 TIR 的循环结构(融合后的循环体更大、中间缓冲更少);TIR 的调度决策(如分块大小)也会反过来影响图级的内存规划。实际调优时经常要在两层之间来回调整。
为何常见:教材按层分开讲,形成了“井水不犯河水”的错觉,就像以为菜单规划与炒菜手法毫无关系,实际上“一锅出”的决策直接改变了灶台操作。
误区Relay 里的类型推断就是检查张量形状有没有写错,对性能没有任何影响,跑通就行了。
正解类型推断结果(shape 与 dtype)直接决定算子选择——比如能否走 int8 量化内核、卷积算法选哪种,还影响 FuseOps 能否融合以及 Lowering 后 TIR 的循环边界。它是一切后续决策的地基。
为何常见:“类型推断”名字听起来像静态检查,容易被低估;实际上它像菜谱上的分量标注,写错了整桌菜都会乱套。
面试题
Q1. TVM 中的两层中间表示是指哪两个?
- A.Relay IR 与 TIR← 正确答案
- B.ONNX 与 TFLite
- C.LLVM IR 与 MLIR
- D.CUDA 与 OpenCL
答案:A。Relay IR 与 TIR 分别承担计算图级和张量级表示,是 TVM 特有的两层 IR。ONNX/TFLite 是模型格式,LLVM IR 是后端通用 IR,CUDA/OpenCL 是编程模型,都不属于 TVM 的两层中间表示。
Q2. Relay IR 主要负责哪类优化?
- A.侧重计算图级优化← 正确答案
- B.侧重循环级优化
- C.侧重指令级优化
- D.侧重寄存器分配
答案:A。Relay 关注计算图中的算子、数据流与类型信息,优化目标是减少冗余算子和中间访存。循环、指令、寄存器层面的优化属于 TIR 及后端代码生成的职责范围。
Q3. 将 Relay 计算图翻译成 TIR 循环结构的过程叫什么?
- A.类型推断与标注
- B.Lowering 降级翻译← 正确答案
- C.图分割与切分
- D.常量折叠预计算
答案:B。Lowering(降级)把高层 Relay 计算图逐算子翻译成 TIR 的循环与存储结构,相当于把菜单蓝图翻译成工序卡。类型推断是前置步骤,图切分是分布式部署手段,常量折叠是图级优化之一。
Q4. FuseOps 算子融合发生在哪一层?
- A.Relay 图级层面← 正确答案
- B.TIR 张量级层面
- C.后端机器码层面
- D.前端框架解析层面
答案:A。FuseOps 是 Relay 图级优化 pass,它在图结构上决定哪些相邻算子合并为 fusion 节点,从而减少中间张量读写。TIR 只是负责把融合后的节点实现为合并的循环,融合决策本身发生在 Relay 层。
Q5. 以下哪些属于 Relay 图级支持的优化?
- A.常量折叠与算子融合← 正确答案
- B.代数化简与冗余消除← 正确答案
- C.循环向量化与分块
- D.公共子表达式消除← 正确答案
答案:ABD。常量折叠、代数化简、公共子表达式消除都发生在 Relay 图级,目标是优化图结构。循环向量化与分块属于 TIR 张量级优化,它不改变图结构,只改变循环与存储安排。
Q6. TIR 的核心表示形式是什么?
- A.张量表达式与循环结构← 正确答案
- B.数据流图与算子节点
- C.神经网络逐层结构
- D.编译正则与文法规则
答案:A。TIR 以张量表达式描述计算语义,以循环结构刻画执行过程,以存储缓冲映射数据位置。数据流图与算子节点是 Relay 的表示形式,神经网络分层结构是框架模型的概念,正则文法与编译原理无关。
Q7. 循环向量化优化属于哪一级的职责?
- A.TIR 张量级优化职责← 正确答案
- B.Relay 图级优化职责
- C.前端模型解析职责
- D.类型推断与标注职责
答案:A。循环向量化把多个标量操作打包成 SIMD 指令,属于 TIR 张量级优化的手段。Relay 只做图级结构优化,前端负责模型格式转换,类型推断为优化提供 shape/dtype 依据,都不直接处理循环。
Q8. 关于两层 IR 的关系,下列说法正确的有哪些?
- A.Relay 负责图级数据流规划← 正确答案
- B.TIR 负责硬件循环与存储← 正确答案
- C.两层 IR 彼此独立可省略
- D.Lowering 桥接两层表达← 正确答案
答案:ABD。Relay 偏重“做什么”(数据流与图结构),TIR 偏重“怎么做”(循环与存储),Lowering 是连接两层的翻译桥,三者协同构成完整编译管线。两层 IR 是上下游关系,缺一不可,不能省略任何一层。
Q9. 常量折叠优化为什么能加速推理?
- A.把常量计算挪到编译期完成← 正确答案
- B.增加网络层数与参数规模
- C.将模型精度降到 INT8 表示
- D.减少训练时反向传播次数
答案:A。常量折叠在编译期把全常量子表达式算完并用结果替换,运行时少做计算和访存,类似高汤提前熬好。它不改变网络结构、不降低精度、与训练过程无关,后三者是其他优化手段或错误描述。
Q10. 以下哪些属于 TIR 级循环与存储优化?
- A.循环分块与并行化处理← 正确答案
- B.向量化访存指令利用← 正确答案
- C.常量折叠与冗余消除
- D.存储布局与缓冲重排← 正确答案
答案:ABD。循环分块、并行化、向量化、存储布局重排都作用于 TIR 的循环与存储结构,提高缓存命中率与计算并行度。常量折叠与冗余消除发生在 Relay 图级,优化对象是图结构而非循环执行。