[论文解读] 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 Techniques参考文献 53被引用 82

一句话总结

Ansor 是一个张量程序生成框架，能够自动构建大规模层次化搜索空间，采样完整程序，并使用带学习成本模型的进化搜索，在 CPU 与 GPU 上优化 DNN 子图，相比现有方法实现显著的加速。

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.

研究动机与目标

• 为在多种硬件平台上实现 DNN 的高性能张量程序的需求提供动机。
• 通过自动生成和探索大规模张量程序搜索空间，消除手工模板设计。
• 通过调度和资源感知选择实现 DNN 子图的端到端优化。
• 证明 Ansor 在 CPU 和 GPU 目标上相比最新方法具有更优性能。

提出的方法

• 引入一个分层搜索空间，将高级结构（sketches）与低级细节（annotations）解耦。
• 通过对表示计算图的有向无环图（DAG）进行基于推导的枚举，自动生成 sketches。
• 对 sketches 随机标注以获得用于评估和调优的完整程序。
• 应用带学习成本模型的进化搜索来微调采样的程序。
• 使用基于梯度下降的任务调度器来优先考虑对端到端 DNN 性能提升最大的子图。
• 利用用梯度提升树训练的成本模型来预测程序吞吐量并引导搜索。

实验结果

研究问题

• RQ1 Ansor 是否能够在没有手动模板的情况下，自动生成一个大规模、全面的张量程序搜索空间？
• RQ2 带学习成本模型的进化搜索是否能够在 CPU 和 GPU 后端有效识别高性能程序？
• RQ3 通过调度器优先考虑子图是否能相较于逐算子优化提高端到端 DNN 性能？
• RQ4 在多种硬件平台上，Ansor 的性能与最先进的基于搜索的框架相比如何？

主要发现

• Ansor 在 Intel CPU、ARM CPU 和 NVIDIA GPU 上的 DNN 执行性能分别提高至多 3.8×、2.6×、1.7×。
• 大多数最优程序超出现有方法的搜索空间。
• 尽管搜索空间更大，Ansor 仍以更高的搜索效率在更短时间内找到更高性能的程序。
• Ansor 可以在少一个数量级的搜索时间内达到与最先进方法相当的性能。
• Ansor 通过仅需要数学定义，而非手动模板，即可实现对新算子的自动扩展。