[Paper Review] Ansor: Generating High-Performance Tensor Programs for Deep Learning

Lianmin Zheng, Chengfan Jia+10 more|arXiv (Cornell University)|Jun 11, 2020

Parallel Computing and Optimization Techniques53 references82 citations

TL;DR

Ansor is a tensor program generation framework that automatically builds a large hierarchical search space, samples complete programs, and uses evolutionary search with a learned cost model to optimize DNN subgraphs across CPUs and GPUs, achieving substantial speedups over state-of-the-art methods.

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 imes$, $2.6 imes$, and $1.7 imes$, respectively.

Motivation & Objective

• Motivate the need for high-performance tensor programs across diverse hardware platforms for DNNs.
• Eliminate manual template engineering by automatically generating and exploring a large search space of tensor programs.
• Enable end-to-end optimization of DNN subgraphs through scheduling and resource-aware selection.
• Demonstrate that Ansor achieves superior performance compared with state-of-the-art approaches on CPU and GPU targets.

Proposed method

• Introduce a hierarchical search space that decouples high-level structures (sketches) from low-level details (annotations).
• Automatically generate sketches via a derivation-based enumeration over a DAG representing the computation graph.
• Randomly annotate sketches to obtain complete programs for evaluation and tuning.
• Apply evolutionary search with a learned cost model to fine-tune sampled programs.
• Use a gradient-descent based task scheduler to prioritize subgraphs that most improve end-to-end DNN performance.
• Leverage a cost model trained with gradient boosting trees to predict program throughput and guide search.

Experimental results

Research questions

• RQ1 Can Ansor automatically generate a large, comprehensive search space for tensor programs without manual templates?
• RQ2 Does evolutionary search with a learned cost model effectively identify high-performance programs across CPU and GPU backends?
• RQ3 Can prioritizing subgraphs via a scheduler improve end-to-end DNN performance compared with per-operator optimization?
• RQ4 How does Ansor's performance compare to state-of-the-art search-based frameworks on multiple hardware platforms?

Key findings

• Ansor improves DNN execution performance on Intel CPU, ARM CPU, and NVIDIA GPU by up to 3.8×, 2.6×, and 1.7× respectively.
• Most best-performing programs found by Ansor lie outside the search space of existing approaches.
• Ansor searches more efficiently and finds higher-performance programs in shorter time despite a larger search space.
• Ansor can match state-of-the-art performance with an order of magnitude less search time.
• Ansor enables automatic extension to new operators by requiring only mathematical definitions, not manual templates.