前面的第一篇"(一)大话深度学习编译器中的自动调优·前言"与第二篇"(二)大话深度学习编译器中的自动调优·DSL与IR"分别介绍了背景与一些相关概念,这第三篇我们开始切入正题,看下现代深度学习编译器中的自动调优(Auto-tuning)方法。Schedule的自动生成,一类方法是基于解析模型(Analytical model),然后使用一些经验公式来产生解;另一类方法是多面体编译技术。它将循环嵌套迭代空间建模为多面体,然后用如整数规划等数学方法求出能提升局部性与并行性的循环变换;还有一类就是经验搜索(Empirical search)方法。这类方法一般采用迭代的方法,通过一定策略产生很多的变体,然后对它们进行性能评估,最终挑选出性能最优的schedule配置(Configuration)。这个过程还可以结合机器学习方法,比如强化学习,让机器学习模型学习schedule参数的选取策略。
在基于empirical search的自动调优方向上,由于Halide与TVM发展时间较长,基于它们的方法比较繁荣。接下来就先以它们为主线展开。我们从一系列相关论文看下它们的大体发展脉络。
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早期受到Halide的启发,也是基于编译器的计算框架,主要用于深度学习。经过几年的快速发展,它在功能与性能上都已比较成熟,且有着众多的衍生项目。TVM中的auto-tuning机制经过了几代的发展:
有了Auto-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以及性能评估方法的改进。
其它还有一些搜索方法上的改进与尝试也值得关注:
关于功能上的扩展,业界也有一些很有意思的工作,如:
这几年开始流行用机器学习的方法来解决程序的优化问题,算子的自动调优也不例外。基于学习的方法的基本思想是训练一个机器学习模型从计算任务得到其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外,还有比如:
要得到高性能算子实现,基于解析的方法虽然计算快,可解释性强,但往往建模困难,模型的假设较多使其在应用上受限。基于规则的方法基于一些专家设计的规则或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预训练,挖掘变体间的相似性信息等等。