Poly基本原理及卷积分析示例

**Poly基本原理介绍 **

考虑到许多读者可能对Poly并不了解，而且许多Poly文献读起来也比较抽象，我们先简单介绍一下Poly的工作原理。我们力图用最简单的代数与几何描述来解释Poly的基本原理。这部分内容参考了文献[12]的图片，我们通过解读这些图片来解释其中原理。

首先，我们假设有如图1所示的一段简单的循环嵌套，其中N为常数。循环嵌套内语句通过对A[i-1][j]和A[i][j-1]存储数据的引用来更新A[i][j]位置上的数据。如果我们把语句在循环内的每一次迭代实例抽象成空间上的一个点，那么我们可以构造一个以(i,j)为基的二维空间，如图2所示。图中每个黑色的点表示写入A[i][j]的一次语句的迭代实例，从而我们可以构造出一个所有黑色的点构成的一个矩形，这个矩形就可以看作是二维空间上的一个Polyhedron（多面体），这个空间称为该计算的迭代空间。

图1 一段简单的代码示例

图2 图1示例代码的迭代空间

我们可以用代数中的集合来对这个二维空间上的Polyhedron进行表示，即{[i, j] : 1 <= i <= N - 1 and 1 <= j <= N – 1}，其中[i, j]是一个二元组，“:”后面的不等式表示这个集合的区间。我们可以给这个二元组做一个命名，叫做S，表示一个语句，那么这个语句的Polyhedron就可以表示成{S[i, j] : 1 <= i <= N - 1 and 1 <= j <= N – 1}。

由于语句S是先迭代i循环再迭代j循环，因此我们可以给语句S定义一个调度（顺序），这个调度用映射表示，即{ S[i, j] -> [i, j] }，表示语句S[i, j] 先按i的顺序迭代再按照j的顺序迭代。

接下来，我们来分析语句和它访存的数组之间的关系，在代数中我们用映射来表示关系。图1中语句S对数组A进行读和写，那么我们可以用Poly来计算出S和A之间的读访存关系，可以表示成{ S[i, j] -> A[i - 1, j] : 1 <= i <= N -1 and 1 <= j <= N- 1; S[i, j] -> A[i, j - 1] : 1 <= i <= N - 1 and 1 <= j <= N -1 } 。同样地，写访存关系可以表示成{ S[i, j] -> A[i, j] : 1 <= i <= N - 1 and 1 <= j <= N -1 }。

基于这个读写访存关系，Poly就可以计算出这个循环嵌套内的依赖关系，这个依赖关系可以表示成另外一种映射关系，即{ S[i, j] -> S[i, 1 + j] : 1 <= i <= N - 1 and 1 <= j <=N - 2; S[i, j] -> S[i + 1, j] : 1 <= i <= N - 2 and 1 <= j <= N- 1 }。

可以注意到，Poly对程序的表示都是用集合和映射来完成的。当我们把语句实例之间的依赖关系用蓝色箭头表示在迭代空间内时，就可以得到如图3所示的形式。根据依赖的基本定理[13]，没有依赖关系的语句实例之间是可以并行执行的，而图中绿色带内（对角线上）的所有点之间没有依赖关系，所以这些点之间可以并行执行。但是我们发现这个二维空间的基是(i, j)，即对应i和j两层循环，无法标记可以并行的循环，因为这个绿色带与任何一根轴都不平行。所以Poly利用仿射变换把基(i, j)进行变换，使绿色带能与空间基的某根轴能够平行，这样轴对应的循环就能并行，所以我们可以将图3所示的空间转化成如图4所示的形式。

此时，语句S的调度就可以表示成{ S[i, j] -> [i + j, j]}的形式。所以Poly的变换过程也称为调度变换过程，而调度变换的过程就是变基过程、实现循环变换的过程。

图3 带依赖关系的迭代空间

图4 变基之后的迭代空间

图4中绿色带和j轴平行，这样在代码中表示起来就方便了。我们说Poly做循环变换的过程就是将基(i, j)变成(i + j, j)的一个过程，也就是说，Poly的底层原理就是求解一个系数矩阵，这个系数矩阵能够将向量(i, j)转换成向量(i + j, j)。

