(三)大话深度学习编译器中的自动调优·Empirical Search

前面的第一篇"(一)大话深度学习编译器中的自动调优·前言"与第二篇"(二)大话深度学习编译器中的自动调优·DSL与IR"分别介绍了背景与一些相关概念,这第三篇我们开始切入正题,看下现代深度学习编译器中的自动调优(Auto-tuning)方法。Schedule的自动生成,一类方法是基于解析模型(Analytical model),然后使用一些经验公式来产生解;另一类方法是多面体编译技术。它将循环嵌套迭代空间建模为多面体,然后用如整数规划等数学方法求出能提升局部性与并行性的循环变换;还有一类就是经验搜索(Empirical search)方法。这类方法一般采用迭代的方法,通过一定策略产生很多的变体,然后对它们进行性能评估,最终挑选出性能最优的schedule配置(Configuration)。这个过程还可以结合机器学习方法,比如强化学习,让机器学习模型学习schedule参数的选取策略。

在基于empirical search的自动调优方向上,由于Halide与TVM发展时间较长,基于它们的方法比较繁荣。接下来就先以它们为主线展开。我们从一系列相关论文看下它们的大体发展脉络。

Halide系

  • 2013年MIT与Adobe的论文《Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines》是早期对于Halide比较系统性的介绍。其中就提到了auto-tuning优化pipeline schedule的方法。它采用stochastic search(具体地,genetic algorithm)来找近似最优解。该方法受PetaBricks的启发。PetaBricks是2009年论文《PetaBricks: a language and compiler for algorithmic choice》提出的用于算法选取的语言与编译器。Halide的auto-tuning方法中起始点可以通过领域专家知识来获得。先验知识可以通过mutation rule来体现。另外,搜索过程中它还会做正确性验证。其缺点是搜索时间较长。针对不同复杂度的情况,tuning的时间在小时到天不等。

  • 2014年MIT论文《OpenTuner: An extensible framework for program autotuning》中的OpenTuner是一个用于构建领域专用程序autotuner的框架。优化可以是多目标(如最小化耗时,最大化准确性等)的。它是一个autotuner的元框架。以往的autotuner往往为特定项目专用,缺少移植性。相比之下,OpenTuner聚焦于通用性,将Search与Measurement分离,中间通过Result dataset(可以是SQL数据库)联接。搜索方法上,它包含了Nelder-mead search,Torczon hillclimber, PSO等多种方法的实现,同时还可以通过AUC Bandit meta 技术同时结合多种搜索方法。实验中它被用于7个不同的项目(GCC/G++, Halide, HPL, PetaBricks, Stencil, Unitary, Mario),以验证其通用性。

  • 2015年MIT的《HalideTuner : generating and tuning halide schedules with Opentuner》
    中的Halidetuner使得Halide用户可以对schedule进行自动调优。它基于OpenTuner构建。前面提到过,OpenTuner是一个较为通用的搜索框架。它基于以往的搜索结果以及运行时找更好的配置。但要将它应用到具体问题,需要定义相应的搜索空间。HalideTuner起的作用就是将Halide schedule搜索问题按OpenTuner要求的进行描述和编码。它将schedule space定义成100个transformation的序列。在搜索的每一轮中,这些transformation应用到base-schedule上,然后schedule verifier用于验证transformation或者schedule的合法性。通过验证的话,经过编译,运行,将运行时信息给到OpenTuner进行学习。该框架将搜索时间降到小时级。但对于复杂的pipeline无法收敛到好的解(比手工实现慢很多倍)。

  • 2016年CMU,Google与Stanford的论文《Automatically Scheduling Halide Image Processing Pipelines》引入了analytical heuristic-based model,扩展了Halide编译器中已有的function bounds analysis进行全局程序变换优化局部性与并行性,适用于CPU和GPU平台。该算法不需要耗时的auto-tuning过程,因此只需要秒级来产生schedule。它使用overlapping tile analysis和greedy grouping/merge算法,在设计空间中快速探索。它使用greedy grouping算法将pipeline中的stages进行分组,然后独立地考虑每个组,通过估计每个stage的arithmetic cost为每个组找最高效的tile。tile size的调节的目标是最小化data load的数量。搜索空间中的tile size限于2的幂(8到256),另外分析不涉及buffer allocation与storage scheduling。因此,它虽然可以在秒级得到schedule,但可能会忽略design space中一些利用到sliding window相关更优解。

  • 2018年MIT,Facebook与UC Berkeley,Google的论文《Differentiable programming for image processing and deep learning in Halide》扩展了Halide,在x86 CPU与NVIDIA GPU平台上,使其能够计算任意Halide程序的梯度(基于reverse-mode automatic differentiation)。使用者只需给算子forward的实现,便可以自动生成backward的实现。为了提升生成算子的性能,它使用了automatic scheduling技术。Automatic scheduling针对当时Halide中对large reduction优化不够,及不支持GPU平台的问题,做了相应的改进。

  • 2019年Eindhoven University of Technology的《Schedule Synthesis for Halide Pipelines through Reuse Analysis》关注通过挖掘图像处理pipeline中的reuse机会来最大化producer-consumer间的访存locality,平均性能超过Halide Auto-Scheduler与PolyMage DSL。它指出当时的SOTA,即Halide Auto-Scheduler只考虑了一小部分的design space(即compute与store设为loop nest的同一level),即限于fully inlined, fully stored等computation与storage在innermost inter-tile level上的情况(overlapping tiles)。本文方法通过分离stage中的computation与storage,可以考虑更多优化机会,如sliding window。该方法主要考虑stage fusion与tile size的选取。在tile size的选取中,通过考虑应用与架构相关的参数的analytical model,该方法可以在秒级输出schedule。

  • 2019年Facebook, UC Berkeley, MIT, Adobe, Google, Stanford的《Learning to Optimize Halide with Tree Search and Random Programs》引入一种于图像处理与深度学习的自动调度算法(扩展到多核CPU)。它采用了beam-search方法(它是greedy方法与DFS的折中,优点是复杂度是线性的)的变体,即基于beam search的backtracking tree search算法,并采用coarse-to-fine的两次pass方式来考虑更多样化的变体。搜索过程中,它使用了基于manually-derived feature和机器学习的cost model来预测intermediate schedule的执行时间,从而对搜索过程中产生的状态进行排序。这个cost model是用对上千个随机生成的Halide program与schedule进行benchmarking来训练的。搜索过程中还会利用一些先验经验进行pruning,如所有多核并行发生在outermost loop level,并行循环的迭代次数在核数与核数的16倍之间等等。另外,由于cost model无法完美地准确预测性能,beam search也无法保证找到最优解,如果使用importance sampling与benchmarking(即autotuning),还可进一步提升性能,但这会耗费更长的时间。在调度空间的设计上,它将schedule表示成specialized loop nest。基于这种表示,从最后的stage开始,每次为一个stage枚举可能的scheduling choice(tiling, compute & storage granularity和annotation)。

  • 2020年Eindhoven University of Technology的论文《Schedule Synthesis for Halide Pipelines on GPUs》主要研究将Halide autoscheduler机制扩展到GPU平台(以往主要关注多核CPU架构)。性能比手工schedule平均快10%,比之前的autoscheduling快2x。与CPU不同,GPU平台对于schedule有一些严格的限制。这篇文章扩展了Halide master中的autoscheduler,添加了一个考虑GPU参数的analytical cost model,和一个analysis pass用于GPU schedule的生成。一方面考虑很大的tiling和kernel fusion空间,另一方面满足硬件约束。优化整个pipeline涉及两个粒度:inter-group的scheduling关注将pipeline切分成stage组,以及将stage内联到它的consumer中。intra-group的scheduling涉及tiling, unrolling,threads/blocks绑定,以及确定producer的计算放到consuming loop的哪层嵌套等等。流程上,scheduler先将简单的(如pointwise consumed,或者low arithmetic cost)stage内联入它的consumer,然后使用greedy algorithm,通过stage间的tiling与fusion,将pipeline切分成group。这一步中,scheduler先确定要tile哪一维,然后通过一些先验公式确定tile size的范围。接着,通过fusion analysis递归地尝试将group进行合并,直到没有对性能有益(即合并后的总cost小于单独group的cost之和)的合并存在。cost考虑memory与arithmetic两方面,在group cost analysis中计算得到。对于每个group,基于前面选择的 tile size将循环做tiling,然后将outer(最多三层)inter-tile变量绑定到block,而将outer intra-tile变量绑定到thread。最里边的intra-tile loop会被展开。另外,该scheduler还支持nested fusion。

