基于自适应静态数据布局策略的深度学习张量程序自动生成框架①

樊 哲 南子渊 郝一帆 杜子东③ 陈云霁

(∗中国科学院计算技术研究所计算机体系结构国家重点实验室 北京 100190)

(∗∗中国科学院大学 北京 100049)

0 引言

深度学习作为一个热门研究方向,其呈现出多样化和轻量化的发展趋势。多样化趋势主要体现在模型和硬件的不断更迭。模型发展如AlexNet[1]、ResNet[2]、CapsNet[3]和Transformer[4]等;硬件发展如中央处理器(central processing unit,CPU)、图形处理器(graphics processing unit,GPU)、DianNao 系列[5-6]、张量处理器(tensor processing unit,TPU)[7]以及各种针对特定算法的专用加速器[8-10]。轻量化趋势主要体现在深度学习应用从云端到边端的迁移,即越来越多的深度学习应用被部署到算力受限的边缘端设备[11-13]。这对深度学习应用的算法优化提出了越来越高的要求。上述发展趋势带来了深度学习应用的编程和性能之间的矛盾。

深度学习张量程序(张量算子的低层级实现[14])的自动生成框架致力于解决编程与性能之间的矛盾。其被广泛用于深度学习编译器(如张量虚拟机(tensor virtual machine,TVM)[15])以自动生成高性能的张量程序。其通过搜索的方式找寻在指定硬件后端上性能最优的张量程序,之后交由对应的硬件编译器生成二进制指令,在减轻用户编程负担的同时尽量保证性能。

为了生成高效的深度学习张量程序,必须考虑静态数据的布局问题。原因主要有3 个方面:(1)对静态数据的布局变换可发生在编译时期,而不占用程序的运行时间;(2)在深度卷积神经网络的前向推理(inference)过程中,静态数据(如权值、偏置等)参与的运算高达总运算量的90%[16-17];(3)数据布局可以提升程序访存局部性和硬件资源利用率[18-20]。因此,静态数据的布局对深度学习张量程序的性能影响十分显著;而如何确定合适的静态数据布局,使得程序执行效率最高,是深度学习张量程序自动生成框架面临的重大挑战。

作为目前应用最广泛、最具前景的深度学习张量程序自动生成框架,Ansor[14]应对上述挑战的方案是:根据预先指定的单一数据布局策略,训练性能预测模型,依据该模型搜索最佳性能的张量程序。但其存在单一策略非最优和性能预测模型不准确的问题。具体地,单一的数据布局策略无法使Ansor在所有任务上都生成性能最优的程序,而对于一个任务,其最优策略无法预先确定,因此预先指定单一策略的Ansor 生成的程序仍有性能提升空间(即单一策略非最优的问题)。另外,指定一种策略后,Ansor 会依据性能预测模型C从任务的程序空间(即搜索空间)中采样一个程序A,用对A进行布局变换得到的程序X去训练C,以指导下一次采样,直到搜索过程结束。最后Ansor 输出搜索过程中遇到的性能最优程序X∗。从优化角度看,X所在程序空间即为优化空间。C并非直接在优化空间中寻优,而是通过在搜索空间中采样,间接在优化空间中寻优。优化空间的程序因经过了数据布局,访存局部性要比搜索空间的程序好,因而优化空间和搜索空间在访存特性上存在巨大差异。这种差异导致Ansor 得到的程序X∗仍有性能提升空间(即性能预测模型不准确的问题)。

针对上述问题,本文提出基于自适应静态数据布局(adaptive layout,AL)策略的深度学习张量程序自动生成框架AL-Ansor。其通过自适应地选取多种静态数据布局策略来解决单一策略非最优的问题;通过多种策略共同训练性能预测模型C来解决模型不准确的问题。具体地,AL-Asnor 可先以直接寻优的布局策略预训练C,将搜索空间限制在静态数据访存局部性较好的程序子空间。然后,再以间接寻优的布局策略(此时搜索空间和优化空间的不一致已被缓解)微调C,使之更容易在搜索空间中采样出程序A∗,其经布局变换后的程序X∗的性能最优。本文在32 核Intel Xeon CPU 上对多个卷积层进行实验验证,结果表明,在同样的搜索次数下,相较于3 种基于单一策略的Ansor,AL-Ansor 生成的张量程序分别有13.81%、12.41%和16.59%的平均性能提升。