根据这样的调度，Poly就可以利用它的代码生成器，生成如图5所示的代码。此时，内层循环就可以并行了。（注：这里示意的是“源到源”翻译的Poly编译器，也就是Poly生成的代码还需要交给基础编译器如GCC、ICC、LLVM等编译成机器码才能运行。也有内嵌在基础编译中的Poly工具。）

图5 Poly变换后生成的代码

当然，我们这里举的例子是一个很简单的例子，在实际应用中还有很多复杂的情况要考虑。Poly几乎考虑了所有的循环变换，包括Interchange（交换）、Skewing/Shifting（倾斜/偏移）、Reversal（反转）、Tiling/Blocking（分块）、Stripe-mining、Fusion（合并）、Fission/Distribution（分布）、Peeling（剥离）、Unrolling（展开）、Unswitching、Index-set splitting、Coalescing/Linearization等，图6～8[14]中给出了几种Poly中实现的循环变换示意图，右上角的代码表示原输入循环嵌套，右下角的代码表示经过Poly变换后生成的代码。图中左边的集合和映射关系的含义分别为：J代表原程序语句的迭代空间，S表示输入程序时的调度，T表示目标调度，ST就是Poly要计算的调度变换。

图6 Poly中skewing变换示意图

图7 Poly中Fusion变换示意图

图8 Poly中Tiling变换示意图

深度学习应用的Poly优化

让我们以图9中所示的二维卷积运算（矩阵乘法）为例来简单介绍Poly是如何优化深度学习应用的。

图9 一个2D卷积示例

Poly会将循环嵌套内的计算抽象成一个语句。例如图9中S1语句表示卷积初始化，S2代表卷积归约；而S0和S3则分别可以看作卷积操作前后的一些操作，比如S0可以想象成是量化语句，而S3可以看作是卷积后的relu操作等。

为了便于理解，我们以CPU上的OpenMP程序为目标对图9中的示例进行变换。Poly在对这样的二维卷积运算进行变换的时候，会充分考虑程序的并行性和局部性。如果我们对变换后的程序并行性的要求大于局部性的要求，那么Poly会自动生成如图10所示的OpenMP代码；如果我们对局部性的要求高于并行性，那么Poly会自动生成如图11所示的OpenMP代码。（注：不同的Poly编译器生成的代码可能会因采用的调度算法、编译选项、代码生成方式等因素而不同。）

图10 Poly生成的OpenMP代码——并行性大于局部性

图11 Poly生成的OpenMP代码——局部性大于并行性

通过对比图10和图11，两种生成的代码采用的循环fusion（合并）策略不同：图10中所示的代码采用了({S0}, {S1, S2, S3})的合并策略，图11中生成的代码则使用了({S0,S1, S2, S3})的合并策略，但是必须通过对S2向右偏移99次、S3向右偏移148次，以及循环层次的interchange（交换）来实现这样的合并。显然，图11所示的代码局部性更好。而并行性上，仔细研究后不难发现，图11生成的代码中，只有最外层c0循环是可以并行的，而图10代码中，S0语句的c0、c1循环都可以并行，并且包含S1、S2、 S3三条语句的循环嵌套的c0、c1循环也都可以并行，相对于图11代码，图10生成的代码可并行循环的维度更多。

当然，在面向CPU生成OpenMP代码时，多维并行的优势没有那么明显，但是当目标架构包含多层并行硬件抽象时，图9中的代码能够更好地利用底层加速芯片。例如，当面向GPU生成CUDA代码时，而图10对应的CUDA代码（如图12所示）由于合并成了两个部分，会生成2个kernel，但是每个kernel内c0维度的循环被映射到GPU的线程块上，而c1维度的循环被映射到GPU的线程上；图11对应的CUDA代码（如图13所示）只有1个kernel，但是只有c0维度的循环被映射到GPU的线程块和线程两级并行抽象上。为了便于阅读，我们并未开启GPU上shared memory和private memory自动生成功能。从图中也不难发现，Poly也可以自动生成线程之间的同步语句。（注：图中循环分块大小为32，图12中线程块上线程布局为3216，图13中为324*4。）

图12 Poly生成的CUDA代码——并行性大于局部性

图13 Poly生成的CUDA代码——局部性大于并行性

值得注意的是，为了充分挖掘程序的并行性和局部性，Poly会自动计算出一些循环变换来实现有利于并行性和局部性的变换。例如，为了能够达到图11和图13中所有语句的合并，Poly会自动对S2和S3进行shifting（偏移）和interchange（交换）。