TVM系

TVM早期受到Halide的启发,也是基于编译器的计算框架,主要用于深度学习。经过几年的快速发展,它在功能与性能上都已比较成熟,且有着众多的衍生项目。TVM中的auto-tuning机制经过了几代的发展:

  • AutoTVM:在2018年OSDI上的论文《TVM: An Automated End-to-End Optimizing Compiler for Deep Learning》与2018年的论文《Learning to Optimize Tensor Programs》中有相关的介绍。该方法需要使用者为算子提供schedule template,即schedule的模板,其中留出一些可调的参数(如tiling factor),然后通过搜索方法(模拟退火)进行自动调参。这些可调参数称为knob,它们的赋值组合称为configuration。为了加速,它使用cost model来指导搜索过程。Cost model从low-level的abstract syntax tree中抽取特征,然后通过机器学习模型(Gradient boosted tree或者TreeGRU)进行学习。另外,像ProTuner,FlexTensor等项目也是使用了这种基于template的方式定义搜索空间。
  • Ansor:在2020年论文《Ansor: Generating High-Performance Tensor Programs for Deep Learning》与官方介绍Introducing TVM Auto-scheduler (a.k.a. Ansor)中有相关介绍。第一代的AutoTVM中一个不太易用的地方是需要开发者写schedule template。Ansor的优点之一是不用写schedule template,而是通过一些规则来生成schedule。然后通过evolutionary algorithm进行搜索。另外,它的改进之处还在于通过层次化的搜索设计扩大了搜索空间,另外引入了task scheduler可以将tuning时间在整个网络的子图中进行更好地分配。之前在学习过程中也写了一些笔记可供参考:《TVM中的auto-scheduling机制(Ansor)学习笔记》。
  • Meta Schedule:可参见社区的介绍[RFC] Meta Schedule (AutoTensorIR)。它指出之前方式的缺点,如手工方式费时费力,AutoTVM需要写schedule template。Ansor不易扩展到新的schedule primitives(如tensorize,software pipelining)。Meta schedule提供了实现以上几种机制的统一API。它将search rule视为composite schedule,然后将composite schedule应用到TensorIR的每个block上,从而定义design space。

