Red-blue Pebble Game

Red-blue Pebble Game是一种两级内存访问模型，被用于分析程序的I/O复杂度。该模型建立在一个有向无环图（DAG）上，其中某些节点具有红色的石子，有些节点具有蓝色的石子，有些节点具有两种，还有些节点两种都不具有。该模型有以下四个规则：

红色石子可以放在任何一个有蓝色石子的节点。
蓝色石子可以放在任何一个有红色石子的节点。
如果一个节点的多有前驱都有红色石子，那该节点可以放置一个红色石子
一个石子可以重任何一个节点上移除

几个定义如下：

configuration（构型）：一对仅包含具有红色石子或具有蓝色石子的顶点的子集，具有两种石子的节点属于两个子集的交集。
input（输入）：没有前驱的节点的子集。
output（输出）：没有后继的节点的子集。
initial（初始）：仅有蓝色石子的输入节点。
terminal（终止）：仅有蓝色节点的终止节点。
transition（推导）：一对有序的configuration，后者是由前者根据某一条规则得来的。
calculation（计算）：一系列configuration，一个完整的计算表示由最初的（initial）构型到最终的（terminal）构型。

该模型完整的推到案例见 DPHPC: Red
-Blue Pebble Game。该模型可以模拟在快慢两级内存（fast and slow memory）模型上的计算，fast memory容量较小，如寄存器；slow memory无限制，如DRAM。顶点表示一个操作的结果，有向边表示前驱的结果是后继的操作数。只有操作数都在fast memory中才能被执行，与规则3相联立。规则2表示将一个结果从fast memory存到slow memory，规则1表示将一个结果从slow memory加载到fast memory。DAG中的红蓝石子总数表示fast memory和slow memory可以利用的大小。因此，我们定义最少的I/O次数为一个完整的计算中的最少推导次数。

文章基于Red-blue Peddle Game提出了S-partition的有向无环分割模型，具体描述如下：

$G(V,E)$ 表示一个DAG，$V$ 表示顶点集合，$E$ 表示边的集合，当对G的一种分割满足以下四个条件时，成为S-partition：

$V$ 被分成 $h$ 个子集 $V_1, V_2, …,V_h$，各个子集没有交集，且它们的合集是 $V$。

对于任意一个 $V_i$，都有一个支配集 $D_i$，对于 $V_i$ 以外的任何一个顶点都存在 $D_i$ 中的一个顶点，使两顶点的边在 $E$ 中。$D_i$ 包含顶点数最多为 $S$。

对于任何一个 $V_i$，都有一个 $M_i$，对于 $M_i$ 中的任何一个顶点，都不存在后继节点属于 $V_i$。$M_i$ 包含定点数最多为 $S$。

各个子集 $V_1, V_2, …,V_h$ 不存在循环依赖性。（如果存在一条边连接 $V_j$ 中的某顶点到 $V_i$ 中的某个顶点，则 $V_i$ 依赖于 $V_j$）。

任一DAG经过S分割后都有一个子集个数。令 $P(S)$ 表示任何S分割的子集数的最小值。给定一个容量为S的fast memory和一个可以用DAG描述的算法，两级内存之间的最小I/O时间Q满足

$\begin{equation} Q \geq S \cdot (P(2S) - 1) \end{equation}$

Challenges

对于winograd快速卷积算法这类多步骤的算法，不适合使用DAG推断合适的I/O下界，原因如下：

winograd每个步骤都需要在输入执行加载，在输出执行存储，而在一个完整的DAG上，数据可以通过fast memory从前一个子计算传递到后一个子计算。
后继步骤必须要在前一个步骤全部完成后才能进行，且fast memory大小有限，因此总的DAG分割方法会影响数据移动的复杂性。

另一个难点在于winograd算法的内存访问、线程模式、特定输入尺寸的组合存在巨大的配置空间，很难利用一个I/O边界实现最优的自动调优算法。

Lower Bound Theory

Red-Blue Pebble Game Re-exploration

Basic Idea

令 $\mathbb{P}_S$ 为包含DAG $G(V, E)$ 所有S分割后的顶点子集集合，令

$\begin{equation} H(S) = \min\limits_{\{V_1,\cdots,V_h\}\in \mathbb{P}_S}\frac{|V|}{\mathop{max}_{1\leq i \leq h}|V_i|} \end{equation}$

