Advertisement

TVM: End-to-End Optimization Stack for Deep Learning

阅读量:

why TVM

为了实现 AI 芯片对深度学习架构(包括 TensorFlow、MXNet、Caffe 和 PyTorch 等)的支持,在设计时需要考虑芯片自身的指令集(如汇编语言或 C 语言)以及如何将这些深度学习架构中的代码转换为相应的指令集进行部署。这就要求从系统层面进行全面的设计工作,并制定一套全方位优化方案以确保系统的高效运行。

这里写图片描述

TVM的主要目标在于实现深度学习架构在各种不同的内存结构和计算单元配置的硬件设备上的易部署性,在这一目标下涵盖了从GPU(如NVIDIA Tesla系列)、FPGA(Field-Programmable Gate Arrays)、ASIC(特别是谷歌TPU)以及各种嵌入式设备等多种应用场景。

这里写图片描述

TVM设计了一个全链路优化架构,并通过一种将算法描述、调度策略与硬件接口分离的方式实现了对深度学习工作负载的适应性增强。通过将调度机制与目标硬件内部函数的实现细节区分开来并进行扩展性设计,在支持新型专用加速器及其内部函数方面取得显著进展。

Optimizing Computational Graphs

Computation Graph: 计算图,在人工智能系统中用于表示整个计算流程。构建计算图时会预先分配内存空间以支持复杂运算需求。优化过程则是生成具有相同功能但效率更高的计算流程。例如,在神经网络训练过程中可以通过预处理某些运算并在后续阶段调用这些中间结果来显著提升运行效率(如上所示)。下文将详细展示这一类优化技术的具体实现方式及其实际应用效果:

这里写图片描述

Operator Fusion:运算符融合技术通过将多个运算符整合到同一计算过程中从而避免了中间结果存储在全局缓存中的步骤这一优化方法可以在GPU及某些加速设备上实现以减少运算开销。

这里写图片描述
  1. 数据布局转换:在机器学习框架中进行数据组织转换是实现高效计算的关键步骤之一。其中一种重要的运算类型是Tensor运算(Tensor Operations),这类运算是构建计算图的基础模块(core components)。由于不同类型的运算对数据的组织需求各不相同,在实际应用中往往需要根据具体场景动态地调整数据排列方式以提高计算效率。例如,在许多深度学习加速器中倾向于采用4×4大小的张量进行处理(tackle),因此在这种情况下将输入的数据划分为若干个互不重叠的子块并分别对其进行处理能够显著提升局部缓存效率(cache efficiency)。为了更好地理解这一机制,请参考图1所示的具体实现方案。这种结构设计特别适用于执行2×2大小的Tensor运算(tensor operations of size 2x2),其核心优势在于能够在有限资源条件下最大化并行计算能力(parallelism)的同时保证计算精度(accuracy)。在这一过程中,默认的做法是针对每一个特定的操作节点(operation node)选择最适合其需求的数据排列模式(data organization pattern),而这些模式的选择直接决定了最终硬件实现的质量表现。
这里写图片描述
  1. Limitations of Graph-Level Optimizations

优化张量算子

Tensor Expression Language:一种用于表示张量运算的语言体系,在TVM框架中采用了这一独特的表示方法;如图所示,则展示了TVM系统中所采用的这种基于张量运算的语言结构;其中每一项运算都涉及输出张量的形状信息以及具体的运算步骤;而每个表达式则明确了如何计算出每个输出张量的具体元素

这里写图片描述

张量表达语言包含C等普通编程语言中的运算和运算符。TVM提供了一个可交换的简化算子,并被应用于多线程调度。随后计划引入一个高层次的扫描操作符。这个高级扫描操作符将基础计算运算结合成循环计算。

  1. Schedule Space:调度空间域中,在任何硬件终端上实现其功能仍然具有挑战性。下图展示了在CPU、GPU以及深度学习加速器上的优化应用。
这里写图片描述

每个经过底层优化的程序与不同硬件后端之间建立了各自的调度策略整合,并由此给Kernel设计者带来了相当大的负担。因此,在借鉴Halide的过程中,我们实施了分离计算描述以及调度器优化过程两种方式进行改进。具体而言,在运行时系统中,在基于硬件后端所设定的具体规则下将计算描述转化为经过优化后的硬件实现;而TVM中对调度算法进行了改进,并提供了有效的调度基元以减少计算量与原来经过改进后的调度算法相比相差很多

  1. 合作式嵌套并行化(Nested Parallelism with Cooperation),一种用于提升深度学习工作负载中密集型计算效率的关键技术。由于大多数GPU都具备较高的多线程处理能力,在编程实现时可以通过采用多线程方式来显著提高运算效率。TVM框架通过引入基于异步执行的设计理念,在单个操作中实现多个子操作的协同执行:每个子操作都可以递归地细分为更小的任务以实现多层次线程处理,并且这些独立运行的任务将在join阶段(即当一个子任务完成,并且后续阶段需要上一阶段生成的数据时)进行数据交换。这种设计模型明确禁止同一操作层次内的不同线程之间进行协作互动,在矩阵乘法操作示例中可直观体现其局限性:单独的一个矩阵乘法操作涉及大量独立的小矩阵相乘步骤,并且这些步骤之间彼此脱节无法相互影响。”
这里写图片描述