有了Auto-tuning机制,通过编译生成的算子性能有了质的提升,但还有一些值得改进的地方,比如:

  • Tuning时间过长。
  • 对特殊计算硬件(如Tensor Core)的支持。
  • 对dynamic shape的支持。

Tuning加速

Tuning时间过长一直以来都是比较突出的问题。由于搜索空间巨大,为了找到更优的schedule,需要在搜索过程中评估大量的变体,而编译运行这些变体的时间会成为tuning中耗时的大头。尽管这是一次性的开销,但还是会影响部署的复杂度与效率。对于该问题业界有不少尝试,比如以下方向:

  • 预tuning查询:比较intuitive的想法就是既然tuning花时间,那就在一些常见配置先tuning好,然后把结果存起来,要用时直接查询。如Upstream版本中实现了tophub机制。它提供了一些case下的schedule供查询。Lorien(https://pypi.org/project/lorien/)是一个构建于TVM上用于搜索schedule的系统,它通过更大规模的数据库来更多更全面的schedule。

  • 并行化:Lorien中使用了基于master-worker模式的分布式的tuner。PGRCA(Parallel Genetic RAFT Classification Algorithm)基于并行化的遗传算法。它将搜索空间分为几块,每块可以独立演进。最后优中选优。TVM Server Instance(TSI)Scaling基于tuning时用少量GPU资源能得到鲁棒性更好的schedule的观察,利用和扩展了NVIDIA MPS,使GPU多路复用进行tuning提高资源利用率。

  • 时间分配:另一个角度是优化tuning过程中的time allocation。2020年Microsoft的论文《AdaTune: Adaptive Tensor Program Compilation Made Efficient》基于同模型和硬件上variance差异较大的观察,提出考虑variance来自动调整性能测量次数,另外在surrogate model中考虑方差,并基于此动态均衡exploration和exploitation。2021年Microsoft的论文《DynaTune: Dynamic Tensor Program Optimization in Deep Neural Network Compilation》主要针对time allocation问题。即如何充分利用编译的时间,时间花可能大提升性能的地方,从而提升整个模型的tuning收敛速度。分析了DL编译器几点挑战:1. 现有方法关注单算子收敛速度(非整模型)。2. 静态调度不够。3. 即使用动态信息,难以extrapolate estimated performance。针对这些问题,DynaTune使用MAB(Multi-Armed Bandit)模型,目标是设计scheduler使cumulative regret最小。UCB(upper confidence bound)用来处理time-slot-based optimization中的决策问题。模型中需要知道maximum latency reduction的算子,而它实际中无法得到。因此文中设计了Bayesian belief model(+MCMC)来预测每个算子的潜在的performance gain,其中的uncertainty信息可以用于指导搜索。实验中它与static schedule(Random, Round-robin与Linear)和dynamic scheme(Dynamic allocation + Random selection, Dynamic allocation + Round-robin selection)做了比较。

  • Cost model:Cost model对于搜索效率有着至关重要的影响。像《Learning to Schedule Halide Pipelines for the GPU》中提到的one-shot模式和top 5模式。在这些模式下不做编译与性能profiling,而是直接根据cost model进行搜索得到schedule。像这类方法就更依赖cost model的准确性。业界这几年有很多工作是针对cost model以及性能评估方法的改进。

    • Lorien中的perf model使用了Gated Neural Network + Normalized Discounted Cumulative Gain进行训练。Tuning时间可以从小时级减少到秒级,性能相比最好的tuned schedule达到其95%。但是需要较多的训练数据。
    • Graph-Level Scheduling Optimization with Polyhedral Analysis for Tensor Programs中使用基于polyhedral分析的analytical模型替代ML-based cost model。实验任务中,在CPU上,相比AutoTVM中的XGBoot tuner,trials数量降低一个量级,且tuning出来的程序性能很接近。
    • 2020年HKUST和Alibaba的论文《A History-Based Auto-Tuning Framework for Fast and High-Performance DNN Design on GPU》提出一种tuning框架可以有效利用以前的tuning历史来探索schedule空间。它维护历史数据库。对于给定输入层,根据相似度判断历史信息是否能重用。如果能,就可以避免从头训练perf model。相比TVM中的XGBoost tuner和GA tuner在搜索时间上有倍数的减少。
    • RAFT(rootfline and fast autotune)(RAFT: Accelerating the Tuning Process for AutoTVM)用于预测性能。它结合roofline模型,从理论上对随机选取的配置进行排序。另外用random forest来分类config是否可执行(因为AutoTVM中很多生成的程序无法执行)。搜索时间在不同GPU上平均有倍数的加速,另外在模型的推理性能上也有提升。
    • 2022年Beihang University的论文《FamilySeer: Towards Optimized Tensor Codes by Exploiting Computation Subgraph Similarity》是一个基于TVM的auto-tuning框架。文中指出现有方法的问题是training sample与time budget无法很好地利用。FamilySeer将subgraph聚类成family,同类中的subgraph可以共享training sample与time budget。Ansor中的cost model对所有的subgraph是一个,忽略了subgraph的不同特性。它对Ansor进行了改进,基本思想是挖掘subgraph间的相似性,将subgraph合并成family,对每个subgraph family构建cost model。这样一个subgraph的tuning也能用于同一faimily的另外subgraph。另外,还实现了cost model的并行训练,以及GPU上的cost measurement。与Ansor相比,FamilySeer在搜索效率上可以有平均约2-3x的性能提升,同时达到相同性能。
    • 2021年UC Berkeley, OctoML与CMU的论文《TenSet: A Large-scale Program Performance Dataset for Learned Tensor Compilers》也是聚焦于cost model。Ansor中的cost model是以on-the-fly的方式在搜索的过程中进行数据的收集与训练的。文中的TenSet通过一个大型的数据集以pre-training的方式构建cost model。该数据集包含来自Intel CPU,ARM CPU,NVIDIA GPU等六个平台的五千多万条记录。将该cost model整合进Ansor,能够在达到相关质量的情况下减少搜索时间10x。
    • 2021年POSTECH的论文《MetaTune: Meta-Learning Based Cost Model for Fast and Efficient Auto-tuning Frameworks》主要从cost model入手。它先对程序表示的AST扩展(成Super-graph argmentation),然后引入graph convolutional network(GCN),使用model-agnostic meta-learning(MAML)方法进行训练。搜索方法实现了Batch Bayesian Optimization与RL。实验表明meta-learning based cost model比machine-learning based model(TVM中的transfer learning based)推理性能平均有8~13%的提升。
    • 2019年的Qualcomm的论文《Simulating Execution Time of Tensor Programs using Graph Neural Networks》基于TVM,提出学习surrogate model来克服搜索configuration space耗时的问题。模型基于AST进行训练,使用graph convolutional network(GraphNN)来挖掘graph中的结构信息。它的优势是使用可学习的基于图的处理比基于heuristic的特征提取有优势。AST中的每个节点会通过一个shared encoder编码成一个固定长的特征,节点间的信息通过GraphNN来传递。这些信息会整合进一个固定长的向量,最后通过一个预测函数对性能进行预测。
    • 2021年Amazon的论文《Tuna: A Static Analysis Approach to Optimizing Deep Neural Networks》使用张量计算性能的静态分析(编译时)方法,基于analytical cost model预测tensor program性能。它基于TVM实现(重用了TE, IR与codegen)。静态分析的优点是可以在多核CPU机器上运行。它基于目标硬件特点构建特征。Tuna既分析program IR也分析low-level生成代码(assembly codes),来提取目标硬件相关的特征(如SIMD指令数量、CPU cache局部性、GPU shared memory使用等)。如GPU cost model会考虑PTX指令数量,thread level parallelism, workload per thread, SM occupancy, warp latency hiding, shared memory bank conflict。Program的performance score这些特征的线性模型(其参数是per-hardware的,通过profiling确定)。根据这些特征就可以预测tensor program上一组transformation的相对性能。搜索方法使用了evolution strategy(ES)。实验中的metrics包括compilation time/cost与inference latency。与AutoTVM相比,它可以只用1.65%的编译时间达到相当的性能。
    • 2022年POSTECH(Pohang University of Science and Technology)的论文《One-Shot Tuner for Deep Learning Compilers》从NAS中启发,使用一些随机产生的code sample预训练基于Transformer的predictor model,避免了传统auto-tuning中重复的search与hardware measurement的开销。为了克服预训练模型带来的挑战,该方法 1) 将显式的task-specific特征(op type, input/output shape, filter size)纳入考虑,因为以往方法没有考虑或者隐式地考虑这些信息,导致模型难以跨任务。2) 采用了Prior-Guided Task Sampling(PGS)挖掘现有模型的分布信息,从而提高sampling efficiency。3) 采用了Exploration-Based Code Sampling(EBS)机制使predictor model更聚焦于高性能的代码,减少花在不合法或者低性能代码上的hardware measurement时间。最后,在搜索过程中的top-K候选会在放到硬件上跑,以确定最优的版本。实现基于TVM,替换了其中的XGB模型。实验中与Auto-TVM,Chameleon,AdaTune,Ansor比较,在编译时间上有数量级上的减少,在推理时间上也有提升。
  • 其它还有一些搜索方法上的改进与尝试也值得关注:

    • 2021年UIUC, NVIDIA与Facebook的论文《Mind Mappings: Enabling Efficient Algorithm-Accelerator Mapping Space Search》提出一种基于可微分的方法解决algorithm-accelerator的mapping space search问题(即从问题,或称算法的参数化实例,到硬件的optimal mapping),比如每个数据结构分配多少buffer,计算任务到processing elements的对应关系等等。该问题难点在于搜索空间通常很大,而且非凸不光滑。因此很多工作使用black-box optimization,如SA, GA等方法来解决。Mind Mapping的基本思想是用一个可微的surrogate model来近似一个不可微的accelerator cost function。这种做的最大好处是可以使用gradient-based的算法,使得搜索更为高效。整个过程分两个阶段:第一个阶段通过一些训练数据训练一个基于MLP的surrogate model。这个阶段中将网络中的weight作为变量,基于problem与mapping预测其cost。第二个阶段通过gradient search找到高质量的mapping。这个阶段将第一个阶段中学习到的weight作为常量,以一个随机的初始mapping为起点,基于梯度更新该mapping。
    • 2020年Berkeley与MIT的论文《ProTuner: Tuning Programs with Monte Carlo Tree Search》中的ProTuner基于Halide通过MCTS(Monte Carlo Tree Search)来解决深度学习与图像处理中的高性能算子的调优问题。它主要对2019年论文中的《Learning to optimize halide with tree search and random programs》中的基于beam search的算法做了改进。它基于learned cost model的beam search的方法来找schedule。它的主要问题一是本质上是基于贪心的思想,即每步只根据当前已有信息做判断。另外cost model无法准确预测partially scheduled program的性能。为了解决这些问题,本文将scheduling问题建模成MDP,然后用MTCS求解。该方法有几个优点:1. 具有理论上的收敛性。2. UCB(upper confidence bound)可以平衡exploration和exploitation。3. 决策时可以“往前看”,避免贪心算法的局限。4. 基于完全scheduled program的cost做决策。完整schedule的cost通过一个基于随机program(完整scheduled)训练的cost model给出。同时它允许我们结合真实硬件上执行时间测量与cost model进行预测。
    • 2021年Netherlands eScience Center的论文《Bayesian Optimization for auto-tuning GPU kernels》采用Bayesian optimization(BO)用于GPU平台上的kernel auto-tuning。文中引入了一种contextual variance exploration factor与一种新的acquisition function,结合acquisition function选择机制来克服搜索空间离散且复杂,以及包含invalid configurations等挑战。BO包含几个基本组件:objective function,surrogate model与acquisition function。由于objective function未知且评估通常很费时,Surrogate model用于模拟objective function,一般求值没有objective function那么费时。本文方法使用Gaussian processes(GP)作为surrogate model。GP中的covariance function中的kernel使用Matérn kernel。Acquisition function在surrogate model上优化,给出下一次要在objective function上评估的parameter configuration。它需要考虑exploration与exploitation间的trade-off,原则上它会选取未知的区域或是有性能提升潜力的配置。Acquisition function通常有Probability of Improvement, Expected Improvement和Upper Confidence Bound。Acquisition function中的exploration factor一般置为常数,文中采用contextual variance在每一次评估中根据surrogate model的状态设置该参数。文中提出multi与advanced multi两种acquisition function,一开始,会有多个acquisition functions,在搜索过程中当给出重复建议时会将其忽略。这样就可以根据特定问题选取最优的acquisition function。实验中与Kernel Tuner框架中其它的方法(Simulated Annealing, Multi-start Local Search, Genetic Algorithm),以及其它的BO框架(BayesianOptimization,Scikit-optimize)作了比较。在OpenCL GEMM,CUDA 2D Convolution,及异构point-in-polygon三个GPU kernel上,本中方法有更优的表现。

