TVM 通过在每个硬件后端生成许多有效的实现并选择优化的实现,为每个算子生成高效的代码。这个过程建立在 Halide 将描述和计算规则 (或调度优化) 解耦的想法的基础上,并将其扩展为支持新的优化 (嵌套并行、张量和延迟隐藏) 和 广泛的硬件后端。
引入了张量表达式语言来支持自动代码生成,与高级计算图表示不同,张量算子的实现是不透明的,每个算子都用索引公式表达语言来描述。以下代码展示了计算转置矩阵乘法的示例张量表达式。
![WX20221026-112920@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6xpOCeV7I9X9ToyqzAdynL21tvqibbBLj0uNYSmXiawH6ticx0t8eYHNKA/640)
每个计算操作都指定输出张量的形状 和 描述如何计算它的每个元素的表达式。张量表达式语言支持常见的算术和数学运算,并涵盖常见的 DL 算子形式。该语言没有指定循环结构和许多其他执行细节,它提供了为各种后端添加硬件感知优化的灵活性。采用来自 Halide 的解耦 计算/调度原则,使用调度来表示从张量表达式到低级代码的特定映射。
通过增量应用保持程序逻辑等价的基本转换 (调度原语) 来构建调度。图5展示了在专用加速器上调度矩阵乘法的示例。在内部,TVM 应用调度转换,使用数据结构来跟踪循环结构和其他信息。然后此信息可以帮助按给定的最终调度生成低级代码。
![WX20221026-114454@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6q0Rx81RagIBUBozzib1Sa9ZE4h9ZURym2LQZg9BAiaJjGE8EDouJtjmg/640)
作者的张量表达式借鉴了 Halide、Darkroom 和 TACO。为了在许多后端实现高性能,必须要支持足够多的调度原语来涵盖不同硬件后端的各种优化。图6总结了 TVM 支持的算子代码生成过程和调度原语。作者重用了有用的原语和来自 Halide 的低级循环程序 AST,并且引入了新的原语来优化 GPU 和 加速器的性能。新的原语是实现最佳 GPU 性能所必需的,也是加速器所必需的。CPU、GPU 和 类TPU加速器是三种重要的深度学习硬件。
![WX20221026-114551@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6cH1gfxMoicr3RUZRTmdSoQDFuVLibSeictjUtIZjm9DnOJ1NIsvZ2HfBw/640)
并行性是提高 DL 工作负载中计算密集型内核效率的关键。现代 GPU 提供了大规模的并行性,要求咱们将并行模式融合到调度转换中。大多数现有的解决方案都采用一种称为嵌套并行的模型,这是一种 fork-join 的形式。该模型需要并行调度原语来并行化数据并行任务;每个任务可以进一步递归地细分为子任务,以利用目标架构的多级线程层次结构 (例如 GPU 中的线程组)。将此模型称为无共享嵌套并行,因为一个工作线程无法在同一并行计算阶段查看其兄弟的数据。
无共享方法的替代方法是协作获取数据。具体来说,线程组可以协作获取它们都需要的数据并将其放入共享内存空间中。这种优化可以利用 GPU 内存层次结构,并通过共享内存区域实现跨线程的数据重写。TVM 使用调度原语来支持这种众所周知的 GPU 优化,以实现最佳性能。以下 GPU 代码示例对矩阵乘法进行了优化。
![WX20221026-133648@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C695X768kvwB23c4sddiaF6guVNxxjzsNYhh6viaFfaXheygy3NzCSSiaew/640)
图7展示了这种优化的影响,将内存范围的概念引入调度空间,以便可以将计算阶段 (代码中的 AS 和 BS) 标记为共享。
![WX20221026-134013@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6ibVib51d554kkK6yWAQ5HgGaIibTicObYIo406GFfrCKyoDcfiakl9kUsxg/640)
如果没有显式内存范围,自动范围推理会将计算阶段标记为 thread-local。共享任务必须计算组中所有工作线程的依赖关系,此外还必须正确插入内存同步屏障,以保证共享加载的数据对消费者可见。最后,除了对 GPU 有用之外,内存范围还让咱们可以标记特殊的内存缓冲区,并在针对专门的 DL 加速器时创建特殊的降低规则。
DL 工作负载具有很高的算术强度,通常可以分解为张量算子,如矩阵乘 或 一维卷积。这些自然分解导致了最近添加的张量计算原语的趋势。这些新的原语为基于调度的编译创造了机遇和挑战,使用它们可以提高性能,所以编译框架必须无缝的集成它们。这种张量化类似于 SIMD 架构的矢量化,但又有着显著差异。指令输入是多维的,具有固定或可变的长度,并且具有不同的数据排布。更为重要的是,这种原语不能是固定的,因为新的加速器有可能会出现张量指令的变体,因此需要一个可扩展的解决方法。
通过使用张量内在声明机制将目标硬件内在特性 与 调度分离,从而使张量化变得可以扩展。作者使用相同的张量表达式语言来声明每个新硬件内在行为 和 与之相关的降低规则。以下代码展示了如何声明一个 8x8 张量的硬件内在函数。
![WX20221026-135905@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6UVx916RPmJricTdPkeecaia88z720XbceEXlkDPAdeuBu6tPSaNDvkHA/640)
此外,引入了张量调度原语,以用相应的内在函数替换计算单元。编译器将计算模式与硬件声明相匹配,并将其降低到相应的硬件内在特性。张量化将调度与特定的硬件原语分离,从而 可以轻松扩展 TVM 以支持新的硬件架构。生成的张量调度代码符合高性能计算特性:将复杂的算子分解为一系列 micro-kernel 调用。咱们还可以使用 tensorize 原语来利用手工设计的 micro-kernel,这在某些平台上可能是有益的。例如,通过利用 bit-serial 矩阵向量乘法 micro-kernel 为移动 CPU 实现超低精度算子,这些算子对一比特或两比特宽的数据类型进行操作。这个 micro-kernel 将结果累积成越来越大的数据类型,以最大限度地减少内存占用。将 micro-kernel 表达为 TVM 固有的张量 与 非张量化版本相比 可产生高达1.5倍的加速。
延迟隐藏是指将内存操作与计算重叠以最大限度地利用内存和计算资源的过程。它需要不同的策略,具体取决于目标硬件后端。在 CPU 上,内存延迟隐藏是通过同步多线程 或 硬件预取 来隐式实现的。GPU 依赖于许多线程的快速上下文切换。相比之下,诸如 TPU 之类的专用 DL 加速器通常倾向于使用解耦访问执行 (DAE) 架构记性更加精简的控制,并将细粒度同步的问题转移给软件。
图9显示了减少运行时延迟的 DAE 硬件管道。与单片硬件设计相比,流水线可以隐藏大部分内存访问开销,几乎可以充分利用计算资源。为了实现更加高的利用率,指令流必须增加细粒度的同步操作。没有它们,就无法强制执行依赖关系,从而导致错误执行。因此,DAE 硬件流水线需要在流水线阶段依赖于细粒度的入队/出队操作,以保证正确执行,如图9的指令流所示。
![WX20221026-142833@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6MeaJL31dHxRwYdia7KVUUmBgFgKMRlMOhicJxKf42IAXtIibvtWWvQgfQ/640)
对需要显式低级同步的 DAE 加速器进行编程很困难。为了减少编程负担,引入了虚拟线程调度原语,让程序员可以指定高级数据并行程序,就像他们指定支持多线程的硬件后端一样。然后 TVM 会自动将程序降低为具有低级显式同步的单个指令流,如图8所示。该算法从高级多线程程序调度开始,然后插入必要的低级同步操作以保证每个线程的正确执行。接下来,它将所有虚拟线程的操作交付给单个指令流中。最后,硬件恢复指令流中低级同步所规定的可用流水线的并行度。
Hardware Evaluation of Latency Hiding 展示了延迟隐藏在基于 FPGA 的定制加速器设计上的有效性。在加速器上运行 ResNet 的每一层,并使用 TVM 生成两个调度:一个有延迟隐藏,一个没有。具有延迟隐藏的调度将程序与虚拟线程并行化 以进行管道并行,从而隐藏了内存访问延迟。结果在图10 中显示为 roofline图表;roofline性能图可以深入了解给定系统在不同基准测试中使用计算和内存资源的情况。总体而言,延迟隐藏提高了所有 ResNet 层的性能。峰值计算利用率从没有延迟隐藏的 70% 提高到了隐藏延迟的 88%。
![WX20221026-145123@2x.png](http://mmbiz.qpic.cn/mmbiz_png/b0xUbKsIMqP9wTJwok7cWrOJ2LscE4C6rODcET66JwAUxd0Qiajib5iarvqLRfN4KyErJAiaxibejjGZq4vs7wcX2ug/640)