EE290-2 LAB3 tiling and optimization for accelerator

Mapping

systolic array为我们提供了一种很好的并行化方法，然而在实际运行中可能存在一些问题：比如systolic array的尺寸过小，无法载入整个权重等。这时就需要插入一些mapping的技巧。

我们先贴出原版的systolic array代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


for (m = 0; m < M; m++)
{
    spatial_for(n = 0; n < N; n++)
    {
        OA[n, m] = 0;
        spatial_for(k = 0; k < K; k++)
        {
            OA[n, m] += IA[m, k] * W[k, n];
        }
    }
}

case 1: systolic array的尺寸小于K or N

这意味着weight无法一次性载入到整个systolic array中，以下为分块载入的代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


for (m = 0; m < M; m++)
{
    for (n1 = 0; n1 < N1; n1++)
    {
        OA [n1 * N0:(n1 + 1) * N0, m] = 0;
        for (k1 = 0; k1 < K1; k1++)
        {
            spatial_for(n0 = 0; n0 < N0; n0++)
            {
                spatial_for(k0 = 0; k0 < K0; k0++)
                {
                    OA[n1 * N0 + n0, m] += IA[m, k1 * K0 + k0] * W[k1 * K0 + k0, n1 * N0 + n0];
                }
            }
        }
    }
}

这里N0*N1=N，K0*K1=K，systolic array的尺寸大小为N0*K0，取块的方式如下：

/img-20230116171209665

case 2：weight buffer小于systolic array的size，也就是说，weight不能被一次性载入到systolic array，下为代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


for (m = 0; m < M; m++)
{
    // IA buffer stores: 1*K
    mvin(IA [m:m + 1, 0:K]);
    for (n2 = 0; n2 < N2; n2++)
    {
        // W buffer stores: N1*N0*K < NK
        mvin(W [0:K, n2 * N1 * N0:(n2 + 1) * N1 * N0]);
        OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m:m + 1] = 0;
        for (n1 = 0; n1 < N1; n1++)
        {
            for (k1 = 0; k1 < K1; k1++)
            {
                spatial_for(n0 = 0; n0 < N0; n0++)
                {
                    spatial_for(k0 = 0; k0 < K0; k0++)
                    {
                        OA[n2 * N1 * N0 + n1 * N0 + n0, m] += IA[m, k1 * K1 + k0] * W[k1 * K0 + k0, n2 * N1 * N0 + n1 * N0 + n0];
                    }
                }
            }
        }
        mvout(OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m:m + 1]);
    }
}

这里只考虑了weight buffer在N维度上不够的情况，即N=N0*N1*N2。映射方式如下：

case 3：input buffer小于systolic array的size，也就是说，input不能被一次性载入到systolic array，下为代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


for (m2 = 0; m2 < M2; m2++)
{
    // IA buffer stores: M1*K
    mvin(IA [m2 * M1:(m2 + 1) * M1, 0:K]);
    for (n2 = 0; n2 < N2; n2++)
    {
        // W buffer stores: N1*N0*K
        mvin(W [0:K, n2 * N1 * N0:(n2 + 1) * N1 * N0]);
        OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m2 * M1:(m2 + 1) * M1] = 0;
        for (m1 = 0; m1 < M1; m1++)
        {
            for (n1 = 0; n1 < N1; n1++)
            {
                for (k1 = 0; k1 < K1; k1++)
                {
                    spatial_for(n0 = 0; n0 < N0; n0++)
                    {
                        spatial_for(k0 = 0; k0 < K0; k0++)
                        {
                            OA[n2 * N1 * N0 + n1 * N0 + n0, m2 * M1 + m1] += IA[m2 * M1 + m1, k1 * K1 + k0] * W[k1 * K0 + k0, n2 * N1 * N0 + n1 * N0 + n0];
                        }
                    }
                }
            }
        }
        mvout(OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m2 * M1:(m2 + 1) * M1]);
    }
}

解决了上述的特殊情况之后，我们还需要解决一个问题：什么样的循环写法是最优的，例如对于case 3的代码，我们可以有如下两种写法：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


for (m2 = 0; m2 < M2; m2++)
{
    // IA buffer stores: M1*K
    mvin(IA [m2 * M1:(m2 + 1) * M1, 0:K]);
    for (n2 = 0; n2 < N2; n2++)
    {
        // W buffer stores: N1*N0*K
        mvin(W [0:K, n2 * N1 * N0:(n2 + 1) * N1 * N0]);
        compute_matmul(*W, *IA, *OA,
                       N2, N1, N0,
                       M2, M1,
                       K1, K0,
                       m2, n2);
        mvout(OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m2 * M1:(m2 + 1) * M1]);
    }
}

for (n2 = 0; n2 < N2; n2++)
{
    // W buffer stores: N1*N0*K
    mvin(W [0:K, n2 * N1 * N0:(n2 + 1) * N1 * N0]);
    for (m2 = 0; m2 < M2; m2++)
    {
        // IA buffer stores: M1*K
        mvin(IA [m2 * M1:(m2 + 1) * M1, 0:K]);
        compute_matmul(*W, *IA, *OA,
                       N2, N1, N0,
                       M2, M1,
                       K1, K0,
                       m2, n2);
        mvout(OA [n2 * N1 * N0:(n2 + 1) * N1 * N0, m2 * M1:(m2 + 1) * M1]);
    }
}

