TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

TL;DR

TVM presents an end-to-end optimizing compiler designed to achieve performance portability for deep learning workloads across diverse hardware back-ends by decoupling high-level representations from low-level hardware specifics. It introduces a tensor expression language, operator fusion, data layout transformation, and tensorization, combined with an ML-based cost model to automatically explore and select optimized operator implementations. The framework supports end-to-end workflows from many DL frameworks to CPUs, GPUs, and accelerators including FPGA-based designs, demonstrated by competitive performance and notable speedups (up to 3.8x) over vendor-tuned libraries. TVM’s open-source stack, distributed RPC optimization, and demonstrated applicability to new back-ends highlight its practical impact for broad deployment and hardware-software co-design in DL systems.

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.

Figures (21)

• Figure 1: CPU, GPU and TPU-like accelerators require different on-chip memory architectures and compute primitives. This divergence must be addressed when generating optimized code.
• Figure 2: System overview of TVM. The current stack supports descriptions from many deep learning frameworks and exchange formats, such as CoreML and ONNX, to target major CPU, GPU and specialized accelerators.
• Figure 3: Example computational graph of a two-layer convolutional neural network. Each node in the graph represents an operation that consumes one or more tensors and produces one or more tensors. Tensor operations can be parameterized by attributes to configure their behavior (e.g., padding or strides).
• Figure 4: Performance comparison between fused and non-fused operations. TVM generates both operations. Tested on NVIDIA Titan X.
• Figure 5: Example schedule transformations that optimize a matrix multiplication on a specialized accelerator.