【TVM教程】创建 Relax
内容提要
TVM 0.21.0 更新，中文文档已对齐。教程演示创建 Relax 函数的三种方式：TVMScript 定义高级运算符和 TensorIR 函数；NNModule API 类似 PyTorch，支持自定义调用；Block Builder API 为底层接口，适合复杂转换。示例包括 MLP 模型构建，并说明各 API 适用场景。
延伸解读
三种 API 的定位与选择
TVMScript 适合直接编写 IR，语法简洁但需熟悉 DSL；NNModule API 类似 PyTorch，适合快速定义模型；Block Builder 是底层接口，灵活但繁琐。普通用户推荐前两者，复杂转换或自定义 Pass 时用 Block Builder。
TVMScript 的语法糖与规范化
TVMScript 支持将多个操作合并为一行（如 `R.matmul(...) + b0`），但打印 IRModule 时采用标准格式，每个绑定只含一个操作。这种差异源于语法糖与规范化表示的不同，理解这一点有助于阅读生成的 IR。
跨层级混合编程
Relax 不仅是图层级 IR，还支持在 Relax 函数中直接调用 TensorIR 函数（如 `R.call_tir`）或打包函数（如 `R.call_dps_packed`）。这允许在高层模型中嵌入底层优化算子，兼顾表达力与性能。
Q&A
TVM 0.21.0 中创建 Relax 函数有哪几种方式？
TVM 0.21.0 中创建 Relax 函数主要有三种方式：使用 TVMScript 定义、使用 NNModule API（类似 PyTorch）定义，以及使用 Block Builder API 定义。
TVMScript 是什么？如何用 TVMScript 定义 Relax 程序？
TVMScript 是用于表示 Apache TVM 中间表示（IR）的领域特定语言（DSL），是 Python 的一种变体。可以用它定义包含 TensorIR 和 Relax 函数的 IRModule。例如，通过 @I.ir_module 和 @R.function 装饰器定义 Relax 函数，并在函数体内使用 R.matmul、R.nn.relu 等高级运算符。
在 Relax 函数中如何调用 TensorIR 函数？
在 Relax 函数中可以通过 R.call_tir 调用 TensorIR 函数。例如，先定义一个 TensorIR 函数（如 relu），然后在 Relax 函数中使用 R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype="float32")) 来调用它。
NNModule API 与 TVMScript 相比有什么特点？
NNModule API 提供了类似 PyTorch 的接口，更直观、更容易使用。它允许用户通过定义 nn.Module 子类来构建模型，然后使用 export_tvm 方法导出为 IRModule。而 TVMScript 更接近底层 IR，需要手动编写脚本。
如何在 NNModule 中调用自定义的 TensorIR 函数或打包函数？
在 NNModule 中，可以使用 nn.extern 调用打包函数，使用 nn.tensor_expr_op 调用 Tensor Expression 操作，使用 nn.tensor_ir_op 调用 TensorIR 函数。例如，nn.tensor_ir_op(tir_linear, "tir_linear", [x, self.fc2.weight, self.fc2.bias], out=nn.Tensor.placeholder((n, 10), "float32"))。
Block Builder API 是什么？适用于什么场景？
Block Builder API 是创建 Relax 程序的底层 IR 构建接口，与 IR 定义紧密耦合。它适用于需要编写自定义 Pass 或进行复杂转换的场景，但使用体验不如 TVMScript 和 NNModule API 友好。普通用户定义和转换机器学习模型时，推荐使用前两种 API。
TVMScript 打印输出的代码与编写的代码为什么不同？
因为打印输出时使用的是标准格式，而 TVMScript 支持语法糖简化输入。例如，可以将多个操作合并为一行，但规范化后的表达式要求一个绑定中只能包含一个操作，所以打印输出会展开为多个绑定。