功能扩展

关于功能上的扩展,业界也有一些很有意思的工作,如:

  • Dynamic shape的支持:2022年University of Toronto、AWS等机构的论文《DietCode: Automatic Optimization for Dynamic Tensor Programs》基于TVM中的Ansor,TVM中现有的自动调优机制Ansor要求workload是static的。对于dynamic-shape workload则会导致tuning时间很长(对每个shape都tuning一把)。一些用于dynamic shape支持的扩展面临不够自动或者产生的kernel性能较差的问题。如Selective tuning需要专家经验。Nimble针对large shape进行tuning,然后将得到的schedule应用到所有的shape上。但在large shape上最优的schedule未必对其它也是最优的(因为padding等原因)。Bucketing会引入padding开销与冗余计算。DietCode基于完整程序是由micro-kernels组成的这个观察,为dynamic shape构建由micro-kernels组成的shape-generic搜索空间,然后对于所有shape在这个统一的空间中进行联合搜索。相应地,文中构建了micro-kernel-based cost model。它分为两个部分:一部分是shape-generic的cost function,预测micro-kernel的cost;另一部分对应shape-dependent的adaption cost function(这部分可通过occupancy与padding ratio计算,不需要feature extraction)。当调优结束,会输出一组micro-kernels。为了将所有可能的shape dispatch到这些micro-kernels,Automatic dispatching机制基于cost model针对特定shape对每个micro-kernel进行投票,然后训练decision tree自动地将选取同一micro-kernel的shape归为一类,并生成源码用于运行时shape到micro-kernel的分派。对于dynamic shape的workload它比现有的auto-scheduling机制tuning时间减少数倍,同时性能上也比现有auto-scheduling与vendor library有所提升。
  • 对于特殊硬件的支持:2021年ByteDance、Rice University和NVIDIA的论文《Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance》提出了Bolt。它提到像Ansor将硬件看作黑盒构建device model。这种做法缺乏硬件信息,因此在non-FP32上无法与cuBLAS,cuCUDNN的性能相比,因为无法有效利用TensorCores。Bolt基于目前硬件厂商越来越多提供模块化和可配置的库(如CUTLASS)的趋势。它利用轻量的perf profiler通过高效的模板搜索生成tensor program。通过合成与扩展模板原语,它使能了更多计算图级别的优化,如通过persistent kernel fusion实现GEMM/Conv间的融合。还有operator级优化 (如automated padding)。相比Ansor,它在一些CNN上平均达到2.5x的推理加速。同时,搜索时间缩短到20分钟左右。

