如果你最近开始关注 Apache TVM大概率会频繁撞见一个名字Relax。TVM 在很长一段时间里都以 Relay 作为上层图 IR但 Relay 在动态形状、控制流、表达抽象和可组合性上越来越吃力社区这几年的重心已经明显转向了 Relax。简单来说Relax 是一套重新设计的、可组合的中间表示IR它不只是给 Relay 换个皮而是把 IR 的构建方式、类型系统和执行方式整体换了一套思路。这篇博客就是写给那些想搞清楚「创建 Relax」到底从哪下手的同学我们先说清 Relax 是干什么的再给三条可落地的创建路径脚本 DSL、BlockBuilder 编程式构建、框架模型导入最后用 VirtualMachine 把模块真正跑起来并整理我自己趟过的坑。1. 先弄清楚 Relax 是什么再谈创建很多人一上来就搜「Relax 教程」但忽略了最关键的一点Relax 不是某个新的算子库也不是一个独立的推理框架。它是 TVM 编译器栈里的 IR 层承上启下。承上接收前端框架模型启下把计算图降低到 TIR张量中间表示和具体代码生成。1.1 从 Relay 到 Relax我为什么觉得换跑道是必然用生活里的场景打个比方Relay 像一张城市地铁图站点算子和线路数据流都画得很清楚适合做整图优化。但一旦遇到临时封路动态 shape、分时段运营控制流这种需求改起来就很痛苦。Relay 的问题主要有三个对动态 shape 支持吃力。Relay 的 shape 推断在静态 shape 场景下很爽但碰到batch_sizeNone、NLP 里那种变长序列整个优化流程会变得非常敏感。图和底层算子之间的边界太僵硬。自动微分、循环、条件分支这些高层控制流在 Relay 里要么表达困难要么只能塞进自定义算子绕了一圈又回到手动写算子。编译流程的可组合性不足。现代编译器的趋势是把 IR 拆成小 pass、小变换让用户可以灵活编排。Relay 设计得过于重做一轮完整的图优化往往牵连太多。Relax 的思路是把“计算图”和“执行逻辑”分开。计算图用 dataflow block 表达执行逻辑用 call、seq、if、while 这类表达式组合。这样一来动态 shape、控制流、算子融合、内存规划都能在一个统一 IR 里做而不是各自为政。1.2 Relax 的核心设计目标落地到三个关键词如果说 Relay 的设计目标是“把图优化做完整”那 Relax 的设计目标就是“让 IR 可表达一切执行策略”。我理解下来主要围绕三个关键词可表达性。Relax 可以表达 dataflow block 里的纯计算子图也可以表达 dataflow block 之外的赋值、条件、循环、闭包甚至嵌套的 call_tir。这让自动微分、量化、重排内存这类需要控制流的变换有了一个标准落地位置。可组合性。整个 IR 构建用的是 Python 算子你可以把一段 IR 当作一个对象传给 pass再得到一段新 IR。脚本 DSL 和编程式 BlockBuilder 是同一套底层区别只是你动不动的灵活性。效率。Relax 的 call_tir 机制可以直接调用 TIR 内核底层还是 TVM 的 CodeGen所以跑起来的速度不输 Relay局部甚至更好尤其是动态场景下的调度。下面这个表可以帮你快速理解 Relay 和 Relax 的差异对比维度RelayRelax图 IR 定位重图 IR自带大量优化规则轻量可执行 IR执行逻辑与计算图分离控制流表达较弱需要特殊算子原生支持 if/while/call 组合动态 shape支持但优化易碎类型系统更灵活shape 可作为 symbolic var自动微分实现复杂可逐步、可组合地插入微分步骤算子融合基于图规则匹配基于 dataflow block 与 pass 链粒度更细构建方式PyTorch 风格模块接口脚本 DSL BlockBuilder 前端导入器我不建议把 Relay 一棍子打死很多生产环境还在用它而且短期内也不会消失。但如果你要写新项目、新编译器模块或者想深度定制模型执行流程现在上车 Relax 是比较清醒的选择。1.3 这篇教程要带你做什么后续内容我会分三条路径创建 Relax 模块用tvm.script的 Relax 脚本 DSL像写伪代码一样写 IR。用relax.BlockBuilder编程式构建 IRModule。用relax.frontend把 ONNX / PyTorch 模型导入成 Relax IRModule。三条路径最终都会进入同一个环节编译成可执行模块用relax.VirtualMachine跑起来。整个过程我会尽量给可复现代码同时把每一步“为什么这么做”讲清楚因为只看代码你不知道踩坑点在哪看懂了原理才能稳定复现。另外先提醒一句TVM 是个迭代非常快的项目Relax 又在快速演进中API 在不同版本之间可能会有差异。我的示例基于当前主流 release 分支的常见写法如果你是源码构建的最新版个别接口名要以实际情况为准。2. 环境准备灶没烧好菜别下锅不管你是想写一个R.add的玩具程序还是导入整个 ONNX 模型都得先把 TVM 环境弄利索。Relax 是 TVM 的一部分不是独立安装包所以别去搜“install relax”核心是装对 TVM。2.1 安装方式怎么选pip 还是源码编译如果只是学 API、跑小模型直接 pip 安装 nightly wheel 就行pip install apache-tvm -i https://pypi.org/simple但你要注意apache-tvm这个包通常是 CPU 版本而且不一定带上最新的 Relax 前端特性。TVM 社区官方也会提供 nightly tar 包地址在官方发布页上需要自己选择。个人建议如果你是做长期项目或做编译器开发最好从源码编译因为你能拿到最新的 pass 和算子实现出问题也方便用调试器跟进。源码编译的常规步骤是git clone --recursive https://github.com/apache/tvm tvm cd tvm mkdir build cp cmake/config.cmake build/ # 编辑 build/config.cmake按需打开 CUDA / LLVM / ROCm cd build cmake .. make -j$(nproc)config.cmake 里最需要关注的是USE_LLVM。Relax 最终要落到 CPU/GPU 代码生成没有 LLVM 你连 CPU 后端都用不了。改成set(USE_LLVM ON)或者指定llvm-config路径编译完检查tvm.target.Target(llvm)能不能正常创建。2.2 冒烟测试先确认 relax 真的可用很多人装了 TVM 之后直接跑import tvm然后开始写 Rela x结果报错AttributeError: module tvm has no attribute relax。这多半是版本不对。检查方式很简单python -c import tvm; print(tvm.__version__) python -c from tvm import relax; print(relax.__name__)只要能打印出版本号且relax模块不报错环境基本没问题。接着跑一个最小冒烟测试创建并构建一个空函数import tvm from tvm import relax as rx bb rx.BlockBuilder() with bb.function(dummy, []): bb.emit_func_output(rx.const(1)) mod bb.get() print(mod.script())如果这步能打印一段含R.function的 IR说明 Relax 的构建链路是通的可以继续往下走。2.3 准备 target 和测试数据所有后续编译都要指定 target。CPU 平台直接写target tvm.target.Target(llvm -mcpuskylake)GPU 平台要写cuda或cuda -archsm_80这种。小教程里建议先从 CPU 开始等整套流程跑通了再换 GPU否则报错时很难判断是 IR 问题还是设备驱动问题。测试数据我用 NumPy 造import numpy as np x_np np.random.rand(4, 4).astype(float32) w_np np.random.rand(4, 4).astype(float32)这里反复强调dtype是float32因为 TVM 对 dtype 非常敏感你用 float64 的 NumPy 数据喂给 float32 IR运行阶段大概率会报 dtype mismatch。3. 创建 Relax 模块一用 tvm.script 写 IR 脚本tvm.script是 TVM 提供的一套基于 Python 语法的 DSL让你可以直接用R.function、R.Tensor、R.add这类标记来写 IR。它的目标很朴素让你写 IR 的时候感觉像在写 Python而不是操作一堆 AST 结点。3.1 先认识 R 门面长什么样Relax 脚本里的“R”命名空间是tvm.script.relax的简化名。你最常碰到的几个标记R.function定义一个 Relax 函数相当于 IRModule 中的一个函数。R.Tensor(shape, dtype)声明张量类型/结构信息。R.add、R.multiply、R.matmul对应的算子调用。R.dataflow声明一个 dataflow block块内是纯计算子图。R.output把 dataflow block 内部变量标记为输出。R.call_tir调用 TIR 内核函数是 Relax 连接底层 TIR 的通道。你不需要背这些IDE 补全 官方示例多写几遍就熟悉了。3.2 一个可运行的最小示例从加法开始新建一个 Python 文件tvm_relax_basic.py写入import tvm from tvm import relax as rx from tvm.script import relax as R R.function def add_matmul( x: R.Tensor((4, 4), float32), y: R.Tensor((4, 4), float32), ) - R.Tensor((4, 4), float32): return R.add(x, y) mod tvm.IRModule({main: add_matmul}) print(mod.script())tvm.IRModule的创建方式有两种一种是直接传函数字典像上面这样另一种是rx.IRModule。都行。打印出来的mod.script()就是标准化之后的 IR。如果看到 IR 里只有add一个操作说明解析成功。这个例子的重点是让你理解函数签名与类型的绑定R.Tensor((4, 4), float32)既规定了输入 shape也是后续 shape 推导的依据。你传 (4, 4) 的float32数据进入函数就返回 (4, 4) 的float32张量。3.3 为什么不建议直接写大表达式而是塞进 dataflow block实际模型往往不是一行算子而是一连串计算。Relax 提供了 dataflow block 来组织这一连串计算。格式是R.function def compute( x: R.Tensor((4, 4), float32), y: R.Tensor((4, 4), float32), ) - R.Tensor((4, 4), float32): with R.dataflow(): a R.add(x, y) b R.multiply(a, y) R.output(b) return bdataflow block 里面的变量是不可跨块引用的这是设计如此只有通过R.output标记为输出的变量才能被块外使用。为什么会有这种限制因为编译器需要知道哪些计算是 DAG 内部的“临时量”哪些是真正要保留的“边界结果”。在纯计算子图内部你可以放心做算子融合、内存复用而不会影响外部执行流。从脚本编写角度dataflow block 让你的代码看起来更接近计算图可读性也更好。后面 Python API 构建出来的 IR跟这种脚本写法是一模一样的所以建议一开始就养成用 dataflow block 的习惯。3.4 从加法升级到矩阵乘法贴近真实的算子形态单算子的示例看不出 Relax 的威力我们写一个稍微像样的模式z x * w然后是result z v。这在 MLP/注意力层里是非常常见的组合逐元素乘再加矩阵乘。R.function def mlp_like( x: R.Tensor((4, 4), float32), w: R.Tensor((4, 4), float32), v: R.Tensor((4, 4), float32), ) - R.Tensor((4, 4), float32): with R.dataflow(): z R.multiply(x, w) result R.matmul(z, v) R.output(result) return result这里你能看到 shape 推导的作用x和w都是 (4, 4)所以z是 (4, 4)z和v做矩阵乘法还是 (4, 4)。如果我把w改成 (4, 8)v改成 (8, 4)result依然是 (4, 4)。Relax 在解析阶段会自动推断这些中间变量的 shape不需要你逐一写注释。不过注意Relax 会推断不代表它不做检查。如果你写出R.multiply(x, v)其中v是 (8, 4)就会因为广播规则在 shape 推断阶段直接失败。所以脚本里写 shape 时要和实际算子语义对齐这是新手最容易翻车的地方。4. 创建 Relax 模块二用 BlockBuilder 编程式构建脚本 DSL 适合人读、人写但如果你想在程序里动态生成 IR比如循环生成不同 batch size 的多个函数脚本就不够灵活了。这个时候要请出relax.BlockBuilder。4.1 BlockBuilder 是干什么的理解它就是理解 pass你可以把 BlockBuilder 想象成一个“IR 脚手架”。它维护内部状态提供一系列方法往当前 IRModule 里添加函数、添加算子、设置输出。所有高级 pass 处理 Relax IR 时底层本质上就是在操作这些结构。所以理解 BlockBuilder不只是为了手写 IR更是为了以后写自定义 pass 时知道每个节点怎么被构造出来的。我习惯把它和脚本 DSL 对比脚本 DSL 是“声明式”BlockBuilder 是“命令式”。声明式写起来快命令式控制力强。二者生成的是同一类 IRModule选择哪条路径取决于场景。4.2 最小 BlockBuilder 实操一步步喂出 IRModule用 BlockBuilder 实现上面那个矩阵乘法函数import tvm from tvm import relax as rx bb rx.BlockBuilder() x rx.Var(x, rx.TensorStructInfo((4, 4), float32)) w rx.Var(w, rx.TensorStructInfo((4, 4), float32)) v rx.Var(v, rx.TensorStructInfo((4, 4), float32)) with bb.function(mlp_like, [x, w, v]): with bb.dataflow(): z bb.emit(rx.op.multiply(x, w)) result bb.emit(rx.op.matmul(z, v)) bb.emit_output(result) bb.emit_func_output(result) mod bb.get() print(mod.script())这段代码做了什么我拆开讲rx.Var创建占位变量x、w、v它们的类型由rx.TensorStructInfo指定。with bb.function(mlp_like, [x, w, v])打开一个新函数的构造上下文后续 emit 的内容都会进入这个函数。with bb.dataflow()进入 dataflow 块上下文块内的 emit 操作被标记为 dataflow 变量。bb.emit(rx.op.multiply(...))生成一个乘法节点并把它加入当前函数体返回这个结果的变量引用。bb.emit_output(result)把result标记为 dataflow block 的输出。bb.emit_func_output(result)把result设置为函数返回值同时自动关闭函数作用域。bb.get()返回完整的 IRModule。打印出来你会发现它和tvm.script写出来的结果如出一辙。所以我不建议你非此即彼地二选一而是理解它们本质相同脚本只是 BlockBuilder 的语法糖BlockBuilder 则是所有 pass 和脚本的底层执行器。4.3 动态生成多个变体一个我经常用到的真实场景有一次我给几个不同batch_size的模型做压测不可能每个 shape 手写一个函数于是直接用 BlockBuilder 循环生成传入不同的 batch size生成同名函数的不同变体放进同一个 IRModule。伪代码如下def make_module(batch_size, dtypefloat32): bb rx.BlockBuilder() x rx.Var(x, rx.TensorStructInfo((batch_size, 4), dtype)) w rx.Var(w, rx.TensorStructInfo((4, 4), dtype)) with bb.function(main, [x, w]): with bb.dataflow(): out bb.emit(rx.op.matmul(x, w)) bb.emit_output(out) bb.emit_func_output(out) return bb.get() for bs in [1, 4, 8, 16, 32]: mod make_module(bs) # 存储 mod后续 build 成动态库这个例子的价值在于当你需要批量实验多种 shape、多种算子组合甚至自动搜索某条 pass 链时BlockBuilder 编程式创建 IR 的威力就体现出来了。手动写脚本 DSL 做不到这种组合爆炸性扩展。再强调一个点BlockBuilder 创建的函数默认都在一个 IRModulemod里你还可以通过mod[main]、mod[mlp_like]访问不同函数这方便你写多函数模块也方便后续按函数粒度做 pass 分析。5. 创建 Relax 模块三从框架模型导入学完手工构建我更推荐你试试前端导入器。真实世界里没人会手写一个 ResNet-50 的计算图都是把 PyTorch、ONNX、TensorFlow 的模型导进来。Relax 的前端导入器就是把框架模型转成 Relax IRModule 的主要入口。5.1 Relax frontend 生态导入器到底做了什么Relax 的frontend包提供了一套按框架划分的导入函数比如from_onnx、from_pytorch、from_tensorflow等。它们的共同任务是把框架的计算图转成 Relax 的 op 和调用图但不负责训练也不负责推理只是一个“翻译器”。翻译过程中框架特有的算子会尽量映射到 TVM 的算子集合里。我给你的建议是先跑通一个 ONNX 导入再用 PyTorch 导入因为 ONNX 的边界更清晰报错更容易定位。PyTorch 导入涉及 trace 或 FX 转换环节更多。5.2 从 ONNX 导入一个模型并查看 IR假设你有一个model.onnx文件。导入代码大致如下import onnx from tvm import relax onnx_model onnx.load(model.onnx) shape_dict {input: (1, 3, 224, 224)} mod relax.frontend.from_onnx(onnx_model, shape_dict) print(mod.script())核心点在于shape_dict。ONNX 文件本身不强制指定输入 shape是动态的所以导入器需要你提供一个参考 shape 来完成图分析和形状推导。不同版本里from_onnx的签名可能有差异少数版本要求传shape_dict为有序字典OrderedDict你可以注意下。导入后打印mod.script()会看到一整段R.function你会发现里面的算子大多被展开成R.call_tir形式。这是因为 Relax 和 TIR 之间有一个绑定关系图级算子最终会对应到具体的 TIR 内核call_tir就是“调用这个内核”的指令。5.3 从 PyTorch 导入的坑trace 和 torch.fx 两种路线PyTorch 导入我单独拎出来因为坑最多。老的实现方式是走torch.jit.trace新版本则会走torch.fx符号跟踪。如果你用的 TVM 版本不一样可能遇见的入口也不一样import torch from tvm import relax # 以 torch.jit.trace 路线为例 scripted_model torch.jit.trace(model, example_input).eval() mod relax.frontend.from_pytorch(scripted_model, {input: example_input.shape})另一个路线mod relax.frontend.from_pytorch(model, {input: example_input.shape})有些版本会自动帮你内部处理 trace有些则要求传入torch.fx.GraphModule。我的经验是如果你的模型里有大量动态控制流、依赖外部 NumPy 数据或者非常规 Python 对象trace 可能不完整相反如果模型是标准 CNN/Transformer两种路线基本都能过。导入器跑通之后不要急着欢呼赶紧打印mod.script()确认关键算子有没有被正确映射。比如你期望看到nn.conv2d结果却出现一个call_tir(tvmgen_default_fused_conv2d...)这其实是正常的因为融合 pass 已经把它包进 TIR 内核了。若看到某些算子没有被识别日志里会有not supported提示这个时候就要先去算子映射表里查看能不能手动替换。6. 编译、运行与验证让 Relax 真正跑起来前几步我们一直在生产 IRModule。IRModule 只是“中间表示”不能直接执行。要让 Relax 模块真正出结果需要走一步编译然后在运行时里加载执行。6.1 用 relax.build 把 IRModule 变成可执行模块传统 TVM 里做图编译是relay.build到 Relax 这里入口是relax.build。基本用法import tvm import numpy as np from tvm import relax target tvm.target.Target(llvm -mcpuskylake) ex relax.build(mod, target)如果走兼容层有些旧版本是relax.vm.build参数基本一致。我在实践里用过两种写法统一建议是先print(hasattr(relax, build))看一眼再决定用哪个函数避免瞎猜。这一步其实做了很多事对 IRModule 跑一遍 pass 流水线包括算子融合、布局转换、内存规划。把 TIR 内核编译成目标机器的机器码。生成虚拟机可用的字节码和元数据。所以relax.build的耗时对你没有直观感知它的产物是一个executable包含了编译后的模块信息。6.2 用 VirtualMachine 执行模块并拿到结果Relax 模块编译完成后我习惯用relax.VirtualMachine来加载和运行dev tvm.cpu() vm relax.VirtualMachine(ex, dev) x_tvm tvm.nd.array(x_np, devicedev) w_tvm tvm.nd.array(w_np, devicedev) out vm[main](x_tvm, w_tvm, v_tvm) # 按函数名调用 print(out.numpy())注意我把 NumPy 数据显式转成了tvm.nd.array并指定了设备。这一步经常被忽略但 TVM 对设备一致性检查很严格你要是拿着 CPU 的 ndarray 去喂给 GPU 上运行的 VM会直接报错。如果当时我们构建的是mlp_like函数调用名就是vm[mlp_like](...)。VM 的执行逻辑可以理解成取每个 VM 指令载入对应的 TIR 内核传参并执行把结果包成新的NDArray。所以你不需要手动管理中间变量的释放VM 的内存池会帮你处理。6.3 性能验证与基本 profiling跑通不代表能上线我一般会做两件事先是“预热 计时”再是用 profiling 定位瓶颈。预热很重要因为第一次调用时有初始化开销直接计时会得到一个虚高的数据。正确姿势是for _ in range(10): out vm[main](x_tvm, w_tvm, v_tvm) import time start time.perf_counter() for _ in range(100): out vm[main](x_tvm, w_tvm, v_tvm) print((time.perf_counter() - start) / 100 * 1000, ms/iter)再进一步你可以拿到 VM 内部 profilereport vm.profile() print(report)它会列出每个算子的调用次数和耗时帮你定位瓶颈是计算密集还是访存密集。比如 matmul 占大头那是正常如果 multiply 这种逐元素算子占了太多比例就要考虑是不是 pass 没做算子融合导致多次内核启动。7. 常见问题与排查技巧实录这块内容是我踩坑踩出来的比前面某些代码更值钱。我按频率排序写几个常见问题。7.1 形状和 dtype 不匹配报错让人看不懂Relax 导入阶段经常出现Check failed: struct_info或shape mismatch这类错误。原因通常是你给R.Tensor((4, 4), float32)传了 (4, 8) 的输入或者 dtype 用了float64。我建议你在写 IR 之前先统一约定 shape 和 dtype。用列表/元组表示 shape 时常量和变量不要混用R.Tensor((batch, 4), float32)里的batch如果是个 Python int 没问题但如果是 symbolic var就要确认脚本解析器支持。排查时打开调试日志最直接export TVM_LOG_DEBUG1跑完后在输出里搜Check failed或raise ValueError一般会定位到具体 op 和结构体。不用全看看前几条关键信息就够了。7.2 不同版本的 API 差异Relax 历史版本里出现过relax.vm.build和relax.build两种写法不同分支里frontend.from_onnx的导入位置也变过。我遇到最多的是旧版本写法新版本写法relax.vm.build(mod, target)relax.build(mod, target)tvm.relax.frontend.from_pytorch(scripted_model, shape)tvm.relax.frontend.from_pytorch(model, shape, use_fxTrue)relax.VMExecrelax.VirtualMachine我的应对策略是动手前列一张表先print([x for x in dir(relax) if build in x])看下当前版本有哪些相关 API。别嫌丑这比搜教程靠谱得多因为教程可能对应三个月前的分支。7.3 运行时出现 link 错误或者 unknown symbol这种情况多是编译目标不对。你在llvmtarget 下构建运行却指定成 CUDA 设备或者反过来都会报符号找不到。另一个常见原因是 device id 没写对多卡机器上tvm.cuda(0)和tvm.cuda(1)的显存数据不能混用要在同一个设备上分配所有输入和输出。建议在构建前就把 target 和 device 配对好写成一个 config 对象传下去而不是在多个函数里各写各的。7.4 前端导入器报不支持的算子导入器报某算子不支持时我会分三档处理如果是不常用算子优先在算子映射表里看有没有替代。如果没有替代就在模型层面把它替换成等价计算组合。如果实在替换不了就回退到先把模型切成几段把不兼容段用原生库计算其余段走 Relax。这个思路在线上的容错方案里很实用也能帮你准确定位 Relax 目前有能力处理的部分。8. 一些实操习惯写给刚开始用 Relax 的朋友最后说点我自己的体会。第一条别一上来就想用 Relax 炼一个大模型。我建议先跑通本文里的最小矩阵乘法再用 ONNX 导入一个小网络最后再碰动态 shape 和自定义 pass。每一步都打印mod.script()尽量把每一步的输出和预期对齐。这能帮你建立对 IR 的直觉。第二条把mod.script()当作理解 IR 的第一工具。手写的 IR 和框架导入的 IR格式上有细微差异通过script()输出你能直观看到dataflow块、call_tir、函数签名是怎么组织的。很多问题在驱动层面报错但根因其实在图层你把 IR 打开看一眼通常就能发现端口没对上。第三条多用dir()和help()去探索 API而不是死记硬背。Relax 迭代快官方示例不一定会同步最新写法但每个模块的 docstring 往往在自己源码里。我现在遇到不确定的函数第一反应是help(relax.build)而不是百度或问 AI因为源码里的参数说明最准确。第四点也算提醒如果项目要长期维护最好把 TVM 固定到一个 commit 或 release 版本并在代码里注释当时的版本号。Relax 的 pass 名、IR 构造方法都可能会变半年后你回看自己代码如果没有版本记录会非常痛苦。这套流程走下来创建 Relax 这件事其实并不神秘。核心就一句话先选好创建路径脚本、BlockBuilder 还是前端导入再构造出 IRModule最后交给 build 和 VirtualMachine 执行。真正难的从来不是“创建”而是创建之后怎么调优、怎么排查、怎么把它接进你自己的编译管线。这些内容我可以后面再写几篇把 pass 定制、自动调优和内存规划逐个拆开聊。