对于任何一个S分割， $\frac{|V|}{\mathop{max}|V_i|}$ 是子集个数的下界，因此有 $P(S) \geq H(S)$，即得

$\begin{equation} Q \geq S \cdot (P(2S)-1) \geq S \cdot (H(2S)-1) \end{equation}$

由此可以通过求 $H(2S)$ 来得到 $Q$ 得下界，而 $H(S)$ 又由 $|V_i|$ 求得，接下来将分析如何求 $|V_i|$ 得上界。

对于一个多步骤算法的DAG图，有如下定义：

定义1：设 $G(V,E)$ 被分解为 $n$ 个子图 $G_1(U_1,E_1),G(U_2,E_2), \cdots ,G_n(U_n, E_n)$，集合 $\{G_1(U_1,E_1), \cdots ,G_n(U_n, E_n)\}$ 称为 $G(V,E)$ 的多步分割（multi-step partition），当且仅当 $G_j(U_j, E_j)$ 的输入节点是 $G_{j-1}(U_{j-1},E_{j-1})$ 的输出节点，且所有节点集合 $U_j$ 互不相交。

要求 $|V_i|$ 的最大值，我们可以先求出 $|V_i \cap U_j|\;(j=1,2,\cdots,n)$ 的上界。假设我们已经知道了 $|V_i \cap U_j|$ 的上界，为了推导出 $|V_i \cap U_{j+1}|$ 的上界，我们应当找到 $U_j$ 和第 $j+1$ 个计算步骤的关系。令 $\widetilde{O}_j$ 为 $U_j$ 的输出集合，$D_i$ 为 $V_i$ 的支配集，引入概念“顶点生成”（vertex generation）表示 $\widetilde{O}_j$ 中和 $V_i \cap U_{j+1}$ 相连的顶点，具体定义为：

定义2：在一个DAG $G(V, E)$ 中，一个顶点集 $U$ 可以生成 $U’$，当且仅当从 $V$ 的输入顶点到 $U’$ 中的一个顶点的路径中包含 $U$ 的顶点。$\Theta(U)$ 表示所有由 $U$ 生成的节点集合。

由以上定义可知，支配集 $D_i$ 可以生成 $V_i$。$V_i \cap U_{j+1}$ 的所有输入节点都属于集合 $\Theta(D_i) \cap V_i \cap \widetilde{O}_j$，因此如果能够求出 $|V_i \cap U_j|$ 和 $|\Theta(D_i) \cap V_i \cap \widetilde{O}_j|$ 的上界，就可得到 $|V_i \cap U_{j+1}|$ 和 $|\Theta(D_i) \cap V_i \cap \widetilde{O}_{j+1}|$。依此递归推导，可以求得所有 $|V_i \cap U_{j+1}|(j=1,2,\cdots,n)$ 的上界，从而求得 $|V_i|$ 的上界。

文章举例介绍了如何求得 $|V_i|$ 的上界：假设一个算法包括两个子步骤，即 $G(U,E)=G_1(U_1,E_1)+G_2(U_2,E_2)$。定义 $k^i$ 表示 $V_i$ 的支配集 $D_i$ 的顶点个数，即 $|D_i| = k^i \leq S$。现令 $k^i = k_1^i + k_2^i$，其中 $k_1^i$ 和 $k_2^i$ 分别表示 $V_i \cap U_1$ 和 $V_i \cap U_2$ 的输入顶点的个数。令 $\varphi_j(k)$ 和 $\psi_j(k)$ 分别表示 $U_j$ 和 $\widetilde{O}_j$ 中由 $k$ 个输入顶点生成的顶点个数的最大值。因此有 $|\Theta(D_j) \cap V_i \cap U_1| \leq \varphi_1(k_1^i)$，并且最多有 $\psi_1(k_1^i)$ 个顶点作为 $V_i \cap U_2$ 的输入，即 $V_i \cap U_2$ 的输入顶点个数最多为 $k_2^i + \psi_1(k_1^i)$。因此可以推导出 $|\Theta(D_i) \cap V_i \cap U_2| \leq \varphi_2(k^i_2 + \psi_1(k_1^i))$。最终可以求出

