多面体编译技术在软硬协同设计中的应用

摘要

多面体模型是一项成熟的编译优化技术，广泛地应用在传统的编译器及AI编译器中。本文关注的是将多面体技术应用到软硬协同设计领域的空间探索(Design Space Exploration)，针对Systolic Array的硬件架构设计多面体的分析模型，借助于现有的多面体编译技术在编译的过程中实现对硬件架构的探索，达到最佳性能；也可以在特定硬件约束下，改变算法设计适配硬件。

多面体模型介绍

多面体模型(Polyhedral)主要关注的是程序设计中的循环优化，两层循环的循环变量的取值范围可以构成一个平面，三层循环的循环变量可以组成一个长方体，如图1所示，因此得名多面体模型。

图1 不同循环层次的多面体表示[1]

多面体编译优化关注的是在确保程序执行正确的前提下重组循环的结构，实现性能最佳。比如图2的循环中，左图表示的原始的二维迭代空间，蓝色箭头表示数据（黑点）之间存在依赖关系，对角线的绿色表示数据没有依赖关系，经过变基操作之后变为右图的表达式及迭代空间，从形状看像是把多面体进行了变形，形象地体现出多面体优化的过程。当然，变形的目的是为了实现并行计算，达到更好的性能，具体分析可以参考[1][2]。

图2 多面体模型的原始循环空间和变基空间[1]

脉动阵列（Systolic Array）介绍

脉动阵列是基于数据流的计算架构，具有十分规则的计算范式，对于两个大小为4x4的矩阵乘投射到一个二维的4x4的脉动阵列上，其计算架构和数据流动可以表示为下图。本文的分析也针对二维脉动阵列展开。

图3 脉动阵列的数据流表示

图中的脉动阵列的特点是数据沿着水平和垂直方向流动，在计算单元(PE：Process Element)中完成乘加操作。

脉动阵列的空时转换模型

针对脉动阵列的空时变换(Space-Time Transformation)的分析模型在上世纪九十年代已经提出[3]。

定义一个二维脉动阵列的大小为

，假设两个矩阵A和B，大小分别为

和

，对于矩阵乘的操作表达为

，该公式的对应的伪代码为

图4 矩阵乘的伪代码表示

基于等价计算的方法，将上述代码的表达转换为自由赋值的符号表达（Assignment-Free Notation）

Eq<1>表示的是输入参数，Eq<2>表示的是计算过程，Eq<3>表示的是输出结果。

针对符号表达，可以将其投射到多面体模型中，

表示迭代的变量（Iteration Variables），他们组成了迭代空间(Iteration Domain)，

每个变量（或称之为Node)对应原始代码表现形式的一个计算实体(Instance)

。在上述第二个公式Eq<2>中

组成了迭代矩阵(Iteration Vector)。为了上述三个公式中的变量统一，将迭代变量k引入到Eq<1>和Eq<3>中。

脉冲阵列的计算可以从空间域和时间域两个维度进行描述，空间域表示每个计算落在脉动阵列的计算单元的位置，时间域表示计算结果输出的时刻。通过观察，在图3的脉动阵列中，第一个计算结果

的输出时刻为t=4，最后一个计算结果的输出的时刻是t=N1+N2+N3。N3可以看成是数据传输延迟，N1和N2是脉动阵列固有的延迟。

将前文提到的迭代矩阵

表示为向量

，定义一个时间向量(Time Vector)

，那么计算实体

在脉动阵列上的发生的时刻t表示为：

为了计算PE的空间坐标

，定义一个投影矩阵(Project Matrix)

，如下

对于一个计算实体

将会映射到具体的PE的坐标是

，计算过程如下

另外，投影方向可以使用任何垂直于投影矩阵的单位向量

表示，其中

注意，投影矩阵和投影方向向量不是唯一的，此处是选择了

。投影方向向量指定了具体的映射方向，且沿着该向量方向的节点投影到相同的PE上。

根据以上关于时间t和空间z的公式表达，那么就可以将脉动阵列的计算过程通过空时变换(Space-Time Transformation) T描述如下，T为投影矩阵和时间向量的组合。

针对图3的脉动阵列架构，空时转换矩阵为

从公式表达上，可以看出

是一个N-1维的矩阵，描述了在计算实体映射到的PE的坐标，

为一维向量，描述了每个计算实体在所投射的PE上的计算时间（执行顺序）。

多面体模型的构建

本节的分析根据论文PolySA [4]展开，论文中多面体模型定义为

。

为迭代空间；迭代矩阵为

；散射函数(Scattering Function)

对应为上述的空时变换矩阵