为了弥补该操作所导致的不足,在这一领域探索了一种新的方法。这一方法的核心在于采用了更为有效的优化方法,在许多现代计算架构中展现出显著的优势。然而还未成功整合到调度机制中。TVM则通过引入记忆化分层机制将这些时间段划分为共享内存段;如果未采用该机制,则这些时间段被视为本地线程。

  1. Tensorization: Generalizing the Hardware Interface,张量化:生成硬件接口。深度学习的工作任务中有高强度的算法运算,这些运算可以分解为张量运算(例如矩阵和矩阵的乘法,一维卷积操作),所以近些年来越来越多的使用张量来进行运算。但是这些计算本质上是很不相同的,例如矩阵与矩阵的乘法,矩阵和向量的product以及一维的卷积操作,这些不同的计算都使张量算子调度更加困难。
    张量的形式与向量的形式特别不同,张量计算的输入是多维的、固定的或者是可变的长度,并且有不同的数据布局。所以我们不能借助于一系列固定的算法就不适用了,因为新的深度学习加速器采用特殊的Tensor指令。为了解决上面的问题,我们从调度器中分离了硬件接口。TVM引进了一套Tensor内联声明机制。可以使用Tensor表达式语言来声明每个新的硬件内联的行为,和给他分配底层原语是一样的。另外TVM引进了一个调度器基元Tensorization来作为基元的计算单元
    如图所示:
这里写图片描述

该语言能够既表示用户准备好的计算描述信息又表示硬件暴露接口的抽象信息。通过Tensorization技术将调度器从硬件基元中解耦出来后,则使得TVM更容易扩展新的 hardware 架构。经过实践表明,在实际应用中使用基于 Tensorization 技术设计出的调度器生成代码具有很高的计算性能:将复杂的操作分解为多个重复的小型 Kernel 函数调用。从而可以看出,在实际应用中使用基于 Tensorization 基元构建的手工优化汇编小型 Kernel 的方式能显著提升性能优势,并且这种方法在许多不同的 hardware 平台上都有很好的应用前景。

  1. Compiler Support for Latency Hiding: 编译器支持延迟隐藏技术。该技术通过并行执行内存计算操作来最大化内存利用率与计算效率的同时提升系统性能表现。在不同类型的硬件设备上需要采用不同的策略来实现这一目标:对于基于CPU的系统而言通常会采用同步多线程的方式;而对于依赖GPU架构的应用则倾向于快速切换线程以充分利用计算单元资源;在一些深度学习加速器中则会将其应用转移到专门设计的堆栈编译器环境中以便进一步优化性能表现。从硬件层面实现这一技术的方法主要包括精确控制硬件资源以确保在特定条件下能够启动多个独立的任务序列从而避免资源冲突;这种方法虽然有效但带来了较高的开销因此开发人员转而寻求更为高效的解决方案:TVM为此提供了一种高效的虚拟线程调度机制能够自动将数据并行程序转换为底层的数据依赖型程序具体的转换流程如图所示
这里写图片描述

上图所示的算法开始是一个高层次并行程序;随后,在每个虚拟线程中插入了必要同步指令以保证每个算子在各自线程中的正确执行顺序;最终将所有虚拟线程的操作整合成为一个单一的实际操作。

Code Generation and Runtime Support

对图和调度进行优化后,剩下的任务就是产生能在目标平台以及能够部署在硬件设备上的代码了。
1. Code Generation:TVM通过对前端框架的硬件如所用到的CPU,GPU进行分析,将高级语言转换成硬件设备可以支持的C语言。
对于一个特定的多元数据流声明,坐标系相关的超图和调度树,我们可以通过迭代遍历调度树的方案生成Lowered代码,并且推断出输入Tensor的依赖范围(使用坐标系相关的超图),接下来生成循环嵌套的low-level代码。low-level代码是通过类C的循环程序具象化而来,在这个过程中使用了一个Halide循环程序数据结构的变体,我们也使用了Halide常用的lowering primitives,像是storage flattening,循环展开,对于GPU和特定深度学习加速器而言,则使用了同步点检测、虚拟线程注入、模块化生成机制。最终,循环程序被翻译成LLVM/CUDA/Metal/OpenCL源代码
2. Runtime Support For GPU programs: TVM会独立的构建host端和device端的模块并且提供一个运行时模块系统来启动Kernel。代码生成算法然后将加速器程序翻译成一系列Runtime API。
3. Autotuning:TVM提供一套新颖的优化框架,这套框架能够为深度学习系统编译高性能底层实现,使得像高维卷积、矩阵乘法和深度卷积这样的复杂操作可以被一套调度器模板进行自动调优。
4. Remote Deployment Profiling:为了嵌入式设备,TVM设计了一套方便性能分析和自动调优的基础设施。传统情况下,嵌入式开发一般是在主机上进行交叉编译然后复制可执行文件到目标设备上执行,上述编译、运行和性能分析的工作都需要手动进行。在编译器堆栈中我们提供了一个远程程序调用:通过RPC接口,我们能够完全在host端完成上面的所有步骤,这样的方式可以极大地加快在嵌入式设备和基于FPGA的加速器上的优化工作。如下图:

这里写图片描述

Evaluation

下图为GPU Evaluation

这里写图片描述

下图为FPGA Accelerator Evaluation

这里写图片描述

全部评论 (0)

还没有任何评论哟~