如何判断两种循环的优劣？我们需要统计据的搬移量：

/img-20230116180336230

针对这两种循环，我们可以设置不同的并行机制（即在不同位置使用spatial_for）。总而言之，设置loop的顺序本质上是一个优化问题，给定：维度数和硬件参数，最优化：能量/延迟。

Sparsity

ReLU函数的存在使得activation中存在大量的0，同理也可以通过剪枝使得weight中存在大量的0。

这为我们存储数值提供了机会：

bitmask：使用0/1数组表示是否为0（会造成独立于密度的固定开销，对高密度不友好）
run-length encoding：游程编码（同样有固定开销）
compressed sprase row：（可以进行快速的行选取）

这其中的col index直接记录其列号，而bound中的第i个数字记录了前i-1行包含的非0元素的数量。例如我们如果需要选取第三行的数字，由于Bounds[2]=3，Bounds[3]=5，我们可以得知在第三行的数字于value中的下标为[3,5)，据此可以直接取用数字。
compressed sprase column：（可以进行快速的列选取）

随后稀疏矩阵和密集矩阵的混合计算。按照操作数的稀疏性和对齐方式分为两种：

/img-20230117002256369

Indirection：使用稀疏矩阵中非0元素的下标去索引密集矩阵中的元素。

场景：算式y[i] = A[i, j] * x[j]，A的第二维是稀疏的，x是密集的。可以想到循环次数要小于len(x)，为len(A_index[1])。
Intersection：寻找权重和激活中均不为0的数值。

场景：算式y[i] = A[i, j] * x[j]，A的第二维是稀疏的，x是稀疏的（想象两个指针在非0下标数组上同时滑动的场景）。
Arbitration：先计算，之后在输出的矩阵中决定这些结果的位置。

场景：算式A[i, j] = B[i, k] * C[k, j]，B和C第二维均是稀疏的，我们先将两者“浓缩”的矩阵进行相乘，然后在输出中决定其位置。

Near-Data Processing

在之前的几节中，我们介绍的主要是Fully Connected和Conv两个深度学习的算子，这两者本质上都可以划归为GEMM（通用矩阵乘）。GEMM中的基本单元是MAC乘加单元，它包括2个operator，load 2个数据，然后将1个结果store进寄存器中。

深度学习中其实还有一些其他的算子，它们在operator、load、store上都和GEMM有很大的不同。举例说明：

Element-wise operation：逐元素运算，例如LSTM的Hadama积，ResNet中的残差相加。它包括1个operator，load 2个数据，然后将1个结果store进寄存器中。

这里需要特别强调的是残差相加，尽管没有引入额外的参数或者是计算复杂度，但仍然会导致减速，这是由于该操作增加了高速缓存的处理量，对之前数据的储存增加了内存负担。

Embedding Layer：嵌入运算。这一运算的本质是将id序列（例如词序号等）转换为one-hot向量之后，与Embedding大矩阵相乘，得到对应的词嵌入。不过实际操作中一般不使用这种类似于GEMM的运算，而是基于id序列对Embedding大矩阵进行一种类似于查找表的运算。因此它需要1个operator，load 2个数据，然后将1个结果store进寄存器中。

emb layer

实际上，上述的一些算子在深度学习之外的一些领域也有用，如GC（Java中的垃圾回收机制）、LINQ（语言集成查询）等。我们可以用上述算子构建Near-Data Processing的专用硬件（这里的Near-Data Processing指的是具有直接访问DRAM的功能，而不需要通过高速缓存的层次结构的硬件）。

In Memory Computing

在介绍IMC之前，我们需要首先回顾一下存储器的分类：

/img-20230117225552792

Random Access Memory(RAM): 随机存取存储器，其“随机存取”指的是当存储器中的消息被读取或写入时，所需要的时间与这段信息所在的位置无关。
- SRAM：静态随机存取存储器，其“静态”指的是这种存储器只要保持通电，里面存储的数据就可以恒常保持。多用于高速cache。
- DRAM：动态随机存取存储器，其“动态”指的是由于晶体管会有漏电电流的现象，导致电容上所存储的电荷数量并不足以正确的判别数据，进而导致数据毁损，所以DRAM需要经常性刷新。多用于廉价内存。
- Flash：闪存，与DRAM和SRAM不同的是掉电后数据不会被清除，但读写速度会慢。
Serial Access Memory(SAM): 顺序存取存储器需要按顺序读取存储数据。
Content Addressable Memory(CAM): 可以根据其内容而不是其名称或位置进行检索。它已被用于固定内容的高速存储和检索。

以下将重点计算在SRAM中进行的存内计算：

在Deep Learning等数据密集型的运算场景中，数据搬移的消耗将超越数据运算。一个直观的思路是在存储器中进行运算，而不是将数据从存储器中搬移出来后进行运算（例如，对于一个weight stationary的dataflow，权重只需要呆在存储器中，无需进行搬运）。这就是存内计算的基本思想。

/img-20230117210804520