Source-linked AI summary
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
TL;DR
Deep learning deployment must span hardware with divergent characteristics, while existing frameworks depend on narrow, manually tuned operator libraries. TVM combines graph rewriting, tensor-operator scheduling, and learning-based search into an end-to-end compiler, and reports 1.2× to 3.8× speedups over framework baselines across several back-ends.
Problem
Existing deep learning frameworks focus on narrow hardware classes and rely on vendor-specific operator libraries, making support for diverse back-ends labor-intensive.
Method
TVM combines graph rewriting, tensor-expression scheduling, automated ML-guided operator search, and hardware-specific code generation in an end-to-end compilation stack.
Results
TVM achieved 1.2× to 3.8× speedups over existing frameworks backed by hand-optimized libraries across server GPU, embedded GPU, embedded CPU, and FPGA-based accelerator workloads.
Takeaways & Limitations
TVM provides portable deep learning performance across diverse hardware back-ends and supports deployment from high-level framework specifications.
Takeaways & Limitations
The reported evaluation omits 3DCGAN and LSTM because the baseline did not support them.
Abstract
from arXiv · showhide
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.
1 Introduction
TVM addresses the difficulty of deploying deep learning across diverse hardware by combining graph- and operator-level optimization with automated code generation and search. It delivers portable performance across several back-ends, including CPUs, GPUs, and FPGA-based accelerators.
- Motivation: Current frameworks rely on high-level graphs and vendor-specific operator libraries, limiting hardware coverage and requiring substantial manual tuning.These frameworks mainly target server-class GPUs, while diverse devices differ in memory organization and compute units.
- Approach: TVM takes high-level deep learning specifications and generates low-level optimized code for diverse hardware back-ends.Its end-to-end approach combines graph rewriting, operator optimization, and hardware-specific code generation.
- Approach: TVM introduces schedule primitives for memory reuse, hardware intrinsics, and latency hiding, alongside an ML-based system that searches optimized tensor operators.The cost model adapts as it collects more data from a hardware back-end.
- Scope: The stack supports workloads from major deep learning frameworks and targets CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators.TVM is open sourced and reported to be in production use inside several major companies.
- Evaluation: 1.2× to 3.8× speedups over existing frameworks backed by hand-optimized libraries were measured across server GPU, embedded GPU, embedded CPU, and FPGA-based accelerator workloads.The evaluation used real-world workloads and reported portable performance across back-ends.
2 Overview
TVM converts models from existing frameworks into optimized computational graphs, generates optimized code for fused operators, and packages the result as a deployable runtime module. Its overview emphasizes an automated path from model description to multiple deployment back-ends.
- Model input: TVM accepts models from existing frameworks and transforms them into computational graph representations before high-level dataflow rewriting.The supported stack includes descriptions from frameworks and exchange formats such as CoreML and ONNX.
- Operator optimization: The operator optimizer generates efficient code for each fused operator using declarative tensor expressions and an ML-based cost model.Execution details are unspecified in the tensor expression, allowing TVM to search a large optimization space.
- Deployment: TVM packages the optimized graph, generated operators, and parameters into a deployable module.The runtime module can be executed on the selected target back-end through TVM’s deployment APIs.
- Deployment: The system supports deployment back-ends through C++, Java, and Python and is designed to be extended to new back-ends.The overview presents a path from high-level model input through compilation to target execution.
3 Optimizing Computational Graphs
TVM optimizes computational graphs through transformations such as operator fusion, constant folding, memory planning, and data-layout conversion. These graph-level transformations reduce execution costs while preparing operators for hardware-specific scheduling and code generation.
- Graph representation: A computational graph represents tensor operations as nodes and data dependencies as edges while leaving each operator’s implementation unspecified.Tensor operations may be parameterized by attributes such as padding or strides.
- Graph optimizations: TVM applies graph-level optimizations including operator fusion, constant folding, static memory planning, and data-layout transformations.These transformations exploit the graph’s global view of tensor operations and dependencies.
- Operator fusion: Operator fusion combines multiple operators into one kernel without storing intermediate results in memory.The paper identifies injective, reduction, complex-out-fusable, and opaque operator categories for applying fusion rules.
- Operator fusion: 1.2× to 2× speedups were observed for fused operators by reducing memory accesses on workloads tested with TVM.The comparison was between fused and non-fused operations on an NVIDIA Titan X.
- Data layout: Data-layout optimization selects layouts suited to target memory hierarchies and inserts transformations when producer and consumer preferences differ.Specialized accelerators may require tiled layouts to exploit matrix operations and access locality.
- Design boundary: Graph-level optimizations depend on available operator implementations, making manually supporting growing fused-kernel patterns difficult across increasing hardware targets.TVM’s lower-level code-generation approach addresses operator implementations beyond the graph rewrite itself.
4 Generating Tensor Operations
TVM separates tensor computations from schedules, then lowers many valid implementations into hardware-specific code using extensible primitives for CPUs, GPUs, and accelerators. Its schedule space supports cooperation, tensorization, and explicit latency hiding.
- Hardware-Aware Scheduling: TVM extends Halide-style scheduling with primitives for nested parallelism, tensorization, and latency hiding across diverse back-ends.These additions target GPU and accelerator requirements while retaining existing low-level loop representations.
- Explicit Memory Latency Hiding: Virtual threading lowers a high-level parallel schedule into one instruction stream containing explicit synchronizations that recover pipeline parallelism for latency hiding.This reduces the programming burden of DAE accelerators that require low-level synchronization.
- Tensor Expression and Schedule Space: TVM describes each tensor operation with an index formula while leaving loop structure and execution details to schedules.Schedules incrementally apply logically preserving transformations to generate low-level code.
- Nested Parallelism with Cooperation: Cooperative thread fetching places shared data in shared memory, enabling reuse across threads and requiring compiler-inserted synchronization barriers.TVM exposes memory scopes so shared stages and their dependencies are represented explicitly.
- Tensorization: Tensorization replaces a matched computation pattern with declared hardware intrinsics, separating schedules from target-specific primitives.The same tensor expression language declares intrinsic behavior and lowering rules, allowing new hardware architectures to be added.
5 Automating Optimization
TVM automates operator optimization by searching schedule configurations with an adaptive ML cost model and measured hardware feedback. Distributed execution and graph-level integration support optimization across diverse workloads and devices.
- Optimization Search: TVM specializes operators for each layer’s input shape and layout, but this creates a large schedule-optimization search problem.Choices include memory access, threading, hardware primitives, loop tiling and ordering, caching, and unrolling.
- Optimization Search: Black-box auto-tuning requires many hardware experiments, while predefined cost models are difficult to build accurately and must be recreated for each hardware target.Modern hardware complicates modeling through memory access, reuse, pipeline dependencies, and threading patterns.
- ML-Based Cost Model: TVM’s ML cost model predicts runtime from lowered loop programs, learns from measurements collected during exploration, and improves with additional trials and related workloads.Schedule templates expose tunable knobs, while the explorer selects promising configurations for real measurement.
- Evaluation: The ML-based optimizer finds better configurations much faster than black-box auto-tuning methods in the reported conv2d comparison.The model begins without training data and uses collected measurements to improve itself.
- ML Model Design Choices: 0.67 ms is the average prediction time of the tree boosting model, thousands of times faster than a real measurement.Tree boosting uses loop-program features; TreeRNN summarizes the program AST, with similar predictive quality but slower prediction.
- Distributed Optimization: A distributed RPC device pool automates compilation, execution, and profiling across multiple devices, supporting both workload optimization and end-to-end graph inference.Remote devices can run cross-compiled modules while results remain accessible from the host script.
6 Evaluation
TVM is evaluated end to end across server-class and embedded platforms, low-precision workloads, mobile GPUs, and an FPGA-based accelerator. It generally outperforms framework or library baselines, while FPGA gains are limited by CPU-executed workload sections.
- Evaluation Setup: TVM was evaluated on server-class GPU, embedded GPU, embedded CPU, and FPGA-based accelerator platforms using real-world inference workloads.Benchmarks included ResNet, MobileNet, LSTM, DQN, and DCGAN workloads.
- Server-Class GPU Evaluation: 1.6× to 3.8× speedups over MXNet, TensorFlow, and TensorFlow XLA were observed on the NVIDIA Titan X.TVM’s gains came from joint graph optimization and automatically generated fused operators; DQN reached the 3.8× speedup.
- Embedded CPU Evaluation: TVM generated faster operators than hand-optimized TFLite versions for both ResNet and MobileNet on an ARM Cortex A53.The result also covered emerging depthwise convolution operators, and TVM outperformed TFLite end to end on three workloads.
- Mobile GPU Evaluation: TVM outperformed the ARM Compute Library on three available Mali-T860MP4 models for both float16 and float32, with speedups from 1.2× to 1.6×.DCGAN and LSTM were not supported by the baseline library.
- FPGA Accelerator Evaluation: 40× acceleration was achieved for convolution layers offloaded to the FPGA-based VDLA, but CPU-only workload sections bottlenecked overall performance.The accelerator used a 16×16 matrix-vector unit on a low-power PYNQ board.
7 Related Work
Related work provides graph-level representations, scheduling abstractions, auto-tuning, and accelerator compilation, but TVM targets operator-level optimization and diverse hardware through an end-to-end stack.
- Deep Learning Frameworks: Existing frameworks provide convenient graph interfaces but commonly depend on vendor-specific tensor operator libraries for execution.These libraries require substantial manual tuning and primarily serve a narrow range of hardware.
- Graph-Level Compilation: High-level graph representations support global optimizations but are too abstract for tensor-operator optimization across diverse hardware back-ends.Prior approaches use hardware-specific lowering rules or vendor-crafted libraries, requiring engineering for each back-end and operator variant.
- Scheduling and DSLs: TVM adopts compute–schedule separation from Halide while adding primitives for GPU and specialized-accelerator scheduling challenges.Its scheduling approach is related to GPU DSLs and polyhedral loop transformation, while TACO and Weld address other tensor or data-processing settings.
- Automatic Optimization: TVM uses domain-aware machine-learning cost modeling to guide optimization search, alongside auto-tuning and prior domain-agnostic search approaches.The model is designed to address the large search spaces created by tensor-program scheduling choices.
- Accelerator Compilation: The VDLA case study demonstrates a concrete compilation approach for TPU-like accelerators, an area where effective compilation stacks remain unclear.The approach could potentially benefit systems that compile deep learning workloads to FPGA targets.
8 Conclusion
TVM presents an end-to-end compilation stack for optimizing deep learning across diverse hardware back-ends. The authors position automated optimization as a basis for further compiler research and software–hardware co-design.
- TVM combines end-to-end automated optimization with support for deep learning across diverse hardware back-ends.The authors describe this as addressing fundamental optimization challenges that have historically required labor-intensive specialization.