高性能计算的内存模型研究

本文首先将建立一个简单的高性能计算内存模型，随后以矩阵乘为例，探讨分块矩阵乘法是如何在内存访问上对整个的矩阵乘计算进行加速的。最后，本文会将该内存模型与Roofline Model结合并进行一个简单的讨论。

1. 基础：一个简单的高性能计算内存模型

本文所述的内存模型具有以下特点：

• 在内存层级方面，只具有两级：fast mem和slow mem；

• 对于slow mem，其读写操作的时间相同，对于fast mem，其读写操作的时间可以忽略；

• 在算法运行开始，所有的数据都在slow mem中；

在以上内存模型的基础上，给定以下参数：

• m：需要在slow mem和fast mem之间搬运的数据量；

• tm：从slow mem搬运一次数据（读或写）的时间；

• f：需要进行的算数运算数量（例如在推理引擎中即为算子的计算量）；

• tf：进行一次算数运算的时间（可认为远小于tm）；

基于以上模型，我们可以认为对于任意算法，其执行时间应由两部分组成，即内存搬运时间以及实际的计算时间：

在以上公式中，我们定义 CI=f/m ，其描述了每从slow mem中搬运一次数据能够进行的计算的数量，并将其称为计算密度。则以上公式可最终修改为：

2. 矩阵乘算法分析

在本章中，我们会以上述内存模型作为基础，对矩阵乘 C=C+A∗B 进行分析。同时为了简化计算，设置输入输出矩阵均为长度为n的方阵。根据计算量公式，以上矩阵乘的计算量f可以很容易得到，为2n^3。

最朴素的矩阵乘

在我们实现矩阵乘时，最为朴素的方式是三层循环，其伪代码以及访存分析如下：

for i = 1 to n
  {read row i of A into fast memory}        total n^2 to read
  for j = 1 to n
    {read C[i,j] into fast memory}          total n^2 to read
    {read column j of B into fast memory}   total n^3 to read
    for k = 1 to n
      C[i,j] = C[i,j] + A[i,k] * B[k,j]
    {write C[i,j] back to slow memory}      total n^2 to write

由以上分析，可以得到上述算法的计算密度：

可以看到，在最朴素的矩阵乘中，其计算密度仅为2。

分块矩阵乘

分块矩阵乘是加速矩阵乘算法的常用手段。这里为了简化分析，我们只对矩阵乘进行一次分块，并设置分块数为N，则其分块大小b为 n/N ，如下图所示：

在以上背景下，分块矩阵乘的伪代码以及访存分析如下：

for i = 1 to N
  for j = 1 to N
    {read block C(i,j) into fast memory}        total N^2 * b^2 = n ^ 2 to read
    for k = 1 to N
      {read block A(i,k) into fast memory}      total N^3 * b^2 = N * n^2 to read
      {read block B(k,j) into fast memory}      total N^3 * b^2 = N * n^2 to read
      C(i,j) = C(i,j) + A(i,k) * B(k,j)  {do a matrix multiply on blocks}
    {write block C(i,j) back to slow memory}    total N^2 * b^2 = n ^ 2 to write

在以上分块策略下，我们可以获得该算法的计算密度：

可以看到，计算密度由之前的2提升到了分块矩阵的大小b。假设fast mem的大小为M，则为了获得最佳性能，我们需要保证在分块矩阵进行计算时，其数据都能够驻留在fast mem中，在最理想情况（即不考虑临时变量以及代码段等占用fast mem）下，分块矩阵的尺寸b最大可为sqrt(M/3)，即计算密度 CI=O(sqrt(M))。

以上结论也符合我们的基本认知，即cache size越大，我们对矩阵分块时的分块大小就可以设置的越大，其相应的计算密度也会越大，整个矩阵乘所需的计算时间就会减小了。

3. 结合Roofline Model进行讨论

如果读者接触过Roofline Model，可能早已发现我们定义的计算密度CI就是Roofline Model的横轴。在上面的矩阵乘示例中，我们通过使用分块矩阵乘，提高了计算密度，对应于Roofline Model，算法的横坐标向右移动，在相同带宽下也就提高了算法的计算速度。

还是以CPU上的矩阵乘法为例，我们在实际的GEMM编写时，会使用多种手段以减少CPU从主存中搬运数据的次数以及数量，包括但不限于数据对齐、对权重矩阵进行重排、调整循环顺序、合理设计分块大小以使内部计算数据均驻留在L2 cache中、适当插入软件Prefetch指令、使用向量指令读写等。可见要从内存的角度提高算法的计算密度并使其接近算法的理论极限，还是相当复杂的。

最后，本文的内容基本上来自课程cs267，如果您对本文讨论的内容感兴趣，可以搜索该课程进行更加深入的学习，也欢迎大家在评论区进行交流。

---