$\begin{aligned} |V_i| &\leq |D_i| + |\Theta(D_i) \cap V_i \cap U_1| + |\Theta(D_i) \cap V_i \cap U_2| \\ &\leq S + \varphi(k_1^i) +\varphi_2(k^i_2 + \psi_1(k_1^i)) \\ &\leq S + \max \limits _{k_1 + k_2 \leq S} (\varphi_1(k_1) + \varphi_2(k_2 + \psi_1(k_1))) \end{aligned}$

Maximum Vertex Generation Functions

对任意整数 $k$、顶点集 $U$ 和支配集 $D$，满足 $|D \cap U_j| + |\Theta(D) \cap \widetilde{O}_{j-1}| \leq k$，定义如下两个顶点生成方程：

$\widetilde{\varphi}_j(U,k)=|\Theta(D) \cap U \cap U_j| \\ \widetilde{\psi}_j(U,k)=|\Theta(D) \cap U \cap \widetilde{O}_j|$

其中 $\widetilde{\varphi}$ 和 $\widetilde{\psi}$ 分别表示顶点集 $U \cap U_j$ 和 $U \cap \widetilde{O}_j$ 中由 $D$ 生成的顶点数量。对于任意顶点 $k$ 和第 $j$ 个子计算，定义最大顶点生成函数为

$\begin{equation} \varphi _j(k) = \max \limits _U \widetilde{\varphi}_j(U,K)\;\text{and}\;\psi _j(k) = \min \limits _U \widetilde{\varphi}_j(U,K) \end{equation}$

$\varphi _j$ 和 $\psi _j$ 分别估计了 $|U_j|$ 和 $|\widetilde{O}_j|$ 的上界。

Estimation of Upper Bound of $|V_i|$

令 $\widetilde{O}_j ^i$ 表示 $\widetilde{O}_j$ 的一个子集，满足对于任意 $v \in \widetilde{O}_j$，由 $V$ 的输入集合到 $v$ 的任意路径中都含有一个属于集合 $\cup _{k=1} ^j(D_i \cap U_k)$ 的顶点。

引理3：$\widetilde{O}_j ^i\; \cup \; (D_i \cap U_{j+1})$ 是 $\widetilde{O}_{j+1} ^i$ 的一个支配集。

引理4：$\widetilde{O}_j ^i\; \cup \; (D_i \cap U_{j+1})$ 是 $V_i \cap U_{j+1}$ 的一个支配集。

依据上述两条引理，可以得出 $|V_i|$ 的上界：

定理5：假设 $\{G_1{U_1,E_1}\},\cdots,G_n(U_n,E_n)$ 是一个DAG图 $G(V,E)$ 的多步分割（多步骤算法的分割），对于 $G(V,E)$ 任意S分割 $\{V_1,\cdots,V_h\}$，$|V_i|$ 存在上限
$\begin{equation} \begin{aligned} T(S) = S+ \max \limits _{\sum_{j=1}^n k_j \leq S}(\varphi_1(k_1)+\varphi_2(k_2+\psi_1(k_1))+\cdots\;\cdots \\ + \varphi_n(k_n+\psi_{n-1}(k_{n-1}+\psi_{n-2}(k_{n-2} \cdots + \psi_1(k_1))))) \end{aligned} \end{equation}$

I/O Lower Bounds of Composite Algorithms

定理6：对于一个描述多步骤算法的DAG图 $G(V,E)$，当fast memory大小为 $S$ 时，fast memory和slow memory之间的最少I/O次数 $Q$ 为
$\begin{equation} Q \geq S \; \cdot \; (\frac{|V|}{T(2S)}-1) \end{equation}$

根据式（6），求一个多步骤算法的I/O下界可分为以下三个步骤：

求 $G(U,V)$ 中顶点的个数 $|V|$。
求 $G(U,V)$ 的第 $j$ 个多步分割的 $\varphi _j$ 和 $\psi _j$，并且由式（5）求出 $T(S)$。
由式（6）求出I/O下界。

依据以上三个步骤，文章求了直接卷积和winograd卷积算法的I/O下界。

I/O Lower Bound of Direct Convolution

直接卷积包含两个步骤，DAG图 $G(V,E)=G_1(U_1,E_1) \cup G_2(U_2,E_2)$。定义 $I_i$ 表示第 $i$ 个大小为 $W_{ker} \times H_{ker} \times C_{in}$ 的输入张量，$K_j$ 表示第 $j$ 个大小为 $W_{ker} \times H_{ker} \times C_{in}$ 的卷积核。直接卷积的第一个步骤是将 $I_i$ 和 $K_j$ 两者相乘，生成 $W_{ker}H_{ker}C_{in}$ 个乘积项。第二步将乘积项相加，形成一颗加法树，最后得到一个输出顶点。