综上,本文的主要贡献有以下3 点:

(1) 分析了Ansor 的工作原理,发现了其单一策略非最优和性能预测模型不准确的问题。

(2) 提出了自适应静态数据布局策略AL 以及基于此策略的张量程序自动生成框架AL-Ansor。

(3) 在典型深度神经网络的多个卷积层上进行实验验证,证明了AL 策略的有效性。

1 背景

1.1 张量程序自动生成框架

由于深度学习模型和硬件的多样化发展,基于搜索的张量程序自动生成框架成为一个重要研究方向。这些框架以描述深度学习模型/算子的计算图(computational graph)和目标硬件信息作为主要输入,以优化硬件执行时间为主要目标执行搜索过程,来自动生成对不同硬件后端具有通用表达能力的张量程序。

张量程序自动生成框架按照其搜索空间的定义方式可以分为2 类:基于模板搜索(如文献[21])和无模板搜索(如文献[14,22]),如图1 所示。基于模板搜索的搜索空间由用户自定义模板来指定;而无模板搜索的搜索空间无需用户参与指定。AutoTVM 是基于模板搜索的一种典型框架。在该框架下,用户需要针对不同硬件自定义不同算子的搜索模板。搜索模板的内容包括算子的实现方式,输入数据应该如何变换布局,以及输入数据以什么方式(循环结构、循环顺序、向量化、并行化等)进行计算等。算子搜索模板内的不确定因素(循环的拆分因子、是否向量化等)定义了输入计算图对应的张量程序搜索空间。但这样以自定义模板方式定义的空间是非常受限的,也不易扩展[14]。Ansor 是无模板搜索的一种典型框架。其搜索空间由内置的与算子解绑的通用派生规则(derivation rules)定义:这些规则可作用于任意输入计算图来实现对计算图的改写,每个改写的计算图对应一种张量程序实现,这些构成输入计算图的张量程序空间,即搜索空间。综上,Ansor 无需为每个计算图专门定义搜索空间,也无需用户参与,因此,其更容易应用于新的算子,也缓解了用户的编程负担,是目前使用最广泛、最有前景的深度学习张量程序自动生成框架。

图1 张量程序自动生成框架示意图

1.2 数据布局优化

数据布局指的是数据在内存中的摆放顺序。改变一个程序所需数据的布局,会影响程序的访存局部性和硬件的计算资源利用率。数据布局优化,即通过改变数据布局,改善程序的访存局部性,使整个程序的性能更好。

图2 展示了矩阵乘程序的一个数据布局优化过程。图2(a)展示了一种实现(M,K)×(K,N) 矩阵乘的程序A。假设2 个输入矩阵U、V以及输出矩阵Z都按行(row-major)存储。对A的输入数据V进行数据布局优化后,会得到2 个程序:图2(b)所示的布局变换程序L和图2(c)所示的布局后程序P。布局变换程序L可根据程序A对数据V的下标访问规则自动生成(Ansor 中的数据布局变换也是这样实现的),使得布局后的数据V’在同样的循环结构,即程序P中拥有更好的访存局部性。布局后程序P和程序A的差别如图2(c)中加粗部分所示,仅在于数据V’。数据布局优化除了使程序P拥有更好的访存局部性外,还可以提升硬件计算资源的利用率。对于一个拥有NI个KI维向量点积运算单元的硬件平台,程序P的最内两重循环(即图2(c)中虚线框部分)可以直接由这些向量点积运算单元处理,而程序A由于对数据V的访存不连续,无法直接利用这些运算单元。如果数据V是编译时可确定的静态数据,则上述数据布局优化称为静态数据布局优化,此时布局变换程序L可在编译过程中执行,不会带来运行时开销。

图2 数据布局优化示例(虚线框部分可进行并行加速,从而提高硬件计算资源利用率;加粗部分展示了P 和A 的差异)

对于深度学习应用而言,数据布局优化是提升程序性能的一个重要优化策略。Caffe[23]在实现卷积操作时,先对输入数据进行im2col 布局变换,再调用高效的通用矩阵乘(general matrix-matrix multiplication,GEMM)接口实现卷积功能。文献[24]为在单指令数据流(single instruction multiple data,SIMD)架构上实现高效的卷积运算,需要将输入布局设置为NCHWc,权值布局设置为KCRSck。LDLNDK(neural network development kit with labled data layout)[25]使用标签标记静态数据布局信息,在编译时进行布局变换,在深度学习加速器上达到数倍性能提升。CDUCA(computation and data unified compile architecture)[26]的数据优化器模块在部署前对数据进行面向硬件平台的数据布局等优化,提升硬件平台的访存和运算效率。NeoCPU[27]针对神经网络在CPU 上的前向推理过程,在图级别对各算子的神经元数据布局进行规划,再对静态权值数据进行相应的布局变换,达到高性能。

