[论文解读] TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

Tianqi Chen, Thierry Moreau+10 more|arXiv (Cornell University)|Feb 12, 2018

Parallel Computing and Optimization Techniques参考文献 38被引用 140

一句话总结

TVM 是一个端到端的编译器，通过将图级优化和算子级优化结合起来，并利用基于 ML 的成本模型进行自动调度，使深度学习工作负载在 CPU、GPU 和加速器上实现性能可移植。

ABSTRACT

There is an increasing need to bring machine learning to a wide diversity of hardware devices. Current frameworks rely on vendor-specific operator libraries and optimize for a narrow range of server-class GPUs. Deploying workloads to new platforms -- such as mobile phones, embedded devices, and accelerators (e.g., FPGAs, ASICs) -- requires significant manual effort. We propose TVM, a compiler that exposes graph-level and operator-level optimizations to provide performance portability to deep learning workloads across diverse hardware back-ends. TVM solves optimization challenges specific to deep learning, such as high-level operator fusion, mapping to arbitrary hardware primitives, and memory latency hiding. It also automates optimization of low-level programs to hardware characteristics by employing a novel, learning-based cost modeling method for rapid exploration of code optimizations. Experimental results show that TVM delivers performance across hardware back-ends that are competitive with state-of-the-art, hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPUs. We also demonstrate TVM's ability to target new accelerator back-ends, such as the FPGA-based generic deep learning accelerator. The system is open sourced and in production use inside several major companies.

研究动机与目标

• 识别在不同硬件后端上实现性能可移植的 DL 工作负载的优化挑战。
• 开发一个端到端的编译栈，将高级 DL 程序映射到优化的低级代码。
• 通过张量表达式语言和调度空间实现联合的高低级优化。
• 利用学习型成本模型自动化算子优化，以探索庞大的优化空间。
• 在 CPU、GPU 和基于 FPGA 的加速器上展示可移植性和性能。

提出的方法

• 引入张量表达式语言和变换原语，以将计算与硬件内置操作分离。
• 将 Halide 启发的计算/调度分离扩展，以支持新型 GPU 与加速器。
• 构建由 ML 基础的成本模型引导的自动优化框架，预测降级程序的性能。
• 结合张量化，通过声明式张量内建操作和 tensorize 原语将计算映射到硬件内置操作。
• 实现内存延迟隐藏和显式内存作用域，以针对多样的内存层次结构进行优化。
• 提供一个端到端栈，将模型从框架（TensorFlow、MXNet、PyTorch 等）编译到特定硬件的优化代码。

实验结果

研究问题

• RQ1 TVM 是否能够在服务器级 GPU、嵌入式 GPU/CPU 与基于 FPGA 的加速器之间实现可移植且具有竞争力的性能，而不依赖厂商特定的算子库？
• RQ2 与黑盒自动调优相比，基于机器学习导向的优化在导航庞大的算子/调度搜索空间方面有多有效？
• RQ3 在不同后端上，将图优化（如融合、数据布局）与算子级代码生成结合起来能带来何种收益？
• RQ4 延迟隐藏和张量化在如 TPU 类设备和 FPGA 的专用加速器上如何影响性能？
• RQ5 TVM 是否能够支持新兴工作负载（如深度卷积、低精度运算）和新兴加速器，同时保持性能可移植性？

主要发现

• TVM 在后端实现可移植性能，相较现有框架中的手工优化库，速度提升为 1.2× 至 3.8×。
• 算子融合和图优化显著减少内存访问，带来显著的运行时提升。
• 通过虚拟线程和显式内存作用域实现的延迟隐藏提升了计算利用率（例如：ResNet 在 FPGA 加速器上的峰值利用率从 70% 提高到 88%）。
• 张量化将硬件内在操作与调度解耦，使对新加速器的支持更容易，并在微内核上实现高达 1.5× 的加速。
• 基于 ML 的成本模型使优化的算子实现更快地被发现，在速度和配置质量方面均优于黑盒自动调优。
• 该系统是开源的并已在产业界部署，展示了在 CPU、GPU 和 FPGA 基于加速器上的实际潜力。