Apache TVM Unity：2022 年 ML 软件和硬件生态系统的愿景

•


Apache TVM Unity 是 TVM 生态系统在 2022 年的路线图。我们看到，面对快速变化的硬件环境，机器学习系统堆栈在优化灵活性和敏捷性方面正在发生更广泛的转变。TVM 将不断发展，打破现有 ML 系统在适应 ML 模型和实现这些模型的加速器快速变化时所受到的限制。

现代 ML 系统堆栈中的边界

image

现代机器学习的系统堆栈由四种抽象组成

  1. 计算图 抽象编码粗粒度张量算子之间的数据流。计算图是用户在 TensorFlow、MXNet 和 PyTorch 中交互的高级抽象。
  2. 张量程序 实现计算图中算子的代码。深度学习编译器为卷积或矩阵乘法等计算生成低级 C++ 或 CUDA 代码。
  3. 类似地，库和运行时 包括预先编写的代码，用于执行和编排张量运算。BLAS 包和 cuDNN 等库为特定硬件目标提供经过广泛调整的算子实现。
  4. 硬件原语 位于堆栈的底部。在这里，低级汇编语言和硬件加速器接口暴露了机器的原始能力。

抽象级别之间存在垂直边界，这些边界禁止级别之间的跨层交互和反馈。在软件堆栈处理中心张量计算级别的两种对立方式之间，也存在水平边界。水平边界划分了张量计算的基于库和基于编译的方法。

image

基于库的框架依赖于预制的、经过仔细调整的算子实现集合作为其计算主力。而基于编译的框架则从头开始生成自己的自定义张量运算代码。现代软件堆栈通常使用其中一种风格，但它们不会将两者结合起来：大多数深度学习框架是基于库的，而大多数深度学习编译器无法整合库和运行时。

在当前的 ML 系统格局中，这些层之间的边界往往是严格的。两种方法都各有优缺点。基于库的堆栈在标准风格的 ML 模型上表现出色，因为它们受益于多年来对常用算子的工程投入。另一方面，基于编译的框架的灵活性和自动化可能更适合需要新算子的新兴模型。

垂直边界存在于两种风格的软件堆栈中。AI 应用程序从堆栈的顶部开始，从上到下遍历各层。框架在图级别选择数据布局和算子融合策略；然后张量计算执行计算图中选择的算子；这些算子映射到一组固定的硬件原语。这是一个一次性的、单向的工作流程：例如，张量程序级别的性能约束无法反馈到影响计算图级别的数据布局。而整合自定义硬件通常意味着手动将新功能传播到所有三层。

垂直和水平边界都在减缓机器学习的创新步伐。新的硬件加速器不断涌现，具有新的能力和性能水平，但驾驭它们将需要 ML 科学家、ML 工程师、硬件供应商之间的流畅协作，而这些边界阻碍了这种协作。为了应对 ML 系统快速变化的步伐，框架需要支持增量式演进：整合新功能应该只需要与变更成比例的努力，而不是在每个级别进行全面的重新工程。

TVM Unity

TVM Unity 的愿景是打破这些障碍。目标是实现跨层交互并自动化其优化。这并不是要将抽象层坍缩成一个整体：对于 AI 程序来说，没有“银弹”式的表示能够同时在每个级别实现优化。相反，TVM Unity 将构建接口，使抽象能够交互和交换信息。

消除系统堆栈中级别之间的严格障碍将实现跨层协同工作的新型优化。统一的系统视图将使 TVM 能够自动协同优化计算图、张量算子和硬件映射中的决策，从而搜索 AI 应用程序的最佳实现。与此同时，TVM Unity 也将充当 ML 科学家、ML 工程师和硬件工程师之间交互的通信基底。这种协作对于适应 ML 硬件加速下一阶段即将到来的快速变化至关重要。

统一抽象

image

TVM Unity 将专注于让 AI 应用程序流畅地跨越算子图、张量程序和硬件原语之间的边界。在 TVM 中，单个 Python 程序可以定义核心张量运算、整合自定义硬件原语，并从更大的算子图中调用该运算。以下示例展示了所有这些功能

import tvm.script
from tvm.script import tir as T, relax as R