2 问题定义

为了更加清晰地进行后续论述,本文在此定义并引入一些符号,并对张量程序自动生成框架的静态数据布局问题进行定义。

2.1 任务及程序空间

任务T由输入数据、输出数据以及输入数据到输出数据的映射规则共同定义。输入数据根据其内容是否为编译时确定又可以分为静态输入数据和动态输入数据。例如,图2 所示为一个将M×K维输入数据和K×N维输入数据进行矩阵乘得到M×N维输出数据的任务MatMul。

程序解决给定任务(实现输入数据到输出数据映射规则)的一个有限可终止的步骤序列(包括访存步骤和计算步骤)。对于图2 所示程序A和P,虽然其计算步骤是相同的,但是其输入(分别为V和V’)和访存步骤不同,因此是2 个不同的程序。

任务T的程序空间AT解决给定任务T的所有程序构成的空间。图2 中程序A和(L+P) 均为任务MatMul 的程序空间AMatMul中一点。对于一个具体的搜索过程(如Ansor 的搜索过程),任务T的程序空间受限于搜索空间的生成方式,但为了描述的方便,会将该搜索过程所能探索到的程序空间记为AT。

任务T的布局分离程序空间ST是对空间AT中每个程序进行数据布局优化所得到的程序构成的空间。显然,该空间是AT的一个子空间。在图2 所示的MatMul 任务中,对程序A∈AMatmul进行数据布局优化后得到的组合程序(L+P)为SMatMul空间中一点。

任务T的布局后程序空间PT是对空间AT中每个程序进行数据布局优化所得到的布局后程序构成的空间。因此,该空间的程序访存局部性都很好,与空间AT在程序访存特征上存在较大差异。根据定义,AT和ST上的点可以在PT上找到点与之对应,反之,PT上的点也都可以在AT和ST上找到点与之对应。但PT并不是空间AT或ST的子空间,因为PT中程序解决的任务并不是T本身,而是T经过布局变换程序后的子任务(该子任务的输入由于数据布局变换,已不同于任务T中对应的输入,如图2 中的V和V’)。图2(c)所示程序P即为PMatMul空间中一点。

图3 展示了上述术语的关系。对于一个任务T,从其程序空间AT中取一程序A,其通过数据布局优化会产生布局变换程序L和布局后程序P。其中,程序(L+P) 属于任务T的布局分离程序空间ST,程序P属于任务T的布局后程序空间PT。虚线展示了程序A到P的数据布局优化路径。反过来,PT中的每一个程序也都可以在ST和AT中找到对应的程序。即对于一个任务T,这3 个空间中存在一一对应的路径。

图3 3 种程序空间的关系以及其中程序的对应关系

2.2 张量程序自动生成框架的静态数据布局

对于张量程序自动生成框架,给定任务T和目标硬件H,其静态数据布局问题定义如下。

定义1在程序空间XT中搜索满足优化目标G的张量程序X∗,使得在H上解决任务T所需的运行时间最短。

根据张量程序自动生成框架所采取的搜索策略(静态数据布局策略)的不同,定义1 中的程序空间XT(优化空间)和输出程序X∗会有所不同。

对于本文关注的Ansor 而言,其共有3 种搜索策略:(1)不进行布局(no layout,NL)策略;(2)运行时布局(layout during running,LR)策略;和(3)编译时布局(layout during compiling,LC)策略。其中第3 种策略只能用于静态数据,而前2 种对于动态数据和静态数据均适用。在Ansor 中,不同策略使用不同的优化空间XT(AT或ST或PT),但却使用同一个搜索空间AT,因为定义于输入计算图上的通用派生规则无法直接生成位于ST和PT中的程序。因此,Ansor 只能从AT中(准确地来说应该是AT-ST)采样程序进行搜索,然后经数据布局变换映射到相应优化空间(AT或ST或PT)上进行测量。这一点体现在优化目标G上。