Reinforcement Learning

这几年开始流行用机器学习的方法来解决程序的优化问题,算子的自动调优也不例外。基于学习的方法的基本思想是训练一个机器学习模型从计算任务得到其schedule参数,其好处是不用每次都重头搜索。在基于机器学习的自动调优方法中,常用的是强化学习。强化学习是机器学习的一种,在与环境的交互中学习策略,就像动物的学习方式。AlphaGo的主要支撑技术就是强化学习,因为下棋就是典型的策略问题,可以在不断的对弈中提升策略。它会训练一个策略模型,训练好后,基于给定的计算,输出schedule相关参数。

  • 2019年UFPR的《Optimization of Halide Image Processing Schedules with Reinforcement Learning》使用强化学习(具体地,PPO算法)用于Halide中的schedule自动调优。Halide中的directive对应强化学习中的scheduling option,即action。当前已选取的directive集合定义为state。Reward是一个标量,正值代表执行时间减少,负值代表编译或执行出错,否则为0。但其scheduling options仍需由开发者从要优化的pipeline中抽取,因此整个过程是半自动的。

  • 2019年University of California的论文《Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation》采用强化学习并结合adaptive sampling的思想,提出了RELEASE(Reinforcement Learning Compiler with Adaptive Sampling for Efficiency)。假设已经有个搜索空间,为了使得auto-tuning高效,一是需要高效的搜索算法,二是减少耗时的hardware measurement的次数 。针对这两点,RELEASE中分别设计了RL-based Search Agent与Adaptive Sampling Module。前者中,RELEASE使用了强化学习(具体地,PPO算法)达到exploration与exploitation间更好的权衡,从而让收敛速度更快。后者中 ,它基于Adaptive Sampling算法,对前者输出的轨迹进行聚类(具体地,k-means),从而让hardware measurement花在那些代表性的,还未被充分探索过的点上,从而减少hardware measurement的开销。实验中它与AutoTVM相比,以更少的优化时间达到了更好的性能。它和FlexTensor类似,使用RL作为cost heuristic来指导搜索。

  • 2020年University of California与Google Research的《Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation》中的Chameleon引入了两个组件:RL-based Adaptive Exploration与Adaptive Sampling。前者使用强化学习中的PPO算法,通过exploration and exploitation的trade-off使得搜索算法更好适应搜索空间。后者使用adaptive sampling替换了greedy sampling。相比之下,AutoTVM使用的是均匀采样,而Adaptive sampling使用非均匀采样。由于候选配置一般都是非均匀分布,因此它利用聚类算法(K-means),然后将每类的中心作为代表性的configuration进行处理。这样可以使得measurement更多的在有潜力获得性能提升的地方进行,从而减少性能测量的次数。它既减少执行时间也减少优化时间。基于AlexNet, VGG-16, ResNet-18,在Titan Xp GPU上,相比AutoTVM平均有若干倍的加速,输出程序性能也有所提升。

  • 2020年北京大学的论文《FlexTensor: An Automatic Schedule Exploration and Optimization Framework for Tensor Computation on Heterogeneous System》中的FlexTensor是一个用于多种硬件平台(CPU, GPU, FPGA)中针对张量计算的schedule搜索优化框架。给定张量计算的高层描述,FlexTensor自动生成为CPU,GPU和FPGA的底层实现。与AutoTVM相比,FlexTensor不需要手写schedule template。该框架包含了前端与后端两部分:前端的输入是由Python写的张量计算。它通过对于张量计算的静态分析抽取有用的算子统计信息(循环个数,循环次数,顺序)与结构化的信息(图中节点个数,张量个数,输出与输入张量个数,消费者节点个数)。然后,基于这些信息产生硬件相关调度空间。调度空间通过枚举调度原语(split, reorder, fuse, etc.)和原语相应的参数的组合产生。同时,它还会裁减调度空间,以及根据结构化相似度重新调整schedule。后端用于调度空间的搜索并产生优化后的schedule。在调度空间的搜索过程中,每一次中需要做两个决策:一是从空间中的哪一点作为下一步的起点;二是在该起点下沿哪个方向得到新的点。FlexTensor中,前者使用启发式方法(simulated annealing),后者使用基于机器学习(Q-learning)的方法。它应用不同的优化,包括multi-level loop tiling,loop reordering, loop parallelization, memory customization等。搜索过程中,性能评估既可以在真实硬件上进行也可以基于analytical model。搜索完成后,就可以基于该schedule在各个硬件平台进行实现。代码生成部分基于TVM。

  • 2020年Ant Group的论文《Woodpecker-DL: an efficient compiler for accelerating deep learning on heterogeneous computing architectures》中的Woodpecker-DL(WPK)是一个用于深度学习推理的优化框架。它有几个组件组成:Graph optimization接收模型,执行图优化(constant folding, operator fusion与data layout transformation),然后输出算子的spec。Automated search模块基于前一步输出的spec,基于schedule template,然后采用两种自动化搜索方法,在目标平台上寻找最优的配置。一种是基于遗传算法,其中schedule template中参数的configuration被编码成一个参数化的向量,然后就可以应用遗传算法。另一种是基于强化学习(PPO算法),其中状态空间是一个17维的向量,包含batch size,channel等计算描述信息,以及tile size等硬件相关的schedule信息。动作空间中的动作用于更新状态中的值(如tile size)。DNN网络被训练来基于状态预测动作。经过schedule的搜索后,最后用Halide产生代码。Runtime engine基于生成的算子进行优化后计算图的推理计算。此外,生成的算子还可以作为自定义算子集成进TensorRT。实验中,在P100 GPU上,它能得到比cuDNN与TVM性能更好的卷积算子,并在端到端模型推理上性能优于TensorRT。

  • 2021年Facebook的论文《Value Learning for Throughput Optimization of Deep Learning Workloads》 将寻找schedule的过程建模成序列优化选择问题。我们知道基于贪心的方法优点是快,但缺点是每步只考虑当前信息而无法找到最优解。理论上,如果我们有一个“先知”可以在每步基于look-ahead的信息给出准确的估计,那即使我们用贪心的方法也能得到最优解。这篇文章提出一种预测partial schedule期望性能的方法。具体地,它将选取最优schedule的问题建模为Markov Decision Process(MDP),状态 s i s_i si为模型前 i i i层的schedule,即partial schedule。动作 a i a_i ai为给定 s i s_i si下第 i + 1 i+1 i+1层的合法scheduling选项。为了求解MDP,就需要一个值函数的估计。这个值函数的作用是预测给定状态下如果后续都采用最优schedule能达到的最小执行时间。实现中它基于LSTM,输入为两类特征:与schedule无关的intrinsic特征,以及schedule相关acquired特征。然后通过强化学习中的value iteration方法对值函数进行近似。其中对于每个状态,基于之前的值函数版本执行beam search,然后通过benchmarking得到性能,再对值函数进行修正。有了这个值函数后,就可以使用贪心方法进行优化了 ,且不需要在目标平台上执行,从而快速找到高效的schedule。实验中它找到的schedule执行性能优于Halide和TVM,同时搜索时间有2到3个数量级的加速(从小时级到秒级)。

