NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - AI模型在线查看 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 - 3D道路快速建模

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

1、现代机器学习系统栈中的边界

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

  • 计算图抽象（computional graph abstraction）对粗粒度张量运算符之间的数据流进行编码。 计算图是用户在 TensorFlow、MXNet 和 PyTorch 中与之交互的高级抽象。
  • 张量程序（tensor programs）实现计算图中运算符的代码。 深度学习编译器为卷积或矩阵乘法等计算生成低级 C++ 或 CUDA 代码。
  • 类似的，库和运行时（libraries and runtimes）包括预先编写的代码来执行和编排张量操作。 像 cuDNN 这样的 BLAS 包和库为特定的硬件目标提供了广泛调整的运算符实现。
  • 硬件原语（hardware primitives）位于栈的底部。 在这里，低级汇编语言和硬件加速器接口暴露了机器的原始功能。

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

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

在 ML 系统的当前格局中，这些层之间的界限往往很严格。 两种方法都不比另一种好，但它们各有优缺点。 基于库的栈在 ML 模型的标准样式上表现出色，因为它们受益于多年的工程投资通用运算符。 另一方面，基于编译的框架的灵活性和自动化可能更适合需要新运算符的新兴模型。

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

纵向和横向边界都在减缓机器学习的创新步伐。 新的硬件加速器正在出现，具有新的功能和性能水平，但利用它们将需要 ML 科学家、ML 工程师和硬件供应商之间的流畅协作，而这些边界阻碍了这些协作。 为了应对 ML 系统的快速变化，框架需要支持渐进式演化：合并新功能需要与变化成比例的努力，而不是在每个级别进行大规模的重新设计。

2、TVM Unity

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

消除系统栈中层级之间的严格障碍将使跨层协同工作的新型优化成为可能。 整个系统的统一视图将使 TVM 自动共同优化计算图、张量运算符和硬件映射中的决策，以搜索 AI 应用程序的最佳实现。 同时，TVM Unity 也将作为 ML 科学家、ML 工程师和硬件工程师之间交互的沟通基板。 这种合作对于适应 ML 硬件加速下一阶段即将到来的快速变化至关重要。

3、统一抽象

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 中持续的自动张量化开发来张量程序和硬件原语。 这些跨层交互为在边界进行增量优化打开了大门。 例如，我们可以构建一个自定义传递到子图的下部到一组运行时库，然后传递到管道的其余部分。

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

4、统一观点

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

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

5、自动化

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

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

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

6、TVM Unity 的新功能

Unity 愿景指导了 TVM 明年发展的技术路线图。 统一的方法将使 TVM 能够提供新形式的自动化和生态系统集成，这在当今的系统堆栈中是不可能的。

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

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

除了 TVM 之外，推动 TVM Unity 的力量同样存在于现代 ML 的理论和实践中。 模型的快速变化、新兴的替代硬件和老化的抽象边界都表明需要一种集成方法。 我们预计 TVM 将引领 ML 系统进入下一个全行业的重大转变。

原文链接：Apache TVM Unity: a vision for the ML software & hardware ecosystem

BimAnt翻译整理，转载请标明出处