Polyhedral Model: DNN Nested Loop 的实现模型

循环是迭代空间的一个点

使用嵌套循环（Nested Loop）抽象不同的 DNN 乘加算子^[1]，使用多面体数学模型（Polyhedral Model）抽象循环的变换优化。

多面体模型里循环可以用迭代向量或者迭代点表示，我们以常见的 Linear Projection Layer 为例分析，该循环的迭代向量表示为 $(b,o,i)\in R^3$

For b in [0, B-1]:
For o in [0, O-1]:
For i in [0, I-1]:
y[b][o] += x[b][i] * w[o][i] // 循环语句 S

则循环被抽象为了迭代向量按某个次序排序，比如这个循环就是 $(b,o,i)$ 依次为 $(0,0,0),(0,0,1),...,(0,0,I-1),...,(B-1,O-1,I-1)$。多面体表示忽略了具体循环语句 S 的执行内容。

为了便于用仿射变换表示后续操作，我们使用向量的增广表示 $X=(b,o,i,1)\in R^4$（后文在涉及仿射变换时使用 $(b,o,i,1)$ 某些分析则使用 $(b,o,i)$ 不做特别区分）。

也有添加多个常数表示迭代向量的^[2]，比如 $(b,o,i,B,O,I,1)$，它们数学本质是一样的

迭代变量的定义域叫做迭代空间（Iteration Space），迭代空间是一组不等式组， e.g.

$$\begin{array}{}\text{(1)}& 0\le i\le I-1\text{(2)}& 0\le o\le O-1\text{(3)}& 0\le b\le B-1\end{array}$$

迭代空间可以用 $D(X)\ge 0$ 统一表示，特别的，当迭代空间是线性时，则可以用一个矩阵表示迭代空间，即 $D(X)=DX\ge 0$，e.g.

$$D=\left[\begin{array}{cccc}1& 0& 0& 0\\ -1& 0& 0& B-1\\ 0& 1& 0& 0\\ 0& -1& 0& O-1\\ 0& 0& 1& 0\\ 0& 0& -1& I-1\end{array}\right]$$

调度是迭代空间上超平面的序列

一个循环可以以多种方式执行，这称作调度（Scheduling）。比如可以先遍历 i $(0,0,0),(0,0,1),(0,0,2),...$ ，也可以先遍历 b $(0,0,0),(1,0,0),(2,0,0),...$ 。调度要在循环上刻下时间的维度，也可以说，调度是循环到时间的映射：

$$F(X)=t$$

对于迭代空间上的每个点，都能通过调度函数映射到时间步 $t$ (timestep)，t 代表了循环的执行次序，t 值越小的，执行越靠前，我们将空间里的各个点按 t 值排序，即将循环序列化。

$F(X)=t$ 本质是描述了一个超平面，因此我们也可以说调度是迭代空间上超平面的序列。e.g. 还是 Linear 层的循环，其对应的调度函数为 $F(X)=b+\frac{o}{B}+\frac{i}{B\times N}=t$。

特别的，满足线性时调度也能用调度变量表示，此时调度函数为调度向量和迭代向量的内积，即

$$F(X)=F\cdot X=t$$

e.g. $F=[1,\frac{1}{B},\frac{1}{B\times I},0]$

调度的超平面并非可以任意选取，要满足数据依赖性，比如以下的例子：

For i in [0, I-1]:
x[i+2] += x[i+1] * x[i] // 循环语句 S

显然，迭代点 $(0)$ 并不能在 $(1)$ 之后执行，也就是调度超平面满足 $ifF(X_1)<F(X_2)then{i}_1<i_2$。依赖性可以进一步扩展到更加复杂的条件约束，这规定了自动化算法的搜索空间。

循环变换

将算法抽象表示并不能直接得到硬件部署，我们要先能够映射到目标平台（Mapping），再进行优化。

不同的算法平台的硬件规格（比如空间计算的维度），执行算法范式（比如存算、脉动阵列）不同，要将循环和目标硬件对齐。

就拿存算举例，Linear 的 nested loop 表示每个循环语句单元是元素乘加，然而存算是向量乘加（向量内积），粒度比 nested loop 更高，且一部分累加要靠后续的累加器完成。

一种处理到存算计算范式的方法是循环分块（Tilling），即将一个循环切换为多个循环。e.g. 对 i 循环切为多个长度为 R 的小块，则每个小块代表执行了一次向量内积。

映射到存算计算范式的方法并不唯一，Tilling 是一种保持数据连续性的映射方法

For b in [0, B-1]:
For o in [0, O-1]:
For ir in [0, floor(I/R)-1]: 
For r in [0, R-1]:
y[b][o] += x[b][ir*R + r] * w[o][ir*R + r] // 循环语句 S

对于原循环，我们定义了迭代空间、约束条件，变换后的循环应当依然满足。比如迭代空间，很显然从新循环中能看出两个不等式 $0\le ir\le floor(I/R)-1$，$0\le r\le R-1$，再结合变换关系 $i=ir\ast R+r$ ，将此不等式带入原先迭代空间和约束条件之中，便能得到变换后的迭代空间和约束条件关系。

常见变换有，Fission, Fusion, Interchange, Tilling, Unrolling 等，其中一部分可以用仿射变换表示，其余不可（比如 Tilling 和 Unrolling）。

综上分析，循环的一部分子集可以用仿射变换表示，这有利于规范统一编译优化器，非仿射变换则要复杂得多。编译器有了初始值、有了迭代方法（变换方法）、有了性能评估（Evaluator），就可以开始迭代优化算法了。本文不讨论 Evaluator 和 Optimize 部分。

控制信号是 Trigger

假设我们通过优化得到某个调度方案 $F(X)$。我们还需要导出复杂的控制流，比如什么时候刷新权重？什么时候清空累加器？什么时候写出数据什么时候输入数据种种。而这都能在多面体的数学模型下规范描述。

控制信号本质是在时间上的脉冲序列：

$${\mathbb{1}}_T(F(X))=1,T\subset ran(F(X))$$

${\mathbb{1}}_T$ 是1函数，$T$ 是一个 $F(X)$ 的真子集（不会是子集，否则控制信号是一个常量，那么也不需要控制了）。这实际上是一个条件判断，当 $F(X)\in T$ 时，输出控制信号。$T$ 也可以按前面的用线性或非线性的表示，具体过程不展开。

---

1. Parashar, A., Raina, P., Shao, Y. S., Chen, Y.-H., Ying, V. A., Mukkara, A., Venkatesan, R., Khailany, B., Keckler, S. W., & Emer, J. (2019). Timeloop: A Systematic Approach to DNN Accelerator Evaluation. 2019 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS), 304–315. https://doi.org/10.1109/ISPASS.2019.00042 ↩︎

2. 《可重构计算》，魏少军，刘雷波，尹首一 ↩︎