[论文解读] TVM: End-to-End Optimization Stack for Deep Learning

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

Parallel Computing and Optimization Techniques参考文献 12被引用 118

一句话总结

TVM 是一个端到端的优化堆栈，可在包括移动设备、嵌入式系统以及 FPGA 和 ASIC 等加速器在内的多种硬件后端上实现高性能的深度学习推理。它通过图级和算子级优化（如算子融合、内存复用和高效的内存延迟隐藏）实现具有竞争力的性能，并采用开源编译器基础设施。

ABSTRACT

Scalable frameworks, such as TensorFlow, MXNet, Caffe, and PyTorch drive the current popularity and utility of deep learning. However, these frameworks are optimized for a narrow range of server-class GPUs and deploying workloads to other platforms such as mobile phones, embedded devices, and specialized accelerators (e.g., FPGAs, ASICs) requires laborious manual effort. We propose TVM, an end-to-end optimization stack that exposes graph-level and operator-level optimizations to provide performance portability to deep learning workloads across diverse hardware back-ends. We discuss the optimization challenges specific to deep learning that TVM solves: high-level operator fusion, low-level memory reuse across threads, mapping to arbitrary hardware primitives, and memory latency hiding. Experimental results demonstrate that TVM delivers performance across hardware back-ends that are competitive with state-of-the-art libraries for low-power CPU and server-class GPUs. We also demonstrate TVM's ability to target new hardware accelerator back-ends by targeting an FPGA-based generic deep learning accelerator. The compiler infrastructure is open sourced.

研究动机与目标

• 解决深度学习框架在异构硬件平台之间缺乏性能可移植性的问题。
• 减少在非服务器 GPU（如移动设备和嵌入式设备）上部署深度学习模型所需的手动工作量。
• 为 FPGA 和 ASIC 等专用加速器提供高效的工作负载编译与优化能力。
• 暴露低级别优化（如内存复用和线程级并行）以提升在各种后端上的性能。
• 提供一个统一且可扩展的编译器堆栈，同时支持高层图融合和底层硬件映射。

提出的方法

• 通过算子融合暴露图级优化，以减少内核启动开销并改善数据局部性。
• 应用低级别优化（如线程块间的内存复用）以最小化内存带宽压力。
• 通过将深度学习算子映射到高效的低级代码生成，支持任意硬件原语。
• 实现内存延迟隐藏技术，以在异构架构上重叠计算与内存访问。
• 使用模块化、可扩展的编译器堆栈，支持为不同硬件后端插拔式代码生成。
• 通过前端接口与现有深度学习框架集成，以支持端到端编译。

实验结果

研究问题

• RQ1 如何在包括移动 CPU 和专用加速器在内的多种硬件后端上高效地编译和优化深度学习工作负载？
• RQ2 在低功耗和嵌入式设备上，哪些优化技术最有效以实现高性能？
• RQ3 统一的编译器堆栈是否能在服务器级 GPU 和低功耗 CPU 上实现与手写优化库相当的性能？
• RQ4 该系统在针对新型硬件加速器（如基于 FPGA 的深度学习加速器）时效果如何？
• RQ5 图级和算子级优化在实现异构平台间性能可移植性方面发挥什么作用？

主要发现

• TVM 在低功耗 CPU 和服务器级 GPU 上实现了与最先进手写优化库相当的性能。
• 该系统成功针对基于 FPGA 的通用深度学习加速器进行了优化，证明了其对新型硬件的可扩展性。
• TVM 通过算子融合和内存复用显著提升了性能，有效缓解了内存带宽瓶颈。
• TVM 中的内存延迟隐藏技术提高了在延迟受限架构上的计算资源利用率。
• 开源编译器基础设施使得在多种硬件平台上快速原型设计和部署深度学习模型成为可能。
• TVM 的端到端优化堆栈减少了在不同硬件后端之间进行手动调优的需求。