Learning to Optimize Tensor Programs

By Tianqi Chen, Lianmin Zheng, Eddie Q. Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy. Neural Information Processing Systems, 2018

501 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Hardware and Architecture, Statistical and Nonlinear Physics

View Paper: https://www.semanticscholar.org/paper/cb91c2f8d3cac0b655a39be318b603334eb18987

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Song Han
2. Ji Lin
3. Chuang Gan
4. Yujun Lin
5. Matei A. Zaharia
6. Joseph E. Gonzalez
7. Yuandong Tian
8. A. Aiken

In Hardware and Architecture

1. Matei A. Zaharia
2. A. Aiken
3. Xinyun Chen
4. Yuandong Tian
5. Oded Padon
6. Koushik Sen
7. Song Han
8. Chuang Gan

In Statistical and Nonlinear Physics

1. Song Han
2. Chuang Gan
3. Lianmin Zheng
4. Yuandong Tian
5. Ji Lin
6. Joseph E. Gonzalez
7. Simon Shaolei Du
8. Jason D. Lee

Citations in the corpus

• TASO: optimizing deep learning computation with automatic generation of graph substitutions (2019, methodology): “TensorFlow, TensorFlow XLA, TensorRT, and MetaFlow use the highly-engineered cuDNN and cuBLAS libraries [10, 12] to perform DNN operators on GPUs, while TVM generates customized GPU kernels for the DNN operators.”
• Learning to Perform Local Rewriting for Combinatorial Optimization (2018, methodology): “Our approach is tested on multiple domains with extensive ablation studies, and could also be extended to other closely related tasks such as code optimization [41, 12], theorem proving [25, 30, 4, 24], text simplification [14, 37, 18], and classical combinatorial optimization problems beyond routin”
• MCUNet: Tiny Deep Learning on IoT Devices (2020, methodology): “…the neural network using neural architecture search based on a given deep learning library ( e.g ., TensorFlow, PyTorch) [42, 6, 45], or (b) tunes the library to maximize the inference speed for a given network [9, 10], MCUNet can better utilize the resources by system-algorithm co-design.”
• Ansor : Generating High-Performance Tensor Programs for Deep Learning (2020, methodology): “We adopt a learned cost model similar to related works [2, 12] with newly designed program features.”
• The Deep Learning Compiler: A Comprehensive Survey (2020, methodology): “ing to optimize performance by tuning the available mapping options. TC can only perform auto-tuning on NVIDIA GPUs, while TVM can apply auto-tuning to CPUs (x86 and ARM), mobile GPUs and NVIDIA GPUs [30]. TVM and TC use different tuning approaches: TC uses genetic search [41] and TVM uses two machi”