Polyhedral模型作为编译优化领域实现循环嵌套优化的一种手段,已经被许多商业和研究型的编译器使用。作为一种抽象程度比较高、理解起来比较复杂的编译模型,Polyhedral编译技术的受众本来是比较少的。不过,深度学习编译器里面使用了Polyhedral模型,让许多人对这个工具产生了极大的兴趣。正是因为Polyhedral模型能够完成那些上层用户真正想让编译器去完成的事情,所以Polyhedral的实现细节成了许多芯片厂商以及系统开发人员比较关注的问题。作为一种编译阶段静态分析和优化的技术,在针对单算子或一个独立的循环嵌套进行优化时,Polyhedral模型并不能轻易地得到与手工优化算子库或其它通过手工指定调度的方法相当的性能,但是在多个算子融合或多个循环嵌套之间快速寻找优化手段方面,Polyhedral却有着相对较大的优势。
前期做了一些在深度学习编译优化领域结合Polyhedral模型的研究,也与不少从事该领域的专家和学者进行了讨论。一些团队开始对Polyhedral编译模型非常感兴趣,也很想入手,借助Polyhedral模型来解决实际遇到的问题,但是许多团队后来都放弃了。交流之后发现,最大的原因在于Polyhedral编译模型具体做了哪些优化、如何实现的这些优化他们并不清楚,只能把Polyhedral的优化当作一个“黑盒”处理,使得优化后的程序效果并不可控。其中,一些对Polyhedral模型相对比较了解的团队还可以通过一些编译选项来选择、控制优化的过程,相当于只了解一些“黑盒”上的开关,让Polyhedral模型的“黑盒”实现或不实现一些优化效果,至于“黑盒”内具体如何实现的,仍然一无所知。
Polyhedral模型实现优化的核心在于其调度算法,上述所谓“黑盒”就是Polyhedral调度算法的具体设计和实现,是一种启发式算法。这种启发式算法允许设计者将目标平台体系结构的特征写入到启发式算法里,让调度算法根据指定启发方向寻找最优解,这种启发方式是通过求解(整数)线性规划方法来达到解决最优化问题实现的,所以启发式算法的设计就在于如何将目标体系结构相关的约束抽象成最优化问题的优化目标和/或约束写入到启发式算法里。
然而,由于Polyhedral模型涉及的基础数学知识和计算机理论抽象问题比较多,使得许多介绍Polyhedral的文章非常晦涩难懂,尤其是关于Polyhedral调度算法的文章,即便是许多研究该领域的研究人员也很难理解其中的含义。为了能够真正弄清楚Polyhedral模型调度算法的原理,结合当下使用范围最广泛的整数线性规划求解工具isl中的实现,作者在这里把对Polyhedral调度算法的理解写成系列文章的形式,希望能够为一些对Polyhedral模型感兴趣的朋友提供一些参考。该系列文章一共分为三个部分,第一部分介绍Pluto算法,第二部分介绍Feautrier算法,第三部分介绍isl库中的调度算法。由于涉及内容较多,每一部分的内容篇幅比较长。
注:本文为大家学习和了解Polyhedral编译技术做参考,由于作者水平有限,文章内容难免有不少错误之处,还请各位读者多谅解。另外,文章内容仅代表个人意见,版权归作者所有,转载请联系作者获取许可。
- Pluto算法的建模过程
为了讲清楚Pluto算法的原理,这里用一个如图1-1所示的一维stencil计算来说明一下Polyhedral编译技术里面的调度是指什么。图1-2是这段代码空间多面体表示及其依赖关系,该迭代空间(t,i)为基,每个黑色的点表示一次循环迭代,也可以理解为
如果只考虑循环迭代之间的并行性,通过分析图1-2不难发现,沿i轴方向上的所有循环迭代之间不存在依赖,也就是处在同一水平面上的循环迭代是可以并行执行的。但是现在没有考虑到如何使用t轴和i轴方向上的数据局部性。当N和T较大时,这种计算方式Cache的命中率就会下降。一个很自然的优化目标是如何让编译器能够自动计算出能够充分利用沿多个轴上的数据局部性,这种任务是通过循环分块来完成的。所以,Pluto算法的一个非常重要的目标是通过变换让上述程序能够实现循环分块。
在Polyhedral编译中,代码生成是以表示程序迭代空间的集合和表示程序调度的映射作为输入生成的代码。其中,表示程序调度的映射必须是从语句实例集合和字典序之间的仿射函数。注意:字典序即可以先简单理解成一个语句循环索引变量构成的向量,对于该例子,可以认为每个语句实例的字典序可以用(t,i)表示。对于图1-2所示代码,Polyhedral模型生成的调度是
在这样的空间上,Polyhedral编译器的代码生成只能是沿着 (t, i) 这个坐标系的基的两个方向进行,因为Polyhedral模型要求在形如
但是,这样的循环分块却不是合法的,因为图中任意两个水平相邻的分块之间存在依赖环,这意味着这种循环分块方式生成的代码会破坏原始程序的语义。要使循环分块合法,就要寻找一种能够不会导致分块之间依赖环的分块形状,而Polyhedral代码生成器又限制了循环分块的形状只能沿 (t, i) 坐标轴的方向切割。
从图1-3中可以看出,t 轴的切割方向与程序中的两个依赖,即指向左上和右上两个方向的依赖都相交,这是导致分块后水平分块之间有依赖的主要原因。如果我们沿着这两个依赖其中的一个方向进行切割并构造新的分块形状,那么我们可以得到如图1-4和1-5的两种分块形状。不难发现,这两种分块形状都不会导致任意方向相邻两个分块之间的依赖环,说明这两种分块形状都是合法的。所以,Pluto算法的任务就是要在保证代码生成器需要的输入满足条件的前提下,构造出如图1-4或1-5的分块形状。
要构造出上述两种分块形状或其中的任意一种,Polyhedral编译技术可以采取的措施是在分块之前对循环的迭代空间实现其它辅助循环变换或多种循环变换的组合,使上述几个以 (t, i) 为基的坐标系变换成另外一组基,其中一个坐标轴仍为 i 保持不变,另外一个坐标轴与指向左上或右上的依赖平行,这样才能得到合法的循环分块。为了能够得到合法的循环分块形状,Polyhedral编译技术还必须能够自动实现对循环分块友好的循环变换及这些循环变换的不同组合。
对于这个例子,Pluto算法总是计算出如图1-4所示的分块形状,该分块形状的斜边与指向左上的依赖平行。如果用依赖距离向量表示该依赖的话,这个指向左上方向的依赖可以用距离向量 (1,-1) 来表示,所以该图中所示的分块形状的斜边与 (1,-1) 这个向量所在的直线平行。如果对依赖距离和依赖相关的知识不了解的,请参考该书籍第二、三章内容。我们需要找到垂直于这个方向的一个向量,这个方向是所有分块形状斜边的法向量。显然,这个法向量可以表示成 (1,1)。所以,如果我们换一个坐标系来看这个分块形状,如图1-6所示,该坐标系的基是 (t, t + i)。其中,纵轴 t 与原始 i 轴垂直,代表了分块形状横向边的法向量对应计算出的坐标轴表达式;横轴 t + i 的表达式是由刚才所说的法向量 (1,1) 与原始坐标系的基 (t, i) 向量的转置相乘得到的。 由于坐标系的基发生了改变,原始迭代空间也由原来的矩形变成了平行四边形,如图1-6所示。 此时,所有的依赖都垂直向上或指向右上方向。将该程序在图1-6坐标系下的顺序关系传递给代码生成器,这时代码生成器可以沿该坐标系的两个轴的方向切割迭代空间,在新的坐标系中得到的循环分块示意图如图1-6绿色分块所示。如果我们将图1-6中的分块再放在以 (t, i) 为基的坐标系中看,就可以得到如图1-4所示的分块形状。
所以,Pluto算法也可以理解成在寻找一组新的坐标系的基,使代码生成器能够在该坐标系上生成带有循环分块的代码。对于我们讨论的例子,原始的迭代空间坐标系的基是 (t, i),变换后的坐标系的基是 (t, t + i),那么 Pluto 算法等价于求解一个整数系数矩阵 M 和常量向量 C, 使得
其中,t 和 i 是原始迭代空间坐标轴的表达式,即循环索引变量,矩阵 M 和(t, i) 的转置的乘积再加 上向量C后与向量(t,t+i)的转置相等,
计算矩阵的过程是按照每一行的方式进行求解的,也就是说先计算(
表示,其中,v 是 n 维整数空间上的一个向量,h 为该仿射超平面的法向量,c 为一维整数空间上的一个常数。
注意:很多读者可能读到仿射超平面的时候就开始读不懂了。实际上仿射超平面就是一个上面的一个
从公式1-2中不难发现,一个向量 v 被映射到哪个仿射超平面上是由 h · v 决定的,n 维空间上相互平行的超平面则对应 h · v 不同的值。对于所有满足 h · v = k 的向量 v ,称这些向量 v 构成一个仿射超平面,h 为该仿射超平面的法向量。当两个向量 v0 和 v1 满足 h · v0 = h · v1 时,称这 两个向量位于同一个超平面上。 上面的例子说明(3,1)和(4,0)在变换之后他们在同一个超平面上。
那么要计算如公式1-1那样的多维形式,Pluto算法就是计算一个形如
的多个仿射超平面或多个一维仿射变换的组合,其中 仿射函数
- Pluto算法的合法性约束
Pluto 算法实现的循环变换是对程序的重排序变换,所以只有在满足依赖的基本定理时,这样的循环变换是合法的。依赖的基本定理描述的一个基本事实是: 当优化编译器对程序实施重排序变换时,它仍然维持原始程序的依赖,这种维持关系可 以通过依赖源点和汇点之间的顺序关系来描述。对于程序中的任意一个依赖,假设其源点语句实例
的关系。一个仿射超平面必须是合法的,图就要在变换后维持这样的关系,那么就有
注意这两个公式中的符号变化,已经将公式1-4中的字典序大于等于
Pluto 算法要求解的是
其中,
注意:Farkas引理只有在
仍然考虑图1-1所示的例子,该代码迭代空间的三个依赖,分别可以用依赖距离向量(1,0)、(1,1) 和 (1,-1) 表示,并假设要求解的系数矩阵 M 和常量向量 C 如公式1-1所示。该循环嵌套中只包含一 个语句
即
显然,仅依靠调度变换的合法性约束无法计算出一个有效解,因为在这样的约束条件下,满足 目标优化问题的可行解有无限多个,因为合法性约束只限定了求解目标问题的下限。我们还要想办 法为所有未知变量
注意:仿射变换中的常量向量在实现循环合并和分布,所以在不涉及这些循环变换时,常量向量为零向量。
- 考虑并行性
Pluto 算法在实施循环分块后,还考虑了更粗粒度的分块之间的并行性。如果我们将一个循环分块看作是一个原子操作,以循环分块对应的原子操作构建一个新的迭代空间,可以得到如图1-7所示的迭代空间。图中每个黑色小方框表示图1-4或1-6中的一个绿色分块,蓝色箭头表示分块之间的依赖。 横坐标和纵坐标分别用
由于已经对循环迭代实施了循环分块,因此这个阶段就不需要考虑数据局部性了,除非需要实 现多级分块。在考虑分块之间的并行性时,Pluto 算法在 (
注意:这里的超平面是图1-7中的超平面,与前面的超平面不是一个坐标系下的超平面。也就是说Pluto算法需要做多次调度。
现在我们再将分块放回到 (t, i) 坐标系中分析循环分块。如图1-8是 (t, i) 坐标系中带有循环分 块编号的示意图,为了便于说明,我们没有显示原来的循环迭代和迭代之间的依赖。根据图1-7所示方式,我们可以得到如图1-9所示的并行实现方式。这里,我们假设图1-8中在同一水平方向上的所有分块由同一个处理器执行,并假设这两个处理器分别为 P0 和 P1。不难看出,分块之间的并行方式采用的是流水并行或称为波前方式。t = 0, 1 时处于流水填充阶段,由于分块 (1,1) 依赖于分块 (1,0),处理 器 P1 空载;t = 2, 3, 4 时,两个处理器同时工作,但每次 t 发生改变时,处理器之间需要进行数据的通信,以保证程序的正确性。例如,t = 3时,处理器 P1 计算分块(2,1),但该分块依赖于由 P0 计算的分块 (2,0),因此需要分块 (2,0) 向分块 (2,1)进行通信数据,图1-9中的蓝色箭头表示数据传输的方向,也表示分块之间的依赖。图1-8中分块 (2,0) 中红线上所有的数据的都要在 t = 3 之前通信给处理器 P1。t = 5, 6是流水并行的排空阶段,由于 P0 已计算完成,这个时间段内 P0 不参与计算。
当然,在实现这些分块之间的并行时,也可以不考虑分块之间的流水并行。在硬件充足的前提 下,可以沿图1-7水平方向
所以,分块的执行方式是只要平行四边形的斜边所在超平面的法向量在(1,1)和(1,-1) 所构成的扇面(cone)范围内,由此构成的分块形状都是合法的。但是,随着平行四边形的斜边越来越倾斜,无论是如图1-8所示考虑循环分块之间流水并行的情况还是图1-10所示采用更多处理器的方法,都无法最小化不同处理器之间的通信数据量。Pluto算法在计算新的语句实例之间的顺序关系时,将最小化通信数据量作为一种目标优化问题,并对该问题进行抽象构造了代价模型。
- Pluto算法的代价模型
事实上,Pluto 调度算法的代价模型非常简单,即构造一个仿射函数
该
给定一段由多个循环嵌套构成的程序片段。如果这些循环嵌套的迭代空间内的循环索引变量 都是有界的,那么
使得
成立。对该程序片段内所有依赖的源点
注意:这里有读者可能不太明白为什么会有这样的上界。实际上,程序中产生依赖的两个语句实例之间肯定会找到一个差函数,这个差函数肯定是和程序循环边界的参数相关的,在单个维度上最大的依赖距离肯定是该循环边界的倍数,否则这种依赖关系将不会有意义。所以,整个循环嵌套上依赖距离肯定是
当然,这个表达式也可以用Farkas引理实现线性化,即
到这为止,Pluto算法的最优化目标问题就已经构建完成了,其目的就是通过寻找
即这些变量按顺序排列之后的字典序最小值来求解调度过程。
注意:既然是字典序最小值,那么公式1-12中各个变量的排序就说明了Pluto算法对不同目标的侧重。Pluto 调度算法将行向量
- 求解方法
求解公式1-12中的目标问题仍然是一个复杂的过程,因为有可能会返回全0平凡解。为了排除这样的平凡解,Pluto算法中引入了
以消除全0平凡解,但该约束也同时将负数系数的解排除在外,所以Pluto算法中是不可以实现循环反转的,也不可以实现带负数系数的循环倾斜,这也就解释了为什么Pluto算法总是计算出图1-4形式的分块而不是图1-5中的分块形状。
现在,Pluto 算法可以使用整数线性规划求解工具来计算目标问题了。Pluto 调度算法针对每个语句计算出一个仿射变换,每个仿射变换都逐行求解。对于被 Pluto算法考虑的任何一个语句 S,整数线性规划求解工具至少要找到与该语句迭代空间维度
注意:这里要求传播这种约束是因为如果不传播前面已经计算出的仿射变换的解,后面在计算其它一维仿射变换或超平面的解的时候,由于考虑的约束和优化目标问题每次都相同,每次计算的一维仿射变换都是相同的解。这也是Pluto算法和Feautrier算法最大的不同。在本系列文章的第2部分内容中我们会再具体介绍。
更具体地,Pluto算法就是使用希尔伯特范式来传递这种前面已经计算出来的解的约束。根据公式1-3多维仿射变换的表达式,假设现在已经有前 r (1 ≤ r ≤ d) 行一维仿射变换已经被计算出来,我们用
给出。具体含义可参考Pluto算法的论文以及回顾希尔伯特范式的意义,这里就不再详细介绍了,读者可以把这个当作一个固定公式计算即可。 其中,
作为计算下一个一维方式变换的约束即可。
仍然考虑图1-1中的实例,假设已经计算出多维仿射变换矩阵的第一行,也就是第一行仿射变换为
删除
那么,在计算第二个一维仿射变换时,根据公式1-15,可以将
- 图1-1中的完整计算过程
首先,我们可以不用考虑公式1-1中的常数向量C,因为该例中只有一个语句,不会涉及到循环合并或分布。根据Pluto算法代价模型的上界,有
即
这里由于依赖距离向量的每个分量都是常数,所以不需要使用Farkas引理。接下来,根据公式1-13,有
所以第一个一维仿射变换要求解的就是
使得
其中,
接下来要求解的是第二个一维仿射变换。与前一个一维仿射变换的情况类似,Pluto算法要求解的是
使得
并得到上述优化问题的字典序最小解为
这个解就是图1-6中的坐标系的基。
注意:在我们的用例中并没有使用Farkas引理。关于使用了Farkas引理来计算调度过程的例子可以参考Pluto算法作者Uday Bondhugula的博士论文。另外,Farkas引理的例子我会在解释Feautrier算法时给出。如果仍然对Farkas引理的使用有兴趣的读者,可以等待本系列文章的第2部分内容。