Source-linked AI summary
Learning to Optimize Tensor Programs
Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
TL;DR
Deep learning systems depend on efficient tensor operators, but hardware-specific libraries support limited devices and impose substantial deployment engineering costs. AutoTVM learns domain-specific statistical cost models to search billions of program variants and transfers models across workloads. The framework improves end-to-end performance by 1.2× to 3.8× over existing frameworks and shows promise for diverse hardware deployment.
Problem
Hardware-specific tensor-operator libraries support a narrow range of devices, limiting deployment flexibility and creating substantial engineering costs for new hardware targets.
Method
AutoTVM uses domain-specific statistical cost models to guide schedule search and transferable program representations to accelerate optimization across workloads.
Results
AutoTVM improved end-to-end performance by 1.2× to 3.8× over existing frameworks on real-world deep learning inference workloads.
Takeaways & Limitations
Learning-based tensor-program optimization shows promise for deep learning deployment across low-power CPUs, mobile GPUs, and server-class GPUs.
Takeaways & Limitations
Considering uncertainty estimates did not improve search, and the choice of acquisition function remains open for further exploration.
Abstract
from arXiv · showhide
We introduce a learning-based framework to optimize tensor programs for deep learning workloads. Efficient implementations of tensor operators, such as matrix multiplication and high dimensional convolution, are key enablers of effective deep learning systems. However, existing systems rely on manually optimized libraries such as cuDNN where only a narrow range of server class GPUs are well-supported. The reliance on hardware-specific operator libraries limits the applicability of high-level graph optimizations and incurs significant engineering costs when deploying to new hardware targets. We use learning to remove this engineering burden. We learn domain-specific statistical cost models to guide the search of tensor operator implementations over billions of possible program variants. We further accelerate the search by effective model transfer across workloads. Experimental results show that our framework delivers performance competitive with state-of-the-art hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPU.
1 Introduction
The paper frames tensor-program optimization as a learning problem motivated by hardware diversity and the engineering burden of manually optimized libraries. AutoTVM searches program variants with statistical models and transfer learning, achieving strong end-to-end improvements.
- Motivation: Tensor operators underpin deep learning systems, but manually optimized libraries support only a narrow range of hardware devices.Equivalent implementations can differ dramatically because of threading, memory reuse, pipelining, and other hardware factors.
- Research question: The paper asks whether learning can automatically optimize tensor operator programs for a given hardware platform.Its affirmative approach uses statistical cost models to predict low-level program runtime and guide exploration.
- Contributions: AutoTVM formalizes learning to optimize tensor programs, proposes a machine-learning framework, and accelerates optimization by 2× to 10× through transfer learning.These contributions address both the search problem and the reuse of information across workloads.
2 Problem Formalization
The paper formalizes tensor-program optimization as choosing hardware-aware schedules for logically equivalent generated programs. It uses a rich search space and learning-based cost models to efficiently find high-performing schedules despite billions of candidates.
- Problem Formalization: Index expressions specify tensor operators while leaving loop order, memory scope, and threading available for low-level scheduling.The compiler generates logically equivalent low-level programs from an expression and a schedule.
- Problem Formalization: The objective is to minimize real hardware runtime f(x), which is unknown analytically but queryable by running experiments.Each schedule s in S_e produces low-level code x = g(e, s).
- Problem Characteristics: Unlike traditional black-box hyper-parameter optimization, tensor programs expose domain-specific structure that can support effective models.The experiments are relatively inexpensive, taking a few seconds to compile and run, so more training data can be collected.
- Requirements: Automatic code generation requires an exhaustive search space covering hardware-aware optimizations and an efficient method for finding an optimal schedule.The paper focuses on schedule optimization rather than improving the underlying search-space design.
- Search Space: The search space includes multi-level tiling, loop ordering, GPU shared-memory caching, unrolling, and vectorization, reaching billions of implementations for one GPU operator.The authors report that this space can contain programs competitive with hand-optimized libraries.
3 Learning to Optimize Tensor Programs
The framework uses statistical cost models to predict low-level tensor-program performance and guide hardware measurements through an efficient search loop. It combines model-based exploration, diversity-aware batch selection, and periodic model updates.
- The framework predicts each low-level program's cost, proposes schedule configurations, measures them on hardware, stores runtime data, and updates the model.
- GBT and TreeGRU models encode low-level program structure differently, with GBTs favoring fast CPU prediction and TreeGRUs offering extensibility without feature engineering.TreeGRU uses batching and GPU acceleration to make training and prediction usable in the framework.
- The search uses simulated annealing with the predicted cost as an energy function instead of enumerating the entire schedule space.Candidate programs are then measured on real hardware and incorporated into the training database.
- The exploration loop retains persistent Markov-chain states across model updates and randomly samples εb candidates to preserve exploration.The random fraction is exemplified as ε = 0.05.
- Candidate batches balance low predicted runtime with diversity across schedule-configuration components through a greedy approximation to a submodular objective.The objective's first term favors low-cost candidates, while its second term rewards coverage of different configuration components.
- Uncertainty-aware acquisition did not improve search performance in this problem, although alternative acquisition functions remain open for further exploration.
4 Accelerating Optimization via Transfer Learning
Transfer learning accelerates optimization across tensor workloads by using representations that remain usable when workloads or search spaces change. The approach combines shared global knowledge with local in-domain adaptation, including structure-aware encodings for tree-based and neural models.
- Transfer learning reuses historical data from previously seen workloads to speed optimization for operators with different shapes and data types.
- Directly representing schedules as configurations is not invariant when the search-space specification changes across workloads.
- The low-level loop AST is invariant to the search space, making it a shared program representation for transfer across workloads.The cost model takes the AST as input, then encodes it into a vector for prediction.
- GBT context relation features model relationships between loop-feature axes so representations can generalize across different loop-nest patterns.The encoding can capture relations such as loop count versus touched memory size.
- The context-encoded TreeGRU replaces workload-specific loop-variable embeddings with context vectors and sums the scattered vectors across loop levels.
- Transfer combines a global model trained on invariant historical representations with a local in-domain model that adapts as target-workload data accumulates.The global model supplies initial predictions before sufficient local data are available.
5 Prior Work
Prior tensor-program optimization used black-box search, hardware-dependent cost models, and polyhedral methods. This work applies domain-specific statistical cost modeling to enable transfer across workloads.
- Existing approaches: Black-box optimization is used in high-performance computing libraries, while alternative approaches build hardware-dependent cost models to guide search.Tensor Comprehensions combines black-box parameter optimization with polyhedral internal-loop generation.
- Existing approaches: Polyhedral methods use integer linear programming to optimize cost, whereas Tensor Comprehensions combines black-box and polyhedral optimization.Tensor Comprehensions chooses thread-block parameters through black-box optimization and generates internal loops polyhedrally.
- Statistical modeling: Both GBT- and TreeGRU-based statistical models converged faster and achieved better results than genetic-algorithm and random-search baselines on NVIDIA TITAN X.The comparison included versions of genetic algorithm and random search with two hardware evaluations per trial.
- Statistical modeling: Statistical cost models had previously been applied to SAT-solver optimization; this work adapts the idea to tensor programs and workload transfer.The proposed domain-specific model is intended to enable effective transfer among workloads.
6 Experiments
Experiments evaluate component choices, transfer learning, and end-to-end deployment across GPUs and CPUs. AutoTVM improves search and inference performance, including 2× to 10× transfer speedups and 1.2× to 3.8× end-to-end gains.
- 6.1 Component Evaluations: Both GBT and TreeGRU statistical models found operators 2× faster than random search, outperforming black-box methods.The authors attribute the result to domain-specific modeling.
- 6.1 Component Evaluations: Rank-based objectives were slightly better than regression-based objectives in most cases, so rank was selected as the default.The authors suggest rank may avoid modeling absolute cost values.
- 6.1 Component Evaluations: Diversity-aware exploration had no positive or negative impact on most workloads but improved C6.The strategy was retained because it can help, has no meaningful negative impact, and negligibly affects runtime.
- 6.1 Component Evaluations: Uncertainty-aware acquisition functions yielded no improvements in the evaluations.Uncertainty was estimated by training five bootstrapped models, using a regression objective.
- 6.2 Transfer Learning Evaluations: 2× to 10× speedup was obtained with transfer learning on target workloads C7, C8, and C9 compared with learning from scratch.Source data came from C1–C6, with experiments on NVIDIA TITAN X, ARM GPU, and ARM Cortex-A53.
- 6.2 Transfer Learning Evaluations: Transferable representations generalized across input shapes and operator types, while cross-domain usefulness depended on domain distance.A preliminary ARM Mali GPU to ARM Cortex-A53 study also showed transfer across devices.
- 6.3 End-to-End Evaluation: AutoTVM generated high-performance tensor programs across server GPU, embedded CPU, and mobile GPU back-ends without external operator libraries.Single-operator baselines included cuDNN, TensorFlow Lite, ARM Compute Library, and Tensor Comprehensions.
- 6.3 End-to-End Evaluation: 1.2× to 3.8× end-to-end performance improvement was measured across ResNet, MobileNet, LSTM, DQN, and DCGAN workloads.The improvements came from tensor-program optimization and operator fusion optimizations.
7 Discussion and Conclusion
AutoTVM automatically optimizes tensor-operator implementations with statistical cost models and model transfer. The authors report promise for deep-learning deployment and opportunities for broader optimization across hardware back-ends.
- Framework: AutoTVM automatically optimizes tensor-operator implementations using a machine-learning-based framework.Its statistical cost model supports sharing between workloads and speeds optimization through model transfer.
- Framework: Model transfer accelerates optimization by sharing statistical cost models between workloads.The conclusion identifies model transfer as a mechanism for speeding the optimization process.
- Implications: The positive experimental results show promise for deep-learning deployment.The conclusion also identifies learning tensor programs as a testbed for neural program modeling, Bayesian optimization, transfer learning, and reinforcement learning.
- Implications: Learning to optimize tensor programs can enable more fused operators, data layouts, and data types across diverse hardware back-ends.The authors characterize these capabilities as crucial to improving deep-learning systems.
A.1 Additional Experimental Results
Additional experiments examine single-operator performance and the effects of cost-model design choices on ResNet-18 convolution operators. The figures cover objective functions, diversity, and uncertainty-aware exploration.
- Additional results: Single-operator performance is evaluated on the ARM Mali T860MP4 mobile GPU.The figure isolates Mali T860MP4 performance.
- Additional results: The effectiveness of the cost model is evaluated across all convolution operators in ResNet-18.This evaluates cost-model behavior at the operator level within a complete network.
- Additional results: The objective function’s impact is evaluated across all ResNet-18 convolution operators.The compared objectives include GBT Rank, TreeGRU Rank, GBT Regression, and TreeGRU Regression.
- Additional results: Diversity-aware exploration is evaluated across all ResNet-18 convolution operators.The figure examines the impact of diversity-aware exploration on these operators.
- Additional results: Uncertainty-aware acquisition is evaluated across all ResNet-18 convolution operators.The figure presents acquisition-function variants including Expected Improvement, Upper Confidence Bound, and Mean.
A.2.1 Loop Context
Loop context features combine loop attributes with buffer access patterns, extracted for variables in the longest AST chain and paired with reuse and top-down relations.
- Loop context records loop attributes and access patterns for all touched inner buffers.
- The method selects the longest AST chain, extracts loop-variable features, and computes touch count versus reuse ratio and top-down relations.
- TreeGRU uses 128-dimensional loop-variable embeddings and a 128-unit hidden state, while planning batches for GBT and TreeGRU are both 64.