Alt text

每一个顶点都表示一个数，输入顶点的个数就是输入张量和卷积核的尺寸，乘积项和加法树中顶点的个数分别为加法和乘法的次数，输出顶点的个数就是输出张量的尺寸。因此在表示直接卷积的DAG图中，顶点的个数为

$\begin{equation} \begin{aligned} |V|=(2W_{ker}H_{ker}C_{in}-1)W_{out}H_{out}C_{out}\\ +W_{in}H_{in}C_{in}+W_{ker}H_{ker}C_{in}C_{out} \end{aligned} \end{equation}$

一个卷积核窗口下输入张量的可重用元素个数 $R=W_{ker}H_{ker}/\mu ^2$，其中 $\mu$ 表示卷积操作的步长。对于直接卷积的两个步骤，可以得出 $\psi _1 = \varphi _1$，$\varphi _1(k_1) \leq 2S \sqrt{Rk_1}$，$\varphi_2(k_2) \leq k_2 - 1$，因此可以求出 $T(S) \leq 4S \sqrt{RS} + S - 1$。

根据式（6）（7），文章建立了直接卷积（DC）的I/O下界：

$\begin{equation} Q_{lower\; DC} = \Omega \; (\frac{W_{ker}H_{ker}C_{in}W_{out}H_{out}C_{out}}{4 \sqrt{2RS}}) \end{equation}$

I/O Lower Bound of Winograd Algorithm

在winograd算法中，三个转换矩阵尺寸很小，可以假设它们都在fast memory中且不会占用空间。令卷积核的尺寸为 $r=W_{ker}=H_{ker}$，在表示winograd算法的DAG中，顶点的个数为

$\begin{equation} |V|=O\;(\frac{W_{out}H_{out}C_{out}C_{in}(e+r-1)^4}{e^2}) \end{equation}$

对于winograd算法中的四个步骤，可以推断出每一个步骤的 $\varphi _j$ 和 $\psi _j$ 的上界如下：

$\begin{aligned} &\varphi_1(k_1) \leq 6k_1(e+r-1)^4/er \\ &\psi_1(k_1) \leq 3k_1(e+r-1)^2/er \\ &\varphi _2(k_2) \leq k_2 \sqrt{k_2} + (e + r -1)^2 S \sqrt{k_2} / e^2 \\ &\psi_2 = \varphi2 \\ &\varphi_3(k_3) \leq k_3 - 1 \\ &\psi_3(k_3) \leq \min \{k_3/2,\; S(e+r-1)^2/e^2\} \\ &\varphi _4(k_4) \leq \min \{ (2k_4-1)e^2, \; (2(e+r-1)^2 - 1)S \} \end{aligned}$

因此可求出winograd算法的 $T(S)$ 和I/O下界 $Q$:

$\begin{equation} T(S) = O\;(2\frac{(e+r-1)^3}{er}S \sqrt{S} + \frac{6(e+r-1)^2}{er}S) \end{equation}$ $\begin{equation} Q_{lower \; WA} = \Omega\;(\frac{W_{out}H_{out}C_{out}C_{in}(e+r-1)r}{e \sqrt{S}}) \end{equation}$

可以发现，当式（8）中 $R=1$，式（11）中 $e=1$，直接卷积和winograd卷积的I/O下界和矩阵乘的相似，都与 $\sqrt{S}$ 成反比。

Near I/O-optimal Strategy

根据式（5）（6），我们可以通过找到I/O下界的最高次项来判断算法哪个部分应当数据重用。在直接卷积中，$\varphi _2$ 为最高次项，表示最小化I/O次数应当最大化输出数据的重用；在winograd算法中，$\varphi _3$ 为最高次项，表示应当最大化第三部中两个临时数组的数据重用。

Dataflow Design for Direct Convolution

从前文的分析得知，直接卷积应当最大化第二部输出数据的重用，因此应将大部分的片上内存分配给输出数据。

Alt text