@tvm.script.ir_module
class MyIRModule:
    # Define a TIR based operation.
	@T.prim_func
	def tir_mm(X: T.Buffer[(n, d), "float32"],
                   W: T.Buffer[(d, m), "float32"],
                   Y: T.Buffer[(n, m), "float32"]):
        for i, j, k  in T.grid(n, m, d):
            with T.block("body"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
		with T.init():
            Y[vi, vj] = 0
        # Can be mapped on to HW intrinsics.
        Y[vi, vj] += X[vi, vk] * W[vk, wj]

	@R.function
	def relax_func(x: R.Tensor[(n, d), "float32"], w: R.Tensor[(d, m), "float32"]):
        with R.dataflow()
            # Invoke the TIR code.
            lv0: R.Tensor[(n, m), "float32"] = R.call_dps((n, m), tir_mm, [x, w])
            lv1: R.Tensor[(n * m,), "float32"] = R.flatten(lv0)
            gv0: R.Tensor[lv2, "float32"] = R.exp(lv1)
            R.output(gv0)

        # Invoke external update rule.
        R.call_packed("custom_inplace_update", gv0)
        return gv0

此代码既有张量程序（tir_mm），也包含它的计算图（relax_func）。高层数据流可以直接调用低层张量操作来构建更大的计算。TVM 运行时统一了算子图和基于编译器的张量计算，以优化整个程序。此代码还使用 call_packed 来调用预先构建的算子——展示了 TVM 如何将基于库的算子与自定义计算平滑地集成在一起。

此外，TensorIR 通过张量化为利用硬件原语打开了大门。张量化将循环级程序转换为映射到特定硬件目标声明的原语的实现。

这里要强调的关键是跨层交互。我们的具体示例展示了以下交互：（1）计算图和张量程序；（2）计算图和运行时库；（3）最后，通过 TensorIR 中正在进行的自动张量化开发，张量程序和硬件原语之间的交互。这些跨层交互为在边界处进行增量优化打开了大门。例如，我们可以构建一个自定义 pass，将子图的较低部分传递给一组运行时库，然后再传递给管道的其余部分。

除了抽象层的统一之外，我们还在努力统一形状表示，以实现跨堆栈的一流符号形状支持。在我们的具体示例中，符号形状维度 (n, m) 可以跨抽象流动，并为动态工作负载实现高级优化。这些附加功能将为训练和推理工作负载优化打开大门。

统一视角

更好的 ML 系统需要 ML 科学家、ML 工程师和硬件工程师之间的协作。即将到来的多样化专业 ML 硬件时代将需要包含所有这三个群体的团队进行协调努力。通过在系统堆栈的各层之间构建丰富、双向的接口，TVM Unity 旨在成为这种协作和迭代发生的媒介。

TVM 中的抽象可以催化 AI 应用程序改进的生命周期。在最高层，ML 科学家可以指定他们构建下一代模型所需的算子。ML 工程师可以在张量计算层工作，使这个新运算高效。最后，这些张量计算可以依赖于硬件工程师编写的硬件原语。每个级别的工作都将通过 TVM 生态系统内的 Python API 进行交互。在 TVM 内部协同工作，而不是每次添加新功能都侵入性地修改框架，这将是在快速发展的硬件面前实现快速迭代的关键。

自动化

统一的 ML 系统创建了一个比具有严格边界的系统堆栈更大、更新的搜索空间。张量计算中的决策会影响算子图的结构，而新的硬件原语可能会彻底改变每个其他层级的最优映射。

TVM Unity 将公开所有这些跨层交互以进行自动化优化。为给定应用程序找到最佳实现将需要学习驱动的优化：使用 ML 来优化 ML，通过探索扩展的联合搜索空间并最小化计算成本。

除此之外，我们还希望在可能的情况下利用领域专家的帮助，并创建机制来有效地整合领域信息，以帮助指导自动优化。

Unity 的新功能

Unity 愿景指导了 TVM 未来一年发展的技术路线图。统一的方法将使 TVM 能够提供当今系统堆栈无法实现的新形式的自动化和生态系统集成。

借助 Unity，TVM 将统一基于库的计算和基于编译器的自动化。AI 应用程序将能够将世界上最著名的常用算子代码与为不完全映射到任何现有算子的计算自动优化的代码相结合。开发人员将能够在两种策略之间平滑过渡，而不会在从内置代码切换到生成的代码时出现陡峭的“性能悬崖”。团队将能够使用为新模型设计编译的代码快速迭代，然后，随着模型成熟和稳定，流畅地整合优化的算子库以最大化性能。通过消除基于算子和基于编译器的堆栈之间的边界，TVM 将实现对两种极端之间权衡空间的自动探索。

TVM 还旨在充当桥梁，统一更广泛的 ML 和硬件生态系统。在 ML 生态系统中，TVM 提供了一个最小的运行时，不会限制团队对框架的选择。TVM 模型将很容易嵌入到其他框架和运行时中，作为训练和推理的子图。通过 ONNX 和 TorchScript 等交换格式，TVM 模型可以流畅地集成到基于任何基础设施构建的更大应用程序中。在硬件生态系统中，TVM 已经是加速器设计人员与 ML 应用程序集成的最佳方式。借助 TVM Unity，硬件供应商将通过一组简单的算子轻松地加入 TVM，然后逐步过渡到基于编译器的集成，以获得更好的灵活性。这样，新的硬件功能就可以开始改进 AI 应用程序，而无需重新发明整个系统堆栈。

image

除了 TVM 本身之外，驱动 TVM Unity 的相同力量也存在于现代 ML 的理论和实践中。模型的快速变化、新兴的替代硬件以及老化的抽象边界都指向了对集成方法的需求。我们期望 TVM 能够引领行业迈向 ML 系统领域下一个伟大的转变。

有关我们 TVM 愿景的更多详细信息，请查看 TVMCon 2021，了解更多演讲和讨论。