「TVM教程」理解 Relax 抽象层
内容提要
TVM 0.21.0 发布，中文文档同步更新。文章介绍 Relax 抽象，用于端到端优化机器学习模型，以两层神经网络为例，展示高层 Numpy 与 Relax 代码，及底层实现。关键元素包括结构信息、R.call_tir 调用底层函数、数据流块标记纯函数区域，便于编译器优化如算子融合。
延伸解读
Relax 抽象层的核心价值
Relax 是 TVM Unity 策略中的图抽象，旨在端到端优化机器学习模型。它通过高层描述模型结构与数据流，同时支持调用底层原始张量函数，实现跨层抽象。这种设计使得编译器能够全局追踪动态形状关系，并执行如算子融合等优化，从而提升模型在多种硬件上的执行效率。
R.call_tir 与目标传递
R.call_tir 是 Relax 中调用底层函数的关键机制，采用目标传递约定：输出张量由外部显式分配，然后传入底层函数执行。这种风格便于高层框架控制内存分配，但并非所有操作都适用，尤其当输出形状依赖输入时。实际使用中，尽量采用此风格编写低层函数通常有益。
数据流块与纯函数
数据流块用于标记无副作用的计算图区域，其中所有操作必须为纯函数。手动标记而非自动推断，是因为自动推断可能不准确，尤其在涉及打包函数时。手动标记能确保编译器准确理解数据流结构，从而不错过关键优化机会，如算子融合。
Q&A
Relax 抽象层在 TVM 中的主要作用是什么？
Relax 是 Apache TVM Unity 策略中使用的一种图抽象方式，用于对机器学习模型进行端到端的优化。它的主要目标是描述机器学习模型的结构与数据流，包括模型不同部分之间的依赖关系与连接方式，以及如何在硬件上执行该模型。
Relax 中的 R.call_tir 是什么？它如何工作？
R.call_tir 是 Relax 中的新抽象，用于在同一个 IRModule 中调用底层的原始张量函数（primitive tensor functions）。它先分配一个输出张量 res，然后将输入和输出一起传入 prim_func，执行后结果写入 res 并返回。这种约定称为 destination passing，便于高层框架控制内存分配。
Relax 中的数据流块（dataflow block）是什么？为什么需要手动标记？
数据流块是 Relax 中用于标记程序中计算图区域的方式，在数据流块内所有操作必须是无副作用的（纯函数），块外可以有副作用。手动标记的原因有两个：一是自动推断数据流块很困难，尤其涉及打包函数时容易不准确；二是许多优化（如算子融合）仅限于数据流块内，手动标记能确保编译器拥有准确信息，实现更有效的优化。
Relax 如何支持动态形状（symbolic shape）？
Relax 原生支持符号形状，例如在 main 函数中张量形状里的 "n"，以及 linear 函数中的 M、N、K。这使得编译器能够在张量操作符和函数调用之间全局追踪动态形状关系。
Relax 中的结构信息（Structure Info）是什么？
结构信息是 Relax 中用于表示 Relax 表达式类型的新概念，可以是 TensorStructInfo、TupleStructInfo 等。在 TVMScript 中，TensorStructInfo 简写为 R.Tensor，用于表示输入、输出和中间结果张量的形状和数据类型。
Relax 中纯函数和有副作用函数有什么区别？
纯函数只读取输入并通过输出返回结果，不会更改程序的其他部分；有副作用的函数会修改已有的中间或输入张量，例如 inplace 操作。在数据流块中，所有操作必须是无副作用的，块外可以有副作用。