，散射函数就是空间映射(Space Mapping)和时间调度(Timing Scheduling)的表达；访问函数（Access Function）

，表示为对内存的访问。

逻辑标签（Logical Stamp）S表示的是空间域信息和时间域信息。

针对脉动阵列的空间探索，散射函数的不同决定了产生不同的脉动阵列架构，而散射函数是由投影矩阵P和时间向量

决定的，P由投影向量

决定（见上节的分析），因此空间探索的问题变为遍历出符合要求的

和

。

根据脉动阵列结构的自身特点，一是每个边上的计算单元存在一个周期的延迟（脉动阵列本身固有特点），因此一个边上的计算时间都需要大于0，

表示位于边上的迭代向量。

二是映射到一个PE的节点无法在同一时刻计算，即前面提到的投影方向向量和时间向量的乘积要大于0。

散射函数存在传递性，可以将其不断投射下去。将多面体模型

定义为N维数组

，根据散射公式转化为N-1的数组

，并类推下去。

比如从

投射到

，投射函数定义为

多面体

的逻辑标签S表示为

脉动阵列空间探索实例

图5 矩阵乘(I=J=K=2)的空间探索分析

图5中的图b和图c是针对2x2的脉动阵列，图d是将矩阵操作投射到1x2的脉动阵列上。

图5中a表示的是多面体的散射函数和访存函数，对于矩阵A，它被访问的迭代变量是i或者k，因此它的访存函数为

图b表示的是沿着k轴

（此处采用

表示投影方向，和前文的

一样）进行投射，那么在i和j轴相同的迭代向量

和

将投射到PE00，

和

将投射到PE11。和时间向量轴垂直的平面的迭代向量在同一时刻输出结果，比如计算实体

，

同时在t=1时刻得到输出结果。

图c表示的也是K轴的映射，空间映射的结构和图b相同，但时间向量选择是

,计算实体

，

结果在同一时刻输出。

图d表示将图c的空时信息进行再一次的投射，改变脉动阵列的结构。新的投射矩阵变为了向量

，时间向量为

，投射方向为

。

构建了多面体的中间表示(IR)之后，可以利用现有的多面体编译技术搭建编译框架，如图6所示。Design Optimizer是针对FPGA上的资源的优化，比如DSP，RAM，LUT等。最终产生HLS代码的配置参数，通过模板的方式生成HLS IP。

图6 论文[4]PolySA的编译框架

软硬协同设计的可能性

1.上述方法将矩阵乘法操作通过多面体的表达，将其映射到脉动阵列的硬件架构上，通过控制不同的投射方向

和时间向量

，得到不同的空间映射和时间表达，二者也形成了空间探索的变量。在特定脉动阵列架构的约束下，可以获得符合约束的投影方向向量

和时间向量

，如果引入具体的内存访问模型，比如对内存层级(Memory Hierarchy)的描述，就可以很好地评估出性能，比如功耗，根据性能选择合适的

和

，从而完成特定硬件架构下最优性能的探索。

2.满足特定的算法性能要求下，尽可能地降低硬件设计的开销，比如降低脉动阵列的维度，可以通过建立多级的散射关系进行空间探索，将一种算法实现投射到最小的硬件结构上。

3.将编译技术引入到不同硬件需求的空间探索中。针对特定架构的硬件，如果追求性能最优，在编译的环节调出性能最优的映射方法；如果追求功耗最低，那么可以调出符合功耗要求的最佳映射方法；如果底层是可配置的硬件，那么从软件和硬件两个维度调优，结合性能和功耗的双重考虑。

结论与思考

本文讨论了针对脉动阵列架构构建多面体模型的分析方法，并将多面体编译技术应用到硬件架构空间的探索中。该方法的启示在于将编译技术融合到软硬协同设计中，在编译的过程中针对特定的需求进行软硬件两个维度的空间探索，达到协同调优。但是目前的多面体分析模型局限于脉动阵列这种具有规则性的硬件架构上，尚且不能映射到所有的架构，这是需要后期研究的工作。

参考文献

[1] Uday Bondhugula, Polyhedral Compilation Opportunities in MLIR

[2] 要术甲杰, Polyhedral Model—AI芯片软硬件优化利器,

https://mp.weixin.qq.com/s/QEooKxP1sm5O90AUiqKQEQ

[3] Space-time transformation and systolic arrays, https://gyires.inf.unideb.hu/KMITT/a52/ch04s02.html

[4] Jason Cong and Jie Wang. PolySA: polyhedral-based systolic array auto-compilation. In ICCAD 2018.

往期推荐

1、深度生成网络新思路：扩散概率模型

2、基于MLIR实现GEMM编译优化

3、机器学习方法在数字芯片设计中的应用