业界关于auto-tuning的项目也越来越多,除了前面提到较多的Halide,TVM外,还有比如:

  • 2019年Netherlands eScience Center的论文《Kernel Tuner: A search-optimizing GPU code auto-tuner》中的KernelTuner是一个基于Python的用于OpenCL, CUDA和C kernel的auto-tuning工具 。它内置了许多的搜索优化算法。该工具负责编译和benchmarking kernels,自动化搜索最优的kernel configuration。实现的算法有basin hopping,differential evolution,firefly algorithm,simulated annealing,particle swarm optimization和genetic algorithm。除了常见的单kernel调优中的tile size, block size等参数外,它还可以对stream的数量(用于计算与数据传输的overlapping),及kernel pipeline(如reduction)进行tuning。
  • 来自OPEN AI LAB的AutoKernel是一个对Halide算子进行自动调优的项目,支持CPU(x86, ARM)与GPU平台(NVIDIA, ARM Mali)。AutoKernel的自动搜索模块包含了greedy algorithm, reinforcement learning等优化算法。生成的调优算子可以一键集成到Tengine这个深度学习推理框架中。
  • 2015年SURFsara HPC centre的论文《CLTune: A Generic Auto-Tuner for OpenCL Kernels》中的CLTune是一个开源的OpenCL kernel的auto-tuner。它在用户定义的参数搜索空间中进行参数的自动调优。这些参数比如workgroup size, vector data-types, tile sizes, loop unrolling factor。它实现了full-search, random-search, simulated annealing与particle swarm optimization这几种搜索策略。实验中将之用于2D Conv与GEMM。GEMM上调优的kernel性能优于clBLAS。
  • 2021年Huawei的论文《AKG: Automatic Kernel Generation for Neural Processing Units Using Polyhedral Transformations》中的AKG使用auto-tuner来找tiling策略。它使用machine learning guided sampling approach来减少tuning空间。AKG先根据一定的策略计算包含合法tiling参数的tuning空间,然后第一轮采样一些样本,并测其性能。这些样本用来训练机器学习模型,并产生第二批样本。第二轮的样本是从第一轮中最好的N个加上改动而来。这个改动以一定概率根据机器学习模型往高性能的方向前进,或者随机选取。第二批样本也会测量其性能并用来更新模型。
  • 2018年Facebook等机构的论文《Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions》中的Tensor Comprehensions使用Auto-tuning来搜索tile size, block size,grid size,unrolling size,fusion strategy,shared and private memory usage等参数。它采用的是generic search方法。为了加速搜索过程,一方面它使用compilation cache保存产生的CUDA或者PTX code,并放到数据库中以便于重用。另一方面它通过多线程、多GPU将多个编译任务并行化。

