用TVM在硬件平台上部署深度学习工作负载的端到端 IR 堆栈

用TVM在硬件平台上部署深度学习工作负载的端到端 IR 堆栈

深度学习已变得无处不在，不可或缺。这场革命的一部分是由可扩展的深度学习系统推动的，如滕索弗洛、MXNet、咖啡和皮托奇。大多数现有系统针对范围狭窄的服务器级 GPU 进行了优化，需要在其它平台，如移动电话、物联网设备和专用加速器（FPGA、ASIC）上部署大量精力。随着深度学习框架和硬件后端数量的增加，建议建立一个统一的中间表示 （IR） 堆栈，以缩小以生产力为中心的深度学习框架与面向性能或效率的硬件后端之间的差距。

 

 

 TVM 是一个新颖的框架，可以：

• 表示并优化 CPU、GPU 和其它专业硬件的常见深度学习计算工作负载
• 自动转换计算图以最大限度地减少内存利用、优化数据布局和融合计算模式
• 提供从现有前端框架到裸机硬件的端到端编译，一直到浏览器可执行的 javascripts。

TVM可以轻松地在移动电话、嵌入式设备甚至浏览器上运行深度学习工作负载，不需要额外的支持。TVM 还为众多硬件平台（包括依赖新型计算原始的专用加速器）的深度学习工作量提供了统一的优化框架。

 

 

 采用编译器方案，提供两个中间表示层，有效地将高级深度学习算法降低到多个硬件后端。

开源TVM包包含x86、ARM、OpenCL、Metal、CUDA和Javascript的原型优化。致力于专业硬件加速和 Nvidia 的 GEMM 优化Volta架构提供支持。

技术细节

TVM 堆栈的目标是提供可重复使用的工具链，以编译从深度学习框架前端到低级机器代码的高层神经网络描述，用于多个硬件后端。以 Apache MXNet 为例，以下代码片段演示，如何使用 TVM 将深度学习模型的高level描述，编译为针对目标硬件定制的优化可执行模块。

 

 

 挑战在于支持多个硬件后端，同时将计算、内存和能量消耗保持在最低水平。借用AICompiler的智慧，以弥合众多深度学习框架和硬件后端之间的差距：

构建了由 NNVM 组成的两级中间层、用于任务调度和内存管理的高级中级代表 （IR） 和用于优化计算内核的具有表现力的低level IR 的 TVM。

堆栈的第一个级别是基于计算图形的表示。计算图是指示式循环图，表示计算作为节点，数据流依赖性表示边缘。此表示非常强大：

将操作属性bake到计算图中，指定转换规则，迭代优化计算图形。这是大多数现有深度学习框架所采用的常见方法，包括 TVM 堆栈中的 NNVM 图形表示、TensorFlow XLA 和英特尔的 ngraph。

 

 

发现仅基于计算图的 IR 不足以解决支持不同硬件后端的挑战。原因是，对于每个硬件的后端，可能会以非常不同的方式映射和优化单个图形算子，如卷积或矩阵乘法。这些针对硬件的优化在内存布局、并行线程模式、缓存访问模式和硬件原始选择方面可能存在巨大差异。希望能够以通用表示方式明确表达这些优化模块，有效引导优化空间。

构建一个低level级别的表示来解决这个问题。此表示基于索引公式，为复发计算提供额外支持。

 

 

 低级 IR 采用现有图像处理语言（如Halid或darkroom）的原则，定制具有表现力的深度学习 DSL。TVM 构建低level级别优化的理论来自循环转换工具，如循环和多面体分析。从 MXNet, TensorFlow, Theano.等深度学习框架中使用的数据流描述语言中汲取灵感。然后在调度阶段处理 TVM 中描述的算法，应用针对目标硬件后端量身定制的转换。

 

TVM includes standard transformation primitives commonly found in CPU optimization frameworks. More importantly, TVM incorporates novel optimization primitives targeted at GPUs, by exploiting thread cooperation patterns, data layout transformations and powerful new compute primitives. Using TVM in combination with NNVM provides an rich opportunity to optimize deep learning workloads across the software stack, enabling joint compute graph-level and operator-level optimizations.

