[Paper Review] Learning to Optimize Tensor Programs

Tianqi Chen, Lianmin Zheng+6 more|arXiv (Cornell University)|May 21, 2018

Parallel Computing and Optimization Techniques36 references124 citations

TL;DR

The paper presents AutoTVM, a learning-based framework that builds domain-specific cost models to automatically search and optimize tensor program implementations for deep learning workloads, achieving competitive performance across CPUs, GPUs, and mobile hardware through transfer learning.

ABSTRACT

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.

Motivation & Objective

• Motivate automatic optimization of tensor operators beyond hand-tuned libraries.
• Formalize the problem of exploring a large space of tensor program schedules to minimize hardware run time.
• Develop a learning-based cost model to guide search over billions of program variants.
• Enable transfer learning to reuse knowledge across workloads and hardware targets.

Proposed method

• Define the schedule search space Se for a given index expression e and a code generator g, with cost f(x) unknown but measurable.
• Develop two cost models: gradient-boosted trees (GBT) using domain features from the low-level AST, and a TreeGRU embedding of the AST.
• Train models with a regression or rank-based objective to predict run time and guide selection.
• Use an exploration loop with simulated annealing and diversity-aware selection to generate candidate schedules and measure their actual performance on hardware.
• Incorporate transfer learning by combining a global invariant representation model with a local in-domain model to speed up optimization for new workloads.
• Report that AutoTVM (TVM-based) achieves end-to-end gains across back-ends with no external operator libraries.

Experimental results

Research questions

• RQ1 Can learning-based cost models effectively navigate the billions-scale schedule space to find fast tensor program implementations on real hardware?
• RQ2 Does transfer learning enable faster optimization across different workloads and hardware targets?
• RQ3 What objective (regression vs ranking) best supports finding fast schedules in this domain?
• RQ4 Is diversity-aware exploration and uncertainty estimation beneficial in this setting?

Key findings

• Statistical cost models (GBT and TreeGRU) outperform black-box baselines in locating faster tensor programs within fewer hardware evaluations.
• Rank-based objectives generally perform as well or better than regression objectives for ordering candidate programs by speed.
• Diversity-aware exploration showed negligible to modest impact, with some gains in specific workloads.
• Transfer learning yields 2x to 10x speedups by reusing knowledge from previously seen workloads across GPUs and mobile devices.
• End-to-end inference improvements with AutoTVM range from 1.2x to 3.8x across NVIDIA TITAN X, ARM Cortex-A53, and ARM Mali back-ends.