结尾

要得到高性能算子实现,基于解析的方法虽然计算快,可解释性强,但往往建模困难,模型的假设较多使其在应用上受限。基于规则的方法基于一些专家设计的规则或heuristics给出解,速度快但往往性能欠优。因此,目前比较流行的方法是通过自动调优的方法。

高性能算子实现的自动调优方法中,一般先设计与定义搜索空间,然后通过某种搜索方法产生一些变体,再对这些变体评估后选出最好的一个。搜索的方法与解决一个典型的组合优化问题时的思路是类似的。最简单的是暴力枚举(Brute-force),为了加速搜索在过程中可以加上剪枝优化。它的好处是可以保证找到最优解,但由于搜索空间巨大,规模稍微大些就抗不住了。另一种思路是贪心法,即先定义评估函数,然后每步都按照评估函数进行最优选择。但贪心法的局限是容易陷于局部最优,导致最终效果较差。这两者中间的一种权衡是beam search,即每步不是只考虑最优那个,而是前n个。另外,此类问题常用的搜索方法还有如SA(Simulated annealing),MCTS(Monte Carlo Tree Search), GA(Generic Algorithm)等。另外,这几年,随着强化学习越来越多地被用于组合优化问题,也出现了不少尝试将之用于算子自动调优。

自动调优带来性能提升的同时,也带来一些挑战,如对于dynamic shape的支持,对于特殊硬件(如张量计算单元)的支持等。它最主要的缺点还是耗时久,这会较大地影响自动调优技术的实用性。如果搜索过程中产生的变体全在真机上评估的话,搜索过程中大量时间会耗费在编译与运行上。为了使得自动调优更加高效,除了对于搜索方法的改进外,业界还有一些其它思路,比如预tuning数据库,并行化,时间分配调度及cost model相关优化等。尤其是对于cost model的改进工作在这几年非常多,如引入图神经网络(GNN),通过benchmark预训练,挖掘变体间的相似性信息等等。

你可能感兴趣的:(深度学习,人工智能,机器学习,Auto-tuning,深度学习编译器)