TVM 包括 CPU 优化框架中常见的标准原型转换。更重要的是，TVM 通过利用线程调度模式、数据布局转换和强大的新原型计算，整合了针对 GPU 的新型优化。使用 TVM 与 NNVM 相结合，为优化整个软件堆栈中的深度学习工作负载提供了丰富的机会，实现了联合计算图形级别和算子级别的优化。

多语言和平台支持

TVM 的众多优势之一，在于对多个平台和语言的丰富支持。框架的两个组件：包含完整优化库以生成优化机器代码的编译器堆栈，以及轻量级lightweight的runtime，以及在不同平台上部署编译模块所需的便携性。

 

 

 

 TVM 当前支持嵌入式编译器堆栈的python和C++接口。设计框架时，考虑到了最大的再利用，以便编译器堆栈改进，可以在 Python 和C++组件之间互换应用。

提供轻量级的runtime，可以在包括 Android、iOS、raspberry pi和 Web 浏览器在内的平台上，以javascript, java, python, C++等语言直接运行 TVM 编译代码。

远程部署和执行

 

 

 

 

 

 TVM 支持使用 TVM RPC 对嵌入式设备进行交叉编译和测试，这是一个轻量级界面，可在远程嵌入式设备上部署和执行 TVM 交叉编译模块。这为 TVM 用户提供了熟悉的高级 Python 界面，以便在各种低级嵌入式设备上远程编译、优化和测试深度学习算法。

性能

TVM能够灵活地探索各种深度学习内核的丰富优化空间，用于多个硬件平台。例如，TVM为最关心的内核和平台定制数据布局和融合模式要求。基线库是针对更通用的问题创建的，而 TVM 的优化内核则针对通过自动调整过程评估的工作负载进行了大量调整。TVM 充当快速生成获取此类专用内核的桥梁。

Raspberry Pi树莓派

在结果的第一部分，将TVM CPU调度与树莓派3B执行网格工作负载上的nnpack进行了比较。使用TVM实现直接卷积，而nnpack用于执行3x3内核的winograd conv。

 

We can find that with TVM’s autotuned kernels, we can obtain performance similar to the hand-optimized kernels in nnpack for the raspberry pi experiments. 可以发现，通过 TVM 的自动调谐内核，可以获得类似于 nnpack 中hand-optimized kernels手动优化内核的性能，用于aspberry pi实验。

GPU 结果

创建了一个端到端的编译管道，可以将 MxNet 模型编译到 TVM 执行图形。通过自动将算子融合在一起，让 TVM 生成融合内核，在图形节点内和图形节点之间应用优化。对mobilenet ImageNet的工作量进行了基准测试，并讨论了以下结果：

 

 

 TVM在速度方面可以优于baseline基线方法。更有趣的是，内核融合带来了额外的加速。值得一提的是，TVM 能够自行生成所有优化的 GPU 内核，无需依赖 CuDNN 等外部库。

开源努力

TVM最初是华盛顿大学保罗·艾伦学校计算机科学与工程系的研究项目。TVM 堆栈旨在支持DLPack，这是多个主要深度学习框架对张力数据结构的共识。收到了来自UW, AWS, Qiho 360, Facebook, HKUST, TuSimple, UCDavis, SJTU以及DMLC开源社区和DLPack倡议成员的早期贡献。今后，该项目将遵循Apache open-source model开源模式，创建社区维护项目。

• Halide：TVM使用HalideIR作为数据结构，用于arithematic简单化和低level lowering。HalideIR来自Halide。还从Halide那里学习，在TVM中实施lowing管道。
• Loopy：使用整数集分析及其循环转换原始。
• Theano：symbolic scan算子的再现设计灵感。
• Github page页面可以在这里找到： https://github.com/dmlc/tvm
• TVM 兼容DLPack，因此可以轻松支持采用该标准的框架，如 MXNet, PyTorch, Caffe2和微小 dnn。

源代码Source code