【TVM教程】Relax

TVM 现已更新到 0.21.0 版本，TVM 中文文档已经和新版本对齐。

Apache TVM 是一个深度的深度学习编译框架，适用于 CPU、GPU 和各种机器学习加速芯片。

在线运行 TVM 学习教程→go.hyper.ai/PEh1Q

Relax 是 Apache TVM 栈中用于图优化和转换的高级抽象层。此外，Apache TVM 将 Relax 和 TensorIR 结合在一起，作为跨层优化的统一策略。因此，Relax 通常与 TensorIR 紧密协作，用于表示和优化整个 IRModule。

什么是图抽象？​

图抽象是将机器学习模型表示为有向图的过程，其中节点代表计算操作（例如矩阵乘法、卷积），边表示操作之间的数据流。这种抽象方式使编译器能够分析模型中不同部分之间的依赖关系与联系。

from tvm.script import relax as R

@R.function
def main(
    x: R.Tensor((1, 784), dtype="float32"),
    weight: R.Tensor((784, 256), dtype="float32"),
    bias: R.Tensor((256,), dtype="float32"),
) -> R.Tensor((1, 256), dtype="float32"):
    with R.dataflow():
        lv0 = R.matmul(x, weight)
        lv1 = R.add(lv0, bias)
        gv = R.nn.relu(lv1)
        R.output(gv)
    return gv

Relax 的关键特性​

Relax 是 Apache TVM Unity 策略中使用的图表示方法，它通过以下几个重要特性支持对机器学习模型的端到端优化：

• 头等符号形状（symbolic shape）：  Relax 使用符号形状来表示张量的维度，使得编译器能够在张量操作与函数调用之间全局追踪动态形状关系。
• 多层次抽象：  Relax 支持从高层的神经网络层到低层的张量操作的跨层次抽象，使得优化可以在模型的多个层级之间协同展开。
• 可组合的变换（Composable transformations）：  Relax 提供了可组合的变换框架，可以选择性地应用于模型的不同部分。这包括部分 lower（向低层转换）和部分特化（partial specialization）等能力，从而提供灵活的定制与优化手段。

这些特性共同赋予 Relax 在 Apache TVM 生态系统中强大而灵活的机器学习模型优化能力。