INTRODUCTION

在一个复杂的多级内存架构下，为了实现高性能矩阵乘，需要精心设计 $C = \widetilde{A}B + C$ 的算法，有两个关注的因素：

  • 浮点数运算速率和浮点数据从Cache到寄存器的传输速率。文章提到，由于浮点计算速率和L2 Cache到寄存器的传输速率比值较小，因此可以将矩阵存放在L2 Cache。
  • TLB可处理数据的大小限制矩阵 $\widetilde{A}$ 大小。

Alt text

文章以列主序存储矩阵，将矩阵乘分为以下八种：

Alt text

为了减少数据在不同层次内存之间的移动，优化访存，矩阵乘算法都将分块实现。GEMM可以被分解为GEPP、GEPM、GEMP，后者又可以继续被分解为GEBP、GEPB、GEPDOT。

HIGH-PERFORMANCE GEBP, GEPB, AND GEPDOT

Alt text

提出三个假设：

a. $m_c$, $k_c$ 足够小，A矩阵和B、C的一行（$B_j$, $C_j$）能够放在Cache中。

b. 如果 $A$, $B_j$, $C_j$ 都在Cache中，$C_j = AB_j + C_j$ 的计算速率能够达到峰值。

c. A能够在其需要被使用的时候能够一直存在Cache中。

计算和数据移动的比例为

在 $m_ck_c$ 个浮点数填充大部分Cache空间，且满足上述三条假设的情况下，该比例应当尽可能的大。

Alt text

文章从以下几个方面对GEBP继续进行优化：

  1. Cache层次
    在满足三条假设的前提下，矩阵A应当放在离寄存器更远的Cache中，因为可以存储更多的数据。文章对L2 Cache进行了分析：令 $R_{comp}$ 和 $R_{load}$ 分别是CPU进行浮点数计算的速率和将数据从L2 Cache传输到寄存器的速率，假设A在L2 Cache中， $B_j$, $C_j$ 在L1 Cache中，同时L1 Cache到寄存器的带宽足够大使 $B_j$, $C_j$ 传输的开销可以忽略。为了将矩阵A的传输开销语计算开销重叠，必须要满足 $2n_rR_{comp} \geq 1/R_{load}$ ，即

  2. TLB
    TLB是页表项的缓存，用于加速从虚拟地址到物理地址的翻译。如果一个虚拟地址在TLB中，可以直接翻译为物理地址，否则需要乡下继续查询页表，即发生了TLB miss。TLB miss和Cache miss的区别是后者可以通过prefetch等计数减少或消除CPU stall，而TLB miss必然引起CPU miss。因此，需要添加两条假设（限制条件）：

    d. $m_c$, $k_c$ 应当足够小，使得矩阵A、$B_j$、$C_j$ 能够在计算 $C_j=AB_j+C_j$ 的时候同时被TLB翻译。

    e. 当A的地址被TLB翻译后，它不在被需要。

  3. packing
    目前要实现得GEBP中，A矩阵在GEMM为一个内存不连续的block。为了减少A需要使用的TLB条目，将A包装成一个连续的数组 $\widetilde{A}$ 。此时 $m_c$, $k_c$ 参数选择受以下两个因素限制：

    • TLB大小。令TLB条目数为 $T$，$\widetilde{A}$, $B_j$, $C_j$ 需要使用的TLB条目数分别为 $T_{\widetilde{A}}$, $T_{B_j}$, $T_{C_j}$，需要满足条件 $T_{\widetilde{A}} + 2(T_{B_j} + T_{C_j}) \leq T$。其中乘数 $2$ 的原因是：当完成计算 $C_j=\widetilde{A}B_j+C_j$ 时, $\widetilde{A}$ 是最近最少使用的，此时加载 $B_{j+1}$ 和 $C_{j+1}$ 时，将会把 $\widetilde{A}$ 在TLB的条目替换掉，因此需要保证TLB中可以同时放置 $B_j$ 和 $B_{j+1}$ 、$C_j$ 和 $C_{j+1}$。同时由于在GEBP中矩阵B也需要被重复访问，因此将B打包成一个连续的数组 $\widetilde{B}$。
    • L2 Cache大小。由于TLB条目数的限制要大于L2 Cache大小的限制，通常不作考虑。

  4. 确保数据访问的连续性。具体做法将在后文阐述。

  5. GEPB和GEPDOT的实现与GEBP相似。

PRACTICAL ALGORITHMS

Alt text

TLB条目的使用情况如下：

  • $\widetilde{A}$ 使用 $T_{\widetilde{A}}$ 个条目
  • 将B打包成 $\widetilde{B}$ 后，$B_j$ 和 $B_{j+1}$ 分别使用1个条目
  • $C_{aux}$ 使用1个条目
  • $C_j$ 使用 $n_r$ 个条目

因此 $T_{\widetilde{A}}$ 被限制在 $T-(n_r + 3)$。$C_j$ 地址不连续并无影响，因为 $C_j$ 并不需要被重复利用。

