Ansor: Generating High-Performance Tensor Programs for Deep Learning

Abstract

High-performance tensor programs are crucial to guarantee efficient execution of deep neural networks. However, obtaining performant tensor programs for different operators on various hardware platforms is notoriously challenging. Currently, deep learning systems rely on vendor-provided kernel libraries or various search strategies to get performant tensor programs. These approaches either require significant engineering effort to develop platform-specific optimization code or fall short of finding high-performance programs due to restricted search space and ineffective exploration strategy. We present Ansor, a tensor program generation framework for deep learning applications. Compared with existing search strategies, Ansor explores many more optimization combinations by sampling programs from a hierarchical representation of the search space. Ansor then fine-tunes the sampled programs with evolutionary search and a learned cost model to identify the best programs. Ansor can find high-performance programs that are outside the search space of existing state-of-the-art approaches. In addition, Ansor utilizes a task scheduler to simultaneously optimize multiple subgraphs in deep neural networks. We show that Ansor improves the execution performance of deep neural networks relative to the state-of-the-art on the Intel CPU, ARM CPU, and NVIDIA GPU by up to $3.8\times$, $2.6\times$, and $1.7\times$, respectively.

OSDI 2020

References in corpus (13)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• cuDNN: Efficient Primitives for Deep Learning
• The State of Sparsity in Deep Neural Networks
• Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions
• Learning to Optimize Tensor Programs
• Deep Learning with Dynamic Computation Graphs
• Optimizing CNN Model Inference on CPUs
• AutoPhase: Juggling HLS Phase Orderings in Random Forests with Deep Reinforcement Learning
• FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems
• ProTuner: Tuning Programs with Monte Carlo Tree Search

Cited by in corpus (10)

• Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference
• FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads
• Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
• CDMPP: A Device-Model Agnostic Framework for Latency Prediction of Tensor Programs
• Integration of a systolic array based hardware accelerator into a DNN operator auto-tuning framework
• IOS: Inter-Operator Scheduler for CNN Acceleration
• Efficient distributed algorithms for Convolutional Neural Networks
• AI Powered Compiler Techniques for DL Code Optimization
• Transfer Learning Across Heterogeneous Features For Efficient Tensor Program Generation
• OpEvo: An Evolutionary Method for Tensor Operator Optimization