对于策略NL,定义1 中的XT是任务T的程序空间AT,输出程序X∗是空间AT中的程序A∗。优化目标G如式(1)所示,其中,timeof(A,H) 表示程序A在硬件H上的执行时间。

对于策略LR,定义1 中的XT是任务T的布局分离程序空间ST,输出程序X∗是空间ST中的程序(L+P)∗。优化目标G如式(2)所示,其中,LT0(A)表示对程序A进行数据布局优化得到的布局变换程序,LT1(A) 表示对程序A进行数据布局优化得到的布局后程序。

对于策略LC,定义1 中的XT是任务T的布局后程序空间PT,优化目标G如式(3)所示,搜索的输出程序是空间PT中的程序P∗。

对于策略LR 和策略LC,由于优化空间XT和搜索空间AT并非同一空间,因此其优化做法——依据在优化空间XT上定义的优化目标G去采样位于搜索空间AT中的程序A,以求A经过变换后得到的程序X能有最小的运行时间timeof(X,H)——是无法达到目标的。尤其是对策略LC,因为其优化空间PT与搜索空间AT中程序的访存特征差异过大。

3 采用自适应静态数据布局策略的AL-Ansor

3.1 Ansor 的工作原理及静态数据布局策略

Ansor 是一种深度学习张量程序自动生成框架,其工作原理如图4 所示。

图4 Ansor 工作原理示意图

Ansor 以待优化算子(记为任务T)的计算图、目标硬件信息、搜索次数以及指定的某一静态数据布局策略(NL、LR 或LC)为输入,经过指定次数的搜索后,输出在目标硬件上执行时间最短的张量程序X∗。在每一轮搜索过程中,Ansor 主要分为4 步进行:(1)程序采样。依据性能预测模型C,从由输入计算图和图上的通用派生规则共同定义的张量程序空间AT中采样若干程序。(2)静态数据布局。依据输入的静态数据布局策略,对上一步采样出的每一个程序A进行相应的变换,得到程序X。(3)编译并执行。将程序X编译至目标硬件平台,并将相应二进制指令部署至目标硬件平台上运行,这个过程中会得到程序X的指令、运行时间等信息。(4)训练性能预测模型C。从上一步的编译及运行信息中提取特征训练性能预测模型C,用来指导下一轮搜索过程中第一步的程序采样。

在Ansor 的一次完整搜索流程中,需要指定固定的一种静态数据布局策略。不同的静态数据布局策略会影响性能预测模型C的训练。

对于策略NL,程序X就是采样程序A。该策略下,性能预测模型C以程序A的信息来训练。此时Ansor 的搜索空间和优化空间均为AT,优化目标为2.2 节中的式(1)。

对于策略LR,程序X是采样程序A经过静态数据布局变换而产生的布局变换程序L和布局后程序P的组合。该策略下,性能预测模型C以(L+P)∈ST的特征信息训练,来进行程序A∈AT的选择。此时Ansor 的搜索空间为AT,而优化空间为AT的子空间ST,优化目标为2.2 节中的式(2)。

对于策略LC,程序X是采样程序A经过静态数据布局变换产生的布局后程序P,因为布局变换程序L被指定在编译时执行。该策略下,性能预测模型C以P∈PT的特征信息训练,来进行程序A∈AT的选择。此时Ansor 的搜索空间为AT,而优化空间为访存局部性更好的布局后程序空间PT,优化目标为2.2 节中的式(3)。优化空间和搜索空间不在同一个空间,且其中的程序访存特征具有较大的差异。

表1 列出了以32 核Intel Xeon CPU 为目标硬件,以ResNet-50 的4 个卷积块(包括输入补pad、2D卷积、加偏置以及做ReLU 激活)为搜索任务,Ansor分别在3 种静态布局策略下搜索1024 次得到的最优程序X∗的运行时间。可以看出,单一的策略NL、LR 和LC 并不总在所有搜索任务上表现得最好。

表1 Ansor 的3 种静态数据布局策略搜索结果

3.2 自适应静态数据布局策略

从Ansor 的工作原理可知,其存在以下2 个问题:

(1)静态数据布局策略需要预先指定,而针对不同的任务,最优的策略不同,因此预先指定的静态数据布局策略往往不是最优的(单一策略非最优)。

