[Paper Review] 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 TechniquesComputer Science12 references118 citations

TL;DR

TVM is an end-to-end optimization stack that enables performance-portable deep learning inference across diverse hardware back-ends, including mobile devices, embedded systems, and accelerators like FPGAs and ASICs. It achieves competitive performance through graph-level and operator-level optimizations, such as operator fusion, memory reuse, and efficient memory latency hiding, with open-source compiler infrastructure.

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.

Motivation & Objective

• Address the lack of performance portability in deep learning frameworks across heterogeneous hardware platforms.
• Reduce the manual effort required to deploy deep learning models on non-server GPUs, such as mobile and embedded devices.
• Enable efficient compilation and optimization of deep learning workloads for specialized accelerators like FPGAs and ASICs.
• Expose low-level optimizations such as memory reuse and thread-level parallelism to improve performance on diverse back-ends.
• Provide a unified, extensible compiler stack that supports both high-level graph fusion and low-level hardware mapping.

Proposed method

• Expose graph-level optimizations through operator fusion to reduce kernel launch overhead and improve data locality.
• Apply low-level optimizations such as memory reuse across thread blocks to minimize memory bandwidth pressure.
• Support arbitrary hardware primitives by mapping deep learning operators to efficient low-level code generation.
• Implement memory latency hiding techniques to overlap computation and memory access on heterogeneous architectures.
• Use a modular, extensible compiler stack that allows pluggable code generation for different hardware back-ends.
• Integrate with existing deep learning frameworks via a frontend interface to support end-to-end compilation.

Experimental results

Research questions

• RQ1 How can deep learning workloads be efficiently compiled and optimized across diverse hardware back-ends, including mobile CPUs and specialized accelerators?
• RQ2 What optimization techniques are most effective for achieving high performance on low-power and embedded devices?
• RQ3 Can a unified compiler stack achieve performance comparable to hand-optimized libraries on both server-class GPUs and low-power CPUs?
• RQ4 How effective is the system in targeting new hardware accelerators, such as FPGA-based deep learning accelerators?
• RQ5 What role do graph-level and operator-level optimizations play in enabling performance portability across heterogeneous platforms?

Key findings

• TVM achieves performance competitive with state-of-the-art hand-optimized libraries on low-power CPUs and server-class GPUs.
• The system successfully targets an FPGA-based generic deep learning accelerator, demonstrating its extensibility to new hardware.
• TVM enables significant performance improvements through operator fusion and memory reuse, reducing memory bandwidth bottlenecks.
• Memory latency hiding techniques in TVM improve utilization of compute resources on latency-bound architectures.
• The open-source compiler infrastructure enables rapid prototyping and deployment of deep learning models on diverse hardware platforms.
• TVM’s end-to-end optimization stack reduces the need for manual tuning across different hardware back-ends.