「TVM教程」创建 Relax
内容提要
TVM 0.21.0 中文文档已对齐。本教程介绍创建 Relax 程序的多种方式：使用 TVMScript 定义含高级算子的 MLP 模型，并可在 Relax 函数中调用 TensorIR；使用类 PyTorch 的 NNModule API 定义模型并导出为 IRModule，支持插入 TE、TensorIR 及打包函数；还介绍更底层的 Block Builder API，适合编写自定义 Pass 等复杂转换场景。
延伸解读
三种创建方式，如何选择？
文章介绍了 TVMScript、NNModule API 和 Block Builder API 三种创建 Relax 程序的方式。TVMScript 适合直接编写 IR，语法简洁；NNModule API 类似 PyTorch，更直观易用，适合定义神经网络模型；Block Builder API 则是最底层的接口，与 IR 定义紧密耦合，适合编写自定义 Pass 等复杂转换。普通用户推荐前两种，需要更高灵活性时再考虑 Block Builder。
TVMScript 的语法糖与规范化输出
使用 TVMScript 编写代码时，可以合并多个操作为一行，例如 `lv0 = R.matmul(data, R.permute_dims(w0)) + b0`。但打印 IRModule 时，TVMScript 会以标准格式输出，每个绑定只包含一个操作，因此会看到多个中间变量。这并非错误，而是规范化结果，有助于理解 IR 的实际结构。
在 Relax 中调用 TensorIR 与外部函数
Relax 支持跨层级表示，可以在 Relax 函数中直接调用 TensorIR 函数，例如通过 `R.call_tir` 调用自定义的 relu 实现。此外，NNModule API 还允许插入 TE、TensorIR 或打包函数，如使用 `nn.extern` 调用外部函数，或通过 `nn.tensor_ir_op` 调用 TensorIR。这为混合使用不同层级的优化提供了灵活性。
Q&A
TVM 0.21.0 中创建 Relax 程序有哪些主要方式？
主要有三种方式：使用 TVMScript 定义 Relax 函数和程序；使用类似 PyTorch 的 NNModule API 定义模型并导出为 IRModule；使用更底层的 Block Builder API 构建 IR。此外，还可以在 Relax 函数中调用 TensorIR 或打包函数。
如何在 TVMScript 中定义一个包含高级算子的 MLP 模型？
使用 @I.ir_module 装饰器定义类，在类中用 @R.function 定义 forward 函数，函数参数为输入张量和权重，函数体内使用 R.dataflow() 块，通过 R.matmul、R.permute_dims、R.add、R.nn.relu 等高级算子构建计算图，最后用 R.output 指定输出并返回。
在 Relax 函数中如何调用 TensorIR 函数？
可以在 IRModule 中先用 @T.prim_func 定义 TensorIR 函数，然后在 Relax 函数中通过 R.call_tir 调用，传入 TensorIR 函数引用、输入张量以及输出张量的结构信息（out_sinfo）。
NNModule API 和 TVMScript 有什么区别？
NNModule API 类似 PyTorch，通过定义类继承 nn.Module 并实现 forward 方法来构建模型，更直观易用，之后通过 export_tvm 导出为 IRModule。TVMScript 是一种 DSL，直接以类似 Python 的语法编写 IR，更底层灵活。两者都可以定义 Relax 程序，但 NNModule API 更适合快速原型设计。
如何在 NNModule 中插入自定义的 TensorIR 或 TE 函数？
在 NNModule 的 forward 方法中，可以使用 nn.extern 调用外部打包函数，使用 nn.tensor_expr_op 调用 TE 算子（如 topi.nn.relu），使用 nn.tensor_ir_op 调用 TensorIR 函数。需要提供输出张量的占位符（out=nn.Tensor.placeholder）。
Block Builder API 适用于什么场景？
Block Builder API 是更底层的 IR 构建接口，与 IR 定义紧密耦合，广泛应用于 TVM 内部逻辑，例如编写自定义 Pass 等复杂转换场景。对于仅需定义和转换机器学习模型的普通用户，推荐使用 TVMScript 或 NNModule API；需要进行更复杂转换的用户可使用 Block Builder API 获得更高灵活性。