(2)对于编译时布局策略LC,搜索空间AT与优化空间PT存在较大差异,主要体现在程序对静态数据的访存特性上,PT上的程序访存特性更好。由此训练得到的性能预测模型难以找到性能最优的程序(性能预测模型不准确)。

针对这些问题,本文提出自适应静态数据布局策略(AL)。其主要特点是,在搜索过程中不是只使用一种固定的静态数据布局策略,而是根据一个自适应策略选择函数Select_Strategy 选择并应用多种静态数据布局策略(NL/LR/LC),来共同训练性能预测模型C。一种可能的自适应策略是:在进行策略LC 前,先使用策略NL 和策略LR 来训练模型C。这样不仅有机会探索到不同策略下的最优程序,还可以缩小待搜索空间。之后再采取策略LC 微调模型C,此时优化空间和搜索空间的差异不那么显著,使得模型C更容易在搜索空间AT中找到程序A∗,其对应的布局后程序P∗在目标硬件上执行时间最短。

自适应静态数据布局策略AL 的伪代码描述如算法1 所示。

算法1 的输入是任务T的张量程序空间AT(也即该算法的搜索空间),目标硬件H,搜索轮数L,以及初始性能预测模型C。输出是在搜索过程遇到的在目标硬件H上执行时间最短的布局后程序P∗。

在每一轮搜索过程中,主要经过了以下过程。

(1)先根据目标硬件H和性能预测模型C在搜索空间AT中采样出候选的程序集合A_Cand(第3行)。

(2)对集合中每个程序进行数据布局变换,得到布局变换程序集合L_Cand 和布局后程序集合P_Cand(第4 行)。

(3)根据自适应策略选择函数Select_Strategy确定本轮搜索应使用的静态数据布局策略(第5行)。该函数的输入可以有当前的搜索轮数Search_Idx、搜索总轮数L以及性能预测模型C等。该函数的实现可以是一种启发式算法,一种预定义规则,或者是一个预训练的决策树或神经网络模型。

(4)根据本轮搜索的静态数据布局策略Strategy,对相应的程序进行编译,并部署到目标硬件H上执行(第6 行)。当Strategy 为NL 时,Build_Run函数只会对A_Cand 和P_Cand 中的程序进行编译并运行得到相应的编译及运行信息A_Infos 和P_Infos。对A_Cand 编译并运行是为了训练性能预测模型C,使其能够锁定一个更小的搜索空间。对P_Cand 编译并运行是为了发现具有最短执行时间的布局后程序。当Strategy 为LR 时,Build_Run 函数只会对L_Cand 和P_Cand 中的程序进行编译并运行得到相应的编译及运行信息L_Infos 和P_Infos。当Strategy 为LC 时,Build_Run 函数只会对P_Cand 中的程序进行编译并运行得到相应的编译运行信息P_Infos。

(5)根据策略Strategy 选择有效的编译及运行信息去训练/微调性能预测模型C(第7 行)。

(6)根据布局后程序集合P_Cand 的编译运行信息P_Infos 中的运行时间信息更新搜索过程中发现的最优布局后程序P∗(第8～11 行)。

与Ansor 所采用的单一数据布局策略相比,ALAnsor 有2 个特点:(1)一个任务T的不同搜索轮次之间可以使用不同的静态数据布局策略(由Select_Strategy 决定);(2)无论何种静态数据布局策略,每一轮都会编译并运行布局后程序集合P_Cand。

3.3 基于AL 策略的AL-Ansor

将自适应静态数据布局策略AL 应用于Ansor的整个搜索过程中,便得到了改进的Ansor——ALAnsor,如图5 所示。和图4 的主要差别在于程序采样和编译执行之间的数据布局变换过程。在ALAnsor 框架中,静态数据布局策略的确定不再依靠一次性的输入,而由内部的自适应策略选择函数决定。输出总是执行时间最短的布局后的程序P∗∈PT,而不是依赖于输入的静态数据布局策略的程序X∗。

图5 AL-Ansor 工作原理示意图

4 实验

4.1 实验设置及任务负载

本文采用的基准Ansor 框架即TVM 库(版本为e718f5a8)的auto_scheduler 模块。根据Ansor 的输入静态数据布局策略的不同,分别将输入NL、LR 和LC 策略的Ansor 记为NL-Ansor、LR-Ansor 和LCAnsor 以方便后续实验展示。Ansor 和AL-Ansor 搜索框架的目标硬件为32 核Intel Xeon Gold 6226R CPU(@2.90 GHz),搜索次数均设置为1024 次,因为这一次数已足够使性能预测模型收敛[14]。实验中AL-Ansor 所采用的自适应策略选择函数Select_Strategy 为一种简单的预定义规则:依次进行NL、LR和LC 策略,且3 种策略搜索次数的占比为5 ∶5 ∶6。