上图展示了一个输出分块中的三个维度，为了减少片外存储访问，应当使 $xyz \approx S/N_p$，$N_p$ 表示处理器个数。由于片上存储空间有限，为了最大化输出数据所占空间，应当分次将input和kernel的子块读入计算，每一次将加载 $x’ \times y’ \times \alpha$ 大小的input块和 $z$ 个大小为 $H_{ker} \times W_{ker} \times \alpha$ 的kernel子块。又因为每个通道的input块只能被同一通道的卷积核重用，因此 $\alpha = 1$，一个大小为 $x’ \times y’$ 的input tile将在一个通道内被循环加载。

对于output得每一个子块，读取数据得I/O开销来自将相应的input子块和kernel子块加载。当 $R=W_{ker}H_{ker}/ \mu ^2$，$x’ \approx \mu x$，$y’ \approx \mu y$，由基本不等式得：

$\begin{equation} \begin{aligned} Q_{DC\;reading} \approx \frac{H_{out}W_{out}C_{out}}{xyz}(H_{ker}W_{ker}C_{in}(z + \frac{xy}{R})) \\ \geq H_{out}W_{out}C_{out}H_{ker}W_{ker}C_{in}(2 \sqrt{\frac{1}{Rxyz}}) \end{aligned} \end{equation}$

当前仅当 $xy=Rz$ 时，读取的I/O可以达到下界。此时可推导出 $x’y’=zW_{ker}H_{ker}$，即一个input tile的最优大小。

在 $xyz \approx S/N_p$ 且 $xy=Rz$ 时，总的读写I/O下界为：

$\begin{equation} Q_{DC} \approx \frac{2H_{out}W_{out}C_{out}H_{ker}W_{ker}C_{in}}{\sqrt{RS/N_p}} + H_{out}W_{out}C_{out} \end{equation}$

通过上式，可以发现：

在单核情况下，即 $N_p=1$，如果 $\frac{H_{ker}W_{ker}C_{in}}{\sqrt{RS}} \geq 1$（各类CNN参数配置均满足改条件），式（13）与式（8）一致，这也印证了将大部分片上存储分配给输出数据是有效的。
在多核下，要想并行化数据移动，应当充分利用片上存储计算部分和（提高 $R$）.

综上，为了降低I/O开销，应当注意几个算法设计细节：

充分考虑input image和kernel的重用。
充分考虑output的重用。每个input tile将沿 $C_{in}$ 维度循环，单个通道下的计算结果都只会在最后一次被写回，因此要沿 $H_{in}$ 和 $W_{in}$ 两个维度加载input。
为了达到I/O下界，应满足 $xy=Rz$, $x’y’=zW_{ker}H_{ker}$。

Dataflow Design for Winograd Algorithm

从前文分析得知，winograd算法应当最大化第三步的临时数据的重用。

Alt text

和直接卷积的设计同理，单个input tile通道上循环加载计算。为了最大化两个临时数组 $\Lambda$ 和 $\Pi$ 的重用，我们将利用大部分的片上存储空间来存储 $2xyz/e ^2$ 个临时数组，即满足 $2\frac{(e+r-1)^2}{e^2}xyz \approx S/N_p$。

当计算output的每一个子块，需要加载大小为 $x’y’C_{in}$ 的input子块和 $z$ 个大小为 $r^2C_{in}$ 的kernel。在winograd中，卷积步长 $\mu=1$，因此有 $x’ \approx y$，$y’ \approx y$，算法总的读取数据的开销为

$\begin{equation} \begin{aligned} Q_{WA\;reading} \approx \frac{H_{out}W_{out}C_{out}}{xyz}(xyC_{in} + zr^2C_{in}) \\ \geq H_{out}W_{out}C_{out}C_{in}(\frac{r}{\sqrt{xyz}}) \end{aligned} \end{equation}$

当且仅当 $xy=zr^2$ 时达到I/O下界。在winograd中，$R=r^2$，因此有 $xy=Rz$，这与直接卷积相同。当 $2\frac{(e+r-1)^2}{e^2}xyz \approx S/N_p$ 时，总的读写I/O下界为

$\begin{equation} Q_{WA} \approx \frac{2H_{out}W_{out}C_{out}C_{in}r(e+r-1)}{e\sqrt{S/N_p}} + H_{out}W_{out}C_{out} \end{equation}$

综上，设计winograd算法数据流的细节如下：

winograd算法重点关注第三步两个临时数组的重用。
提高变量的计算并行性。输出通道上的每个 $e \times e$ 块的更新是并行执行的。为了实现高并行性和数据重用，大部分片上内存用于加载临时数组。