深度学习编译器Relax入门：从TVM到动态图编译实战

1. 从TVM到Relax为什么我们需要一个新的编译器前端如果你接触过深度学习模型部署大概率听说过TVM。这个开源编译器栈在过去几年里几乎成了将PyTorch、TensorFlow模型高效部署到各类硬件从CPU、GPU到各种边缘端AI加速器的代名词。它的核心价值在于通过一套统一的中间表示IR将不同框架的模型转换成可被后端硬件高效执行的代码。然而对于很多开发者尤其是从应用层切入的工程师来说TVM的学习曲线一直是个挑战。它的IR设计严谨但复杂其核心IR我们常说的“Relay”在表达动态计算图、控制流以及混合编程范式时显得有些力不从心。这就是Relax诞生的背景。它不是要取代TVM而是TVM生态的一次重要演进和补充。你可以把Relax理解为TVM的“2.0版本前端IR”或者一个更友好、更强大的“新接口”。它的核心目标是解决现代深度学习模型编译中几个日益突出的痛点动态性支持如今的模型越来越“活”。动态形状Dynamic Shape——比如处理可变长度的文本序列或不同尺寸的图片——已成为常态。传统的静态图IR如早期的Relay在处理这类问题时需要大量变通甚至难以直接支持。Relax从设计之初就将动态性作为一等公民。渐进式 lowering传统的编译流程往往是一个“黑盒”从高层IR到底层代码的转换一步到位中间状态难以观察和干预。Relax引入了明确的、多阶段的Lowering过程。你可以清晰地看到你的模型是如何从直观的Python描述一步步被分解、优化最终变成针对特定硬件的内核代码的。这大大提升了编译过程的可控性和可调试性。对Python生态的深度拥抱Relax允许你直接用Python语法通过TVMScript来定义计算图。这意味着你不再需要学习一套全新的、复杂的IR构建API可以直接用你熟悉的if、for循环、shape变量来构建模型。这对于快速原型验证和算法研究人员来说是个巨大的福音。简单来说Relax让TVM变得更“平易近人”同时赋予了它处理更复杂、更现代模型的能力。它不是一个孤立的工具而是TVM编译流水线中承上启下的关键一环对上它友好地接纳来自PyTorch、TensorFlow、ONNX等框架的模型对下它通过清晰的Lowering步骤将计算任务交给TVM成熟的后端代码生成和优化器。2. 环境准备搭建你的第一个Relax工作空间在开始写第一行Relax代码之前一个稳定、可复现的构建环境是必不可少的。与直接pip install tvm安装预编译包不同为了获得最完整的功能和最新的特性我强烈建议从源码编译TVM包含Relax。这个过程看似繁琐但能让你彻底掌控依赖版本也是深入理解TVM架构的第一步。2.1 系统依赖与源码获取首先确保你的开发环境以Ubuntu 20.04/22.04为例已安装必要的构建工具和库。sudo apt-get update sudo apt-get install -y \ git \ cmake \ build-essential \ python3-dev \ python3-pip \ libtinfo-dev \ zlib1g-dev \ libedit-dev \ libxml2-dev接下来克隆TVM的主仓库。建议使用--recursive参数因为它包含了一些必要的子模块如VTA一个用于硬件加速研究的子项目。git clone --recursive https://github.com/apache/tvm.git cd tvm注意TVM的代码库较大且主分支main开发活跃。如果你追求稳定性可以切换到某个发布版本分支例如git checkout v0.14.0。但为了体验最新的Relax特性我们通常使用main分支。2.2 配置与编译关键选项解析TVM使用CMake进行构建。我们创建一个独立的构建目录并运行CMake进行配置。mkdir build cd build cp ../cmake/config.cmake .现在用文本编辑器打开config.cmake文件。这个文件决定了TVM将包含哪些功能模块。对于Relax开发以下几个选项至关重要# 启用CUDA支持如果你有NVIDIA GPU并打算进行GPU编译和测试 set(USE_CUDA ON) # 启用cuDNN用于加速某些算子如卷积 set(USE_CUDNN ON) # 启用CUTLASS这是一个高效的GPU矩阵计算库TVM可以用它来生成高性能的GPU内核 set(USE_CUTLASS ON) # 启用LLVM。这是**必须**开启的选项因为Relax的很多优化和代码生成特别是针对CPU依赖于LLVM。 set(USE_LLVM ON) # 启用Relax。确保这一行是开启的这是编译Relax模块的前提。 set(USE_RELAX ON) # 启用TVMScript。这允许我们用Python直接编写Relax IR是核心开发接口。 set(USE_TVMSCRIPT ON) # 启用MKL用于加速Intel CPU上的数学运算 # set(USE_MKL ON) # 启用OpenCL用于支持AMD GPU或集成显卡 # set(USE_OPENCL ON)编辑完成后保存文件。接下来执行CMake生成构建文件并开始编译。-j后面的数字表示并行编译的线程数可以根据你的CPU核心数调整以加快编译速度。cmake .. make -j4编译过程可能需要10到30分钟取决于你的机器性能和开启的选项。如果遇到依赖缺失的错误通常错误信息会提示你缺少哪个库按照提示安装即可。2.3 Python环境绑定与验证编译完成后我们需要让Python能够找到TVM的库。最直接的方法是将TVM的Python包路径添加到你的Python环境变量中。# 假设你在tvm/build目录下 cd ../python # 将当前目录tvm/python添加到PYTHONPATH export PYTHONPATHpwd:$PYTHONPATH为了方便你可以将上述export命令添加到你的shell配置文件如~/.bashrc或~/.zshrc中。为了立即验证安装是否成功打开一个Python解释器import tvm print(tvm.__version__) # 尝试导入relax这是我们的主角 from tvm import relax print(“Relax module imported successfully!”)如果没有报错恭喜你Relax环境已经准备就绪。这个从源码构建的过程虽然步骤多但它确保了所有组件版本的一致性避免了预编译包可能存在的二进制接口不匹配问题为后续的深入开发和调试打下了坚实基础。3. 初识Relax IR从一段简单的TVMScript开始理解了“为什么”和“如何搭建”现在让我们直面Relax本身。Relax的核心是一种中间表示IR但与我们过去在纸上画的计算图不同它是一种结构化的、机器可读也可写的文本或内存对象。为了让人也能直观地编写TVM团队创造了TVMScript它允许我们用近乎Python的语法来定义Relax IR。让我们从一个最简单的例子开始实现一个向量加法函数C A B。import tvm from tvm.script import relax as R from tvm.script import tir as T R.function def main( A: R.Tensor((1024,), dtype“float32”), B: R.Tensor((1024,), dtype“float32”) ) - R.Tensor((1024,), dtype“float32”): # R.add 是一个Relax内置的算子Op它执行逐元素加法 C R.add(A, B) # 在Relax中函数的最后一个值会被隐式返回 return C这段代码看起来就像普通的Python函数但被R.function装饰器修饰。我们来拆解其中的关键元素R.Tensor这是Relax中表示张量多维数组的类型注解。(1024,)指定了张量的形状这里是包含1024个元素的一维向量dtype“float32”指定了数据类型。注意这里的形状(1024,)在编译期是已知的是一个静态形状Static Shape的例子。R.add这是一个Relax算子Operator。Relax提供了丰富的内置算子库涵盖数学运算、张量操作、控制流等。这些算子是构建计算图的基本单元。计算图构建C R.add(A, B)这行代码并没有立即执行计算而是在构建一个计算图节点。它描述了“存在一个加法操作其输入是A和B输出是C”。实际的数值计算要等到这个图被编译并运行在具体设备上时才会发生。为了验证和查看这个函数对应的Relax IR我们可以这样做# 获取上面定义的函数对象 func main # 打印其IR文本表示 print(func.script())输出会显示更接近编译器内部表示的文本格式它精确描述了函数的签名和主体。但TVMScript的魅力在于我们几乎不需要直接阅读或编写这种底层IR文本。3.1 动态形状初探让维度“活”起来静态形状对于部署固定尺寸的模型如某些图像分类模型是足够的。但现实世界的数据往往是变化的。Relax如何应对答案是符号变量Symbolic Variable。R.function def dynamic_vector_add( A: R.Tensor((“n”,), dtype“float32”), # 使用字符串“n”作为符号变量 B: R.Tensor((“n”,), dtype“float32”) ) - R.Tensor((“n”,), dtype“float32”): C R.add(A, B) return C看我们把形状从具体的1024换成了字符串“n”。这个“n”就是一个符号变量它代表一个在编译时未知、在运行时才能确定的维度。这个函数现在可以处理任意长度但A和B长度必须相等的一维向量了。这里有一个非常重要的实操心得当你用符号变量定义函数后在编译和运行它时必须提供具体的形状值来实例化这些符号。我们会在后续的编译运行章节详细说明。3.2 引入控制流条件与循环Relax的强大之处还在于它原生支持控制流这使得它能够编译更复杂的模型逻辑例如包含if-else分支或for循环的模型。R.function def conditional_relu( X: R.Tensor((“n”,), dtype“float32”), threshold: R.Tensor((), dtype“float32”) # 标量0维张量 ) - R.Tensor((“n”,), dtype“float32”): # 使用R.if_else来实现条件判断 # 它的逻辑是如果条件成立返回第一个结果true分支否则返回第二个结果false分支 # 这里我们逐元素判断 X threshold return R.if_else( R.all(X threshold), # 条件X中所有元素都大于threshold吗这里简化处理实际可能需逐元素判断 X, # True分支返回X本身相当于ReLU的线性部分 threshold # False分支返回threshold一个简化的处理实际ReLU是max(0,x) )这个例子展示了一个简化的条件ReLU。R.if_else是Relax中表示条件执行的原语。需要注意的是Relax的控制流是数据流层面的控制流它决定的是计算图中不同子图分支的执行路径而不是Python代码的执行流。对于循环Relax提供了R.loop等原语用于处理像RNN中随时间步展开的循环结构。这些高级特性使得Relax能够表达非常复杂的模型计算图。通过这几个例子你应该对Relax IR的“样子”和基本能力有了直观感受。它用接近Python的方式描述了一个静态或动态的、可能包含控制流的计算图。接下来我们要让这个图“动”起来。4. 编译与运行将Relax函数转化为可执行代码编写TVMScript只是第一步就像写好了蓝图。编译Build的过程就是将这张高级的蓝图转化为针对特定硬件如你的CPU或GPU的、可高效执行的机器代码或运行时函数。4.1 构建Relax模块与虚拟机代码生成一个Relax函数通常被组织在一个IRModule中。IRModule是TVM中管理多个函数可能是Relax函数也可能是底层TIR函数的容器。# 我们使用之前定义的 dynamic_vector_add 函数 ir_module tvm.IRModule({“dynamic_vector_add”: dynamic_vector_add}) print(ir_module.script())现在我们需要将这个高级的IRModule进行编译。编译的核心步骤是“Lowering”和“Codegen”。# 1. 首先我们需要一个目标设备。这里以LLVM为后端的CPU为例。 target tvm.target.Target(“llvm”) # 如果是GPU可能是 “cuda” 或 “rocm” # 2. 使用Relax的构建接口进行编译。 # 这一步会执行一系列优化如算子融合、常量折叠、内存规划等并生成底层代码。 ex relax.build(ir_module, target) # ex 是一个可执行模块Executable它包含了编译后的代码以及运行时所需的信息。relax.build内部发生了很多事情Legalization将高层的、抽象的Relax算子如R.add转换为更底层的、可能由多个基础算子组成的实现序列。Lowering将Relax IR逐步降低Lower到TVM的底层IR——TIRTensor IR。TIR更接近硬件显式地描述了循环、内存加载存储等细节。优化在TIR层面进行循环优化、并行化、向量化等硬件相关优化。代码生成调用对应的代码生成器如LLVM、CUDA、OpenCL将TIR转换成目标硬件上的原生代码如.so动态库、.ptxGPU代码。4.2 运行时执行与形状绑定编译完成后我们得到了可执行模块ex。接下来需要创建一个运行时环境来执行它。# 创建一个虚拟机VM运行时来执行编译好的模块。 # Relax目前主要使用基于寄存器的虚拟机VM作为其运行时执行引擎。 vm relax.VirtualMachine(ex, tvm.cpu()) # 如果编译目标是GPU这里第二个参数应为 tvm.cuda(0) # 准备输入数据。我们需要创建符合函数签名要求的TVM NDArray。 import numpy as np n 5 # 这次我们运行一个长度为5的向量 A_np np.random.randn(n).astype(“float32”) B_np np.random.randn(n).astype(“float32”) # 将NumPy数组转换为TVM NDArray A_nd tvm.nd.array(A_np) B_nd tvm.nd.array(B_np) # **关键步骤处理动态形状** # 我们的函数 dynamic_vector_add 接受一个符号形状“n”。 # 在调用时我们需要通过一个“形状字典”来指定“n”的具体值。 shape_dict {“n”: n} # 执行函数 # 注意我们调用的是编译后模块中的函数“dynamic_vector_add”并传入形状字典和输入参数。 result_nd vm[“dynamic_vector_add”](shape_dict, A_nd, B_nd) # 将结果转换回NumPy数组进行验证 result_np result_nd.numpy() expected_np A_np B_np print(“TVM Result:”, result_np) print(“NumPy Result:”, expected_np) print(“Close?”, np.allclose(result_np, expected_np))这里有一个极易踩坑的点对于动态形状的函数调用时必须提供shape_dict且其键值对必须与函数签名中定义的符号变量名和传入张量的实际形状匹配。如果忘记提供或者提供的形状与实际数据形状不符运行时将会报错。对于静态形状的函数则不需要提供shape_dict。4.3 调试与IR变换观察编译过程像是一个黑盒吗在Relax中不是。你可以插入观察点查看Lowering过程中间状态的IR这对于理解和调试优化过程至关重要。# 在build之前我们可以设置一个调试钩子或者手动调用Lowering的各个pass。 # 一个常用的方法是使用 relax.transform.Sequential 来组合并运行特定的优化pass序列 # 并在每一步之后打印IR。 # 首先获取默认的优化pass序列build内部使用的 mod ir_module # 我们可以手动应用一些pass并观察 from tvm.relax.transform import LegalizeOps, ToNonDataflow, CallTIRRewrite from tvm.relax.transform import LowerAllocTensor, VMBuiltinLower, VMShapeLower # 例如先进行算子合法化LegalizeOps mod LegalizeOps()(mod) print(“ After LegalizeOps ”) print(mod.script()) # 然后转换为非数据流格式ToNonDataflow这是为虚拟机执行做准备 mod ToNonDataflow()(mod) print(“\n After ToNonDataflow ”) print(mod.script()) # ... 继续应用其他pass通过这种分步观察的方式你可以清晰地看到你的高级Relax函数是如何一步步被分解、重写、最终变成更底层表示的。这对于定位性能瓶颈或理解某个优化为何未生效非常有帮助。5. 集成真实模型从PyTorch到Relax的完整流程学习一个编译器前端最终目标是为了处理真实的深度学习模型。让我们走通一个完整的流程将一个简单的PyTorch模型通过Relax编译并部署运行。5.1 准备一个简单的PyTorch模型我们创建一个包含卷积、ReLU和全连接层的小网络。import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 16, kernel_size3, padding1) self.relu nn.ReLU() self.pool nn.AdaptiveAvgPool2d((1, 1)) self.fc nn.Linear(16, 10) def forward(self, x): x self.conv(x) x self.relu(x) x self.pool(x) x x.flatten(1) x self.fc(x) return x model SimpleCNN() model.eval() # 切换到评估模式5.2 使用TorchScript导出并导入TVMTVM提供了tvm.relax.frontend.from_pytorch工具链可以将PyTorch模型通过TorchScript导入为Relax IRModule。import torch from tvm import relax from tvm.relax.frontend.torch import from_pytorch # 1. 创建一个示例输入用于追踪模型生成计算图 example_input torch.randn(1, 3, 32, 32) # [batch, channel, height, width] # 2. 使用from_pytorch进行转换 # 这个过程会执行TorchScript tracing并将追踪到的算子转换为Relax算子。 mod: relax.IRModule from_pytorch( model, example_input, keep_params_as_inputTrue # 将模型参数也作为函数的输入便于后续绑定 ) # 查看转换得到的Relax模块 print(“Exported IRModule:”) print(mod.script())转换完成后mod中就包含了模型的计算图以Relax函数的形式以及模型的所有参数通常被提取为模型的附加属性或作为函数的额外输入。5.3 权重绑定与模型编译转换得到的IRModule中模型的权重可能被表示为常量R.const或作为函数参数。我们需要将这些权重与实际的数值绑定。# 假设转换后的主函数名为“main”并且参数包含输入数据和模型权重 # 首先从PyTorch模型中获取权重并转换为TVM NDArray格式 params {} for name, param in model.named_parameters(): params[name] tvm.nd.array(param.detach().cpu().numpy()) # 使用relax.transform.BindParams这个pass将IRModule中的参数占位符绑定到具体的NDArray值。 from tvm.relax.transform import BindParams mod BindParams(“main”, params)(mod) # “main”是你的主函数名 # 现在mod中的“main”函数可能就只剩下一个输入即数据输入了权重已经内联为常量。 print(“IRModule after binding params:”) print(mod[“main”].script())绑定参数后就可以像编译我们自己写的函数一样编译这个模型了。target tvm.target.Target(“llvm”) # 或 “cuda” ex relax.build(mod, target) vm relax.VirtualMachine(ex, tvm.cpu())5.4 运行与验证最后我们生成一些随机数据用编译好的模型进行推理并与PyTorch原生结果对比验证正确性。# 准备输入数据 input_np np.random.randn(1, 3, 32, 32).astype(“float32”) input_nd tvm.nd.array(input_np) # 运行TVM编译后的模型 # 注意由于我们绑定了参数现在“main”函数可能只需要一个输入数据 tvm_output_nd vm[“main”](input_nd) tvm_output_np tvm_output_nd.numpy() # 运行PyTorch原生模型进行对比 with torch.no_grad(): torch_input torch.from_numpy(input_np) torch_output model(torch_input).numpy() print(“TVM output shape:”, tvm_output_np.shape) print(“Torch output shape:”, torch_output.shape) print(“Output close?”, np.allclose(tvm_output_np, torch_output, rtol1e-3, atol1e-3))如果一切顺利两个输出应该非常接近。至此你已经完成了一个完整的“PyTorch模型 - Relax IR - 编译优化 - 部署运行”的流程。对于更复杂的模型流程是类似的但可能会遇到更多算子支持、形状推导、图优化等方面的挑战这需要更深入的学习和调试。6. 性能调优与调试实战让模型跑得更快将模型跑通只是第一步让模型在目标硬件上高效运行才是编译器的核心价值。Relax和TVM提供了丰富的工具链来辅助性能分析和调优。6.1 利用TIR进行底层调度与优化当Relax函数被Lower到TIR层后我们就获得了对循环、内存访问等底层细节的控制权。TVM的TIR层提供了强大的调度原语Schedule Primitives允许我们手动或自动地改变计算循环的执行方式以优化缓存、提高并行度。假设我们有一个简单的Relax函数它执行一个矩阵乘法后接一个ReLU。编译后我们可以提取其对应的TIR函数并进行调度。# 假设我们有一个已经Lower到TIR的模块 lowered_mod # 我们可以从中找到对应的TIR函数 sch tvm.tir.Schedule(lowered_mod) # 找到主要的计算块block通常命名为“T_matmul”或类似 matmul_block sch.get_block(“T_matmul”) # 获取这个块周围的循环 i, j, k sch.get_loops(matmul_block) # 进行调度优化示例 # 1. 循环平铺Tiling将j循环分解为jo和ji以改善缓存局部性 jo, ji sch.split(j, factors[None, 4]) # 将j循环拆分成外循环jo和内循环ji内循环大小为4 # 2. 循环重排序Reordering sch.reorder(i, jo, k, ji) # 3. 并行化Parallel将最外层的i循环标记为并行执行 sch.parallel(i) # 应用调度后重新构建模块 tuned_mod sch.mod手动调度需要深厚的硬件架构和性能优化知识。对于大多数用户TVM的AutoTVM和Ansor等自动调优工具是更实用的选择。它们可以自动搜索最优的调度参数。Relax可以与这些自动调优器协同工作你可以指定需要调优的子图往往是计算密集的算子如卷积、矩阵乘让调优器为其生成高效的TIR代码。6.2 使用Profiler定位性能瓶颈当模型运行速度不如预期时首先需要找到瓶颈在哪里。TVM提供了内置的Profiler。# 在创建VirtualMachine时可以传入一个启用了profiling的运行时。 # 首先编译时需要开启 profiling 支持部分后端支持 # 这里以CPU为例使用“debug”模式构建可能会包含更多profiling信息但更常用的是通过自定义运行时。 # 一个更直接的方法是使用TVM的 runtime.profiler。 from tvm.contrib import utils import tempfile # 创建一个临时目录存储profile数据 tmpdir tempfile.mkdtemp() profile_file utils.tempdir(tmpdir).relpath(“profile.json”) # 目前更详细的profiling通常需要在构建时通过特定选项开启或者使用像“graph_executor”配合自定义计时器。 # 对于初步评估一个简单有效的方法是使用Python的time模块对VM调用进行计时。 import time warmup 10 num_repeats 100 total_time 0.0 # 预热 for _ in range(warmup): vm[“main”](input_nd) # 计时循环 for _ in range(num_repeats): start time.perf_counter() vm[“main”](input_nd) end time.perf_counter() total_time (end - start) avg_time total_time / num_repeats print(f“Average inference time: {avg_time*1000:.2f} ms”)对于更细致的算子级别性能分析你可能需要依赖后端硬件提供的工具如NVIDIA的Nsight Systems用于GPU或Linux的perf用于CPU。TVM生成的底层代码如CUDA内核可以被这些工具识别和剖析。6.3 常见性能问题与排查思路内存拷贝开销过大频繁在CPU和GPU之间或者在TVM NDArray和NumPy数组之间拷贝数据会带来巨大开销。确保你的流水线中数据尽可能留在设备内存中。算子融合未生效TVM/Relax的一个重要优化是算子融合Operator Fusion即将多个连续的小算子如ConvReLUPooling合并成一个大的内核减少中间结果的读写。如果性能不佳检查编译日志或IR变换过程看融合是否成功。有时需要手动定义融合规则或调整图优化pass的顺序。使用了低效的默认实现对于某些算子如果没有为你的特定硬件如带有特定指令集的CPU实现优化版本TVM可能会回退到通用的、低效的实现。确保你开启了正确的编译选项如USE_MKL,USE_CUDNN并考虑使用AutoTVM为关键算子搜索最优实现。动态形状引入的额外开销完全动态的形状会阻止很多静态优化如常量传播、预分配内存。如果可能尽量使用“半动态”形状如批量大小动态但图像尺寸固定或者在运行时根据实际形状缓存多个编译好的内核版本。调试性能是一个迭代过程分析Profiling - 假设Hypothesis - 优化Optimization如调度、融合 - 验证Validation。Relax提供的可观测的Lowering过程是进行这种迭代的强大助手。7. 进阶话题Relax生态与未来展望掌握了Relax的基础和核心工作流后你的视野可以投向更广阔的领域了解Relax如何融入更大的ML编译生态系统以及它正在演进的方向。7.1 与MLC-LLM等大模型编译栈的协同Relax不仅是TVM的一部分它也是新兴的MLC-LLMMachine Learning Compilation for LLMs项目的核心前端。MLC-LLM旨在为大型语言模型LLM提供端到端的、可移植的高效部署方案。它利用Relax的动态形状能力和对控制流的良好支持来编译像GPT、LLaMA这样结构复杂、序列长度可变的模型。在MLC-LLM中Relax负责描述模型整体的计算图结构包括自注意力机制、层归一化、前馈网络等并处理动态的序列长度。然后结合TVM的TIR和AutoTensorization等技术为模型中的关键计算如矩阵乘、注意力评分生成高度优化的硬件代码。如果你关注大模型部署学习Relax是深入理解MLC-LLM工作原理的必经之路。7.2 自定义算子扩展当内置算子不够用时尽管Relax和TVM覆盖了广泛的算子但在研究或部署最新模型时你仍可能遇到不支持的算子。这时你需要扩展Relax。扩展主要有两种方式通过TIR编写自定义算子这是最底层、最灵活的方式。你直接用TIRTensor IR描述算子的计算逻辑和内存访问模式。这需要你熟悉TIR的语法和调度。T.prim_func def my_custom_add(A: T.Buffer[(1024,), “float32”], B: T.Buffer[(1024,), “float32”], C: T.Buffer[(1024,), “float32”]): T.func_attr({“global_symbol”: “my_custom_add”}) for i in T.grid(1024): with T.block(“add”): vi T.axis.spatial(1024, i) C[vi] A[vi] B[vi] T.float32(1.0) # 一个加了偏置的加法然后你可以通过relax.call_tir在Relax函数中调用这个自定义的TIR函数。通过Extern函数调用外部库如果你的算子已经有高度优化的第三方实现如CuBLAS中的某个特殊函数你可以将其封装为一个“Extern”函数。在Relax中声明它的签名然后在运行时链接到对应的库。这种方式将计算委托给外部引擎Relax只负责数据搬运和流程控制。7.3 图优化Pass的定制Relax的编译流程是由一系列“Pass”变换组成的。你可以编写自己的Pass来对Relax IR进行定制化的优化或变换。例如你可以写一个Pass来识别特定的算子模式并将其替换为更高效的实现或者插入一些性能 profiling 的指令。from tvm import relax from tvm.relax import PyExprMutator from tvm.relax.expr import Call class MyCustomFusePass(PyExprMutator): def visit_call_(self, call: Call): # 遍历计算图如果遇到连续的 add - relu 模式将其融合为一个新的算子 # 这里只是一个框架示例 new_call self.visit(call) if (isinstance(new_call, Call) and new_call.op.name “add”): # 检查下一个节点是否是 relu ... pass return new_call relax.transform.function_pass(opt_level1) def my_fuse_function(mod, ctx): return MyCustomFusePass().visit_module(mod)将这个Pass插入到你的编译流程中就可以实现自定义的图优化。这为高级用户和框架开发者提供了极大的灵活性。从环境搭建到编写第一个动态图函数从编译部署真实模型到性能调优和生态扩展Relax的学习路径是一条从用户到贡献者的道路。它降低了使用编译技术的门槛同时又没有牺牲深度和灵活性。无论是为了部署一个简单的模型还是为了研究最新的编译优化技术Relax都提供了一个坚实而现代的起点。