本文采用的任务负载是一些常见于深度卷积神经网络中的卷积块(包含输入补pad、进行2D 卷积、加偏置以及做ReLU 激活),因为卷积是其主要计算负载。这些卷积块用TVM 的topi 接口进行描述,然后分别调用NL-Ansor、LR-Ansor、LC-Ansor 和ALAnsor 进行搜索优化。所选卷积任务的输入、输出数据如表2 所示。

表2 实验负载

从2 个方面对4 种框架进行比较。一个是不同框架下的搜索时间(或者编译时间,因从计算图到生成张量程序的搜索过程可以视为编译);另一个是它们搜索出的最优程序的运行时间。在进行实验时,独占目标服务器测量时间。对于最优程序运行时间的测量,采取以下方法:对搜索得到的最优程序进行3 组测量取平均值,每组测量100 轮取中位数,每轮运行至少500 ms(根据程序执行时间,每轮的执行次数在几百至几千的范围),以最小化可能存在的测量误差。

4.2 搜索时间对比

图6 展示了NL-Ansor、LR-Ansor、LC-Ansor 和AL-Ansor 的搜索时间对比。其中AL-Ansor 的搜索时间被拆分为3 部分:AL-Ansor-base、AL-Ansor-exb和AL-Ansor-exr。AL-Ansor-exb 和AL-Ansor-exr 分别表示AL-Ansor 的自适应静态数据布局策略AL 所引入的程序P的编译开销和运行开销,因为在搜索过程中无论选取哪一种布局策略,都要生成P并对其进行编译和运行,这一点通过对比图4 和图5 可以看出。但这一开销只会发生在自适应策略选择函数返回策略NL 的搜索轮次;对于返回策略LR 和策略LC 的搜索轮次,程序P本就包含在程序X之中,因此不会产生额外的编译开销和运行开销。ALAnsor-base 所示的时间则为把这部分开销去掉的时间。

图6 NL-Ansor、LR-Ansor、LC-Ansor 和AL-Ansor的搜索时间对比

从图6 可以观察到以下现象:

(1)AL-Ansor-base 在10 个任务上的搜索时间相较于NL-Ansor、LR-Ansor 和LC-Ansor 平均提升了13.86%～18.24%。这些主要来自于布局变换后的程序X的编译和运行时间提升,而程序X是采样程序A经过静态数据布局变换得到的,这表明策略AL下决定采样质量的性能预测模型C要优于其他单一策略。

(2)相较于AL-Ansor-base,在每个搜索轮次引入对程序P的评估(AL-Ansor-exb 和AL-Ansor-exr)会带来平均41.82%的时间开销。如前所述,该部分开销主要是由策略NL 所在搜索轮次引入的。这部分开销中87.67%来自于程序P在目标硬件上的运行过程(AL-Ansor-exr),这是为了估计P部署到目标硬件后的真实运行时间,以在搜索过程中选出最优的布局后程序P∗。AL-Ansor-exr 的开销可以根据需要去调整:如果希望对程序P的评估尽可能准确,则可以将这部分时间延长;如果希望整个搜索过程快一些,则可以将部分时间适当缩短。

(3)在10 个任务上,相较于NL-Ansor、LR-Ansor和LC-Ansor,AL-Ansor 平均上有15.95%～22.17%的搜索时间开销。但是,由于AL-Ansor 内的静态数据布局策略是自适应选择的,而Ansor 中无法预先确定哪一种策略下能找到当前任务的最优张量程序。所以,为了能找到最优的张量程序,Ansor 需要尝试每一种策略。从这个角度看,AL-Ansor 的搜索时间相较于Ansor 平均会有2.5 倍的搜索时间提升。

总体来看,AL-Ansor 在搜索时间上仍然是优于最优静态数据布局策略未知的Ansor 的。虽然相较于确定采用某一种静态数据布局策略的Ansor(NLAnsor、LR-Ansor 或LC-Ansor),AL-Ansor 会有一定的搜索时间开销。但是,对于一个确定的任务而言,搜索过程往往是一次性的,而搜索得到的最优张量程序在部署到目标硬件后则会反复运行使用。因此部署后运行时间的提升更为关键,部署前搜索时间的少量开销是可以接受的。下面将分析4 种自动生成框架得到的最优张量程序运行时间。

