企业官网建设流程全解析

1. 从TVM到Relax：为什么我们需要一个新的编译器前端？

如果你接触过深度学习模型部署，大概率听说过TVM。这个开源编译器栈在过去几年里，几乎成了将PyTorch、TensorFlow模型高效部署到各类硬件（从CPU、GPU到各种边缘端AI加速器）的代名词。它的核心价值在于，通过一套统一的中间表示（IR），将不同框架的模型转换成可被后端硬件高效执行的代码。然而，对于很多开发者，尤其是从应用层切入的工程师来说，TVM的学习曲线一直是个挑战。它的IR设计严谨但复杂，其核心IR（我们常说的“Relay”）在表达动态计算图、控制流以及混合编程范式时，显得有些力不从心。

这就是Relax诞生的背景。它不是要取代TVM，而是TVM生态的一次重要演进和补充。你可以把Relax理解为TVM的“2.0版本前端IR”，或者一个更友好、更强大的“新接口”。它的核心目标，是解决现代深度学习模型编译中几个日益突出的痛点：

1. 动态性支持：如今的模型越来越“活”。动态形状（Dynamic Shape）——比如处理可变长度的文本序列或不同尺寸的图片——已成为常态。传统的静态图IR（如早期的Relay）在处理这类问题时需要大量变通，甚至难以直接支持。Relax从设计之初就将动态性作为一等公民。
2. 渐进式 lowering：传统的编译流程往往是一个“黑盒”，从高层IR到底层代码的转换一步到位，中间状态难以观察和干预。Relax引入了明确的、多阶段的Lowering过程。你可以清晰地看到你的模型是如何从直观的Python描述，一步步被分解、优化，最终变成针对特定硬件的内核代码的。这大大提升了编译过程的可控性和可调试性。
3. 对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 PYTHONPATH=`pwd`:$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内部发生了很多事情：

1. Legalization：将高层的、抽象的Relax算子（如R.add）转换为更底层的、可能由多个基础算子组成的实现序列。
2. Lowering：将Relax IR逐步降低（Lower）到TVM的底层IR——TIR（Tensor IR）。TIR更接近硬件，显式地描述了循环、内存加载存储等细节。
3. 优化：在TIR层面进行循环优化、并行化、向量化等硬件相关优化。
4. 代码生成：调用对应的代码生成器（如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_size=3, padding=1) 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导出并导入TVM

TVM提供了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_input=True # 将模型参数也作为函数的输入，便于后续绑定 ) # 查看转换得到的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, rtol=1e-3, atol=1e-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 常见性能问题与排查思路

1. 内存拷贝开销过大：频繁在CPU和GPU之间，或者在TVM NDArray和NumPy数组之间拷贝数据，会带来巨大开销。确保你的流水线中数据尽可能留在设备内存中。
2. 算子融合未生效：TVM/Relax的一个重要优化是算子融合（Operator Fusion），即将多个连续的小算子（如Conv+ReLU+Pooling）合并成一个大的内核，减少中间结果的读写。如果性能不佳，检查编译日志或IR变换过程，看融合是否成功。有时需要手动定义融合规则或调整图优化pass的顺序。
3. 使用了低效的默认实现：对于某些算子，如果没有为你的特定硬件（如带有特定指令集的CPU）实现优化版本，TVM可能会回退到通用的、低效的实现。确保你开启了正确的编译选项（如USE_MKL,USE_CUDNN），并考虑使用AutoTVM为关键算子搜索最优实现。
4. 动态形状引入的额外开销：完全动态的形状会阻止很多静态优化（如常量传播、预分配内存）。如果可能，尽量使用“半动态”形状（如批量大小动态，但图像尺寸固定），或者在运行时根据实际形状缓存多个编译好的内核版本。

调试性能是一个迭代过程：分析（Profiling） -> 假设（Hypothesis） -> 优化（Optimization，如调度、融合） -> 验证（Validation）。Relax提供的可观测的Lowering过程，是进行这种迭代的强大助手。

7. 进阶话题：Relax生态与未来展望

掌握了Relax的基础和核心工作流后，你的视野可以投向更广阔的领域，了解Relax如何融入更大的ML编译生态系统，以及它正在演进的方向。

7.1 与MLC-LLM等大模型编译栈的协同

Relax不仅是TVM的一部分，它也是新兴的MLC-LLM（Machine Learning Compilation for LLMs）项目的核心前端。MLC-LLM旨在为大型语言模型（LLM）提供端到端的、可移植的高效部署方案。它利用Relax的动态形状能力和对控制流的良好支持，来编译像GPT、LLaMA这样结构复杂、序列长度可变的模型。

在MLC-LLM中，Relax负责描述模型整体的计算图结构，包括自注意力机制、层归一化、前馈网络等，并处理动态的序列长度。然后，结合TVM的TIR和AutoTensorization等技术，为模型中的关键计算（如矩阵乘、注意力评分）生成高度优化的硬件代码。如果你关注大模型部署，学习Relax是深入理解MLC-LLM工作原理的必经之路。

7.2 自定义算子扩展：当内置算子不够用时

尽管Relax和TVM覆盖了广泛的算子，但在研究或部署最新模型时，你仍可能遇到不支持的算子。这时，你需要扩展Relax。

扩展主要有两种方式：

1. 通过TIR编写自定义算子：这是最底层、最灵活的方式。你直接用TIR（Tensor 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函数。

2. 通过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_level=1) def my_fuse_function(mod, ctx): return MyCustomFusePass().visit_module(mod)

将这个Pass插入到你的编译流程中，就可以实现自定义的图优化。这为高级用户和框架开发者提供了极大的灵活性。

从环境搭建到编写第一个动态图函数，从编译部署真实模型到性能调优和生态扩展，Relax的学习路径是一条从用户到贡献者的道路。它降低了使用编译技术的门槛，同时又没有牺牲深度和灵活性。无论是为了部署一个简单的模型，还是为了研究最新的编译优化技术，Relax都提供了一个坚实而现代的起点。