文章提到：

  • B的打包是“memory-to-memory copy”，开销与 $k_c \times n$ 成正比，并且在 $2m \times n \times k_c$ 的计算量上平摊，因此对于每次拷贝，将执行 $O(m)$ 次计算。
  • A的打包时从内存到L2 Cache的拷贝，开销与 $m_c \times k_c$ 成正比，并且在 $2m_c \times k_c \times n$ 的计算量上平摊，因此对于每次拷贝，将执行 $O(n)$ 次计算。

综上两点，这种方式适合m和n都很大，且k不是很小的GEMM运算。

对于其他情况：

  • 基于GEBP实现GEPM
  • 基于GEPB实现GEPP
  • 基于GEBP实现GEMP
  • 基于GEPDOT实现GEPM和GEMP

原理都类似。
Alt text
Alt text
Alt text

其中，图8和图9展示了两种GEBP算法。两者的区别是，图8的算法先将矩阵B进行包装，再每次迭代中将矩阵C从内存中传入和传出，该算法能够通过计算 $C_{aux}=\widetilde{A}B_j$ 来隐藏对矩阵C操作的延迟；图9的算法再每次迭代中将矩阵B从内存传出，计算C的值，再将C累加后对矩阵C进行unpack操作，该算法能够隐藏包装矩阵B的延迟。由于对矩阵C的unpack操作相较于对矩阵B的unpack操作开销更大，因此图8的算法要优于图9。同理，图10的算法要优于图11。

图8算法和图10算法是对称的，当内存以列主序存储数据时，最好是按列访问矩阵中的每一块（图8中的A和图10中的B），因此图8算法比图10更优。

而对于GEPDOT，需要将一个矩阵C的子块放在L2 Cache中，并读取矩阵A的几列和矩阵B的几行，这需要L2 Cache到寄存器的带宽是GEBP和GEPB的两倍，因此其性能较差。

MORE DETAILS YET

上文分析得出结论，在列主序的情况下，GEBP_OPT1是最优的，下文将继续阐述实现GEBP_OPT1的更多细节。

Alt text

  1. 寄存器分块

    在计算 $C_{aux} = \widetilde{A}B_j$ 时，$C_{aux}$ 被分为 $m_r \times n_r$ 大小的子矩阵存在寄存器中。每次迭代有 $2m_rn_rk_c$ 次计算和 $m_rn_r$ 次寄存器到 $C_{aux}$ 的数据拷贝。为了隐藏数据拷贝的延迟，通常将 $k_c$ 设置大一点。

  2. $m_r \times n_r$

    • 确保一半寄存器用于保存矩阵C的大小为 $m_r \times n_r$ 子矩阵。另外一半寄存器用于 $\widetilde{A}$ 和 $\widetilde{B}$ 的预取。
    • 当 $m_r \approx n_r$ 时，平摊加载寄存器的代价是最优的。
    • 为了保证计算能够隐藏 $\widetilde{A}$ 从L2 Cache到寄存器的延时，需要满足条件 $n_r \geq R_{comp} / (2R_{load})$。

  3. $k_c$

    • $B_j$ 需要被重复利用多次，必须存在L1 Cache内，且小于L1 Cache的一半空间，剩下的缓存空间留给 $\widetilde{A}$ 和 $C_{aux}$。
    • $\widetilde{A}$ 占用L2 Cache的大部分。

  4. $m_c$

    减少 $m_c \times k_c$ 大小的 $\widetilde{A}$ 占用TLB和L2 Cache。

EXPERIMENTS

文章的实验环境为Intel Pentium 4 Prescott Processor (3.6 GHz, 64bit)，参数如下：
Alt text

$n_r \geq R_{comp} / (2R_{load}) = 0.97$，EM64T架构包含16个128-bit寄存器，因此 $m_r \times n_r$ 设为 $4 \times 4$。

Alt text

实验结果表示当 $m_c \times k_c = 696 \times 192$ 时，性能接进极大值（$696 \times 196$）。此时 $\widetilde{A}$ 大小约等于1K，占L2 Cache一半，与理论分析相符。

Alt text

当k相对于m和n较小时，实现的时GEPP。上图四条曲线含义如下：

  • 最顶上表示GEBP_OPT1算子的性能
  • 第二条表示使用GEPP算法实现的DGEMM性能
  • 底部两条分别表示 $\widetilde{A}$ 和 $\widetilde{B}$ 的pack操作时间开销占整体的百分比。该操作导致kernel到dgemm的性能下降。

Alt text

当m很小时（GEMP），将B转换为 $\widetilde{B}$ 的开销无法被计算平摊，从而导致性能较差。一种解决方法是跳过B的转换，另一种方法是实现图9的算法；同理，当m=k=2000，n较小时（GEMP），将A转换为 $\widetilde{A}$ 的开销无法被浮点计算给平摊，此时应当实现图11的算法。