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.
翻译:高性能张量程序对于确保深度神经网络的高效执行至关重要。然而,在不同硬件平台上为各种算子获取高性能张量程序极具挑战性。目前,深度学习系统依赖供应商提供的内核库或多种搜索策略来获得高性能张量程序。这些方法要么需要大量工程工作来开发特定平台的优化代码,要么因搜索空间受限和探索策略低效而难以找到高性能程序。我们提出Ansor——一个面向深度学习应用的张量程序生成框架。与现有搜索策略相比,Ansor通过从搜索空间的层次化表示中采样程序,探索了更多的优化组合。随后,Ansor利用进化搜索和学习到的代价模型对采样程序进行微调,以识别最优程序。Ansor能够找到超出当前最先进方法搜索空间的高性能程序。此外,Ansor还采用任务调度器同时优化深度神经网络中的多个子图。实验表明,相对于现有最优方法,Ansor在Intel CPU、ARM CPU和NVIDIA GPU上分别将深度神经网络的执行性能提升了高达$3.8\times$、$2.6\times$和$1.7\times$。