4.3 最优程序运行时间对比

图7 中依次展示了由NL-Ansor、LR-Ansor 和AL-Ansor 搜索出的最优张量程序X,相对于LC-Ansor 搜索出的最优张量程序Y的性能提升百分比y,即y由式(4)得到:

图7 NL-Ansor、LR-Ansor 和AL-Ansor 相对于LC-Ansor的性能提升

其中timeof(Y,H)为程序Y在硬件H上的执行时间。

在图7 中,条形图所处方位展示了对应框架采取策略与LC 的优劣:位于x轴上方表示对应策略优于策略LC,位于x轴下方则表示劣于策略LC。条形图的高度则展示了对应策略与LC 的差距:高度越大,表示对应策略所得最优程序与策略LC 所得最优程序之间的运行时间绝对值之差越大。从图7中可以得出如下结论。

(1)Ansor 原本支持的3 种独立策略NL、LR 和LC 并没有绝对的优劣之分。这再一次表明了由用户预先指定单一的静态数据布局策略是不恰当的。表3 展示了3 种策略的不同优劣排序在任务负载中均曾出现。

表3 3 种策略优劣排序在T1～T10 上的发生情况

(2)策略NL、策略LR 和策略AL 相较于LC 分别平均提升了2.59%、4.13%和16.59%的性能。策略NL 和策略LR 相较于策略LC 其性能平均上有提升,这一点虽然反直觉,但也进一步印证了在策略LC 下,搜索空间AT和优化空间PT之间的差异对于张量程序搜索框架的不利影响。

(3)策略AL 在绝大多数任务上优于其他策略。只在任务T4 和T6 上存在2 处例外,在T4 上,策略AL 相对于LC 损失了1.88%的性能;而在T6 上,AL相对于NL 损失了1.31%的性能。这些微小的损失可能来自于搜索过程中为了跳出局部最小点的探索机制,而此过程是随机不可控的。但值得注意的是,由此为策略AL 带来的影响是很小的。除去这2 个异常点,策略AL 相对于策略NL 提升了0.98%～26.40%(T10 上最小,T2 上最大,见图8),相对于策略LR 提升了3.79%～27.91%(T4 上最小,T5 上最大,见图8),相对于策略LC 提升了3.58%～36.80%(T3 上最小,T7 上最大)。

图8 AL-Ansor 相对于NL-Ansor 和LR-Ansor 的性能提升

图8 单独展示了AL-Ansor 相对于NL-Ansor 和LR-Ansor 的性能提升百分比。在所选的10 个任务负载上,策略AL 相对于NL 平均提升了13.81%,相对于LR 平均提升了12.41%。

综合来看,相较于使用单一静态数据布局策略的Ansor,采用自适应静态数据布局策略的AL-Ansor 框架更加合理和高效。

5 结论

本文研究了静态数据布局在深度学习张量程序自动生成框架中的应用。发现应用最广泛、前景最佳的Ansor 框架存在2 个问题:(1)只能指定单一的静态数据布局策略(NL 或LR 或LC)进行张量程序的搜索,而单一布局策略具有非最优问题;(2)对于专用于静态数据布局的策略LC,其搜索空间和优化空间因其中程序的访存特征差别而存在较大差异,导致性能预测模型不准确的问题。这些问题都会导致搜索得到的张量程序仍有较大优化空间。针对这些问题,本文提出了自适应的静态数据布局策略AL及基于此策略的深度学习张量程序自动生成框架AL-Ansor。实验表明,在张量程序搜索过程中,自适应地使用多种不同的静态数据布局策略,不仅能够缓解如何选择合适布局策略的难题,也能够减小搜索空间和优化空间的差异,因此搜索生成的张量程序具有更好的性能。这为深度学习的高效部署与应用带来了广泛的好处。

虽然本文讨论的主要对象是静态数据,目标硬件平台选取的是CPU,但对于动态数据或者其他目标硬件平台(如GPU、TPU 等),本文所揭示的原理也是适用的。另外,本文所提出的自适应策略选择函数是决定该方法好坏的一个因素,如何选择它仍是有待研究的课题。