【强基固本】通用矩阵乘（GEMM）优化算法

“强基固本，行稳致远”，科学研究离不开理论基础，人工智能学科更是需要数学、物理、神经科学等基础学科提供有力支撑，为了紧扣时代脉搏，我们推出“强基固本”专栏，讲解AI领域的基础知识，为你的科研学习提供助力，夯实理论基础，提升原始创新能力，敬请关注。

来源：https://jackwish.net/2019/gemm-optimization.html

本文采用知识共享 署名-非商业性使用-禁止演绎 4.0 国际许可授权

01

引言

气象预报、石油勘探、核子物理等现代科学技术大多依赖计算机的计算模拟，模拟计算的核心是表示状态转移的矩阵计算。另一方面，计算机图形处理以及近年来兴起的深度学习也和矩阵乘高度相关。而矩阵乘对计算资源消耗较大，除了计算机体系结构的不断更新外，软件优化方面也有大量的研究工作。

本文简要介绍通用矩阵乘（GEMM，General Matrix Multiplication）优化的基本概念和方法、神经网络量化中矩阵乘的优化方法。旨在帮助大家在概念中建立一些直觉，无甚高论。

02

通用矩阵乘概念

矩阵乘通常定义为

其中   三者的形状分别为   。图一是矩阵乘的可视化展示，和计算时为得到一个输出点所要使用的输入数据。

图一：矩阵乘一个输出元素的计算

与之相对应的伪代码表示为：

for (int m = 0; m < M; m++) {
  for (int n = 0; n < N; n++) {
    C[m][n] = 0;
    for (int k = 0; k < K; k++) {
      C[m][n] += A[m][k] * B[k][n];
    }
  }
}

对这样的矩阵乘的算法优化可分为两类：

基于算法分析的方法：根据矩阵乘计算特性，从数学角度优化，典型的算法包括 Strassen 算法和 Coppersmith–Winograd 算法。

基于软件优化的方法：根据计算机存储系统的层次结构特性，选择性地调整计算顺序，主要有循环拆分向量化、内存重排等。

下面将简要介绍几种典型的方法。

03

基于算法分析的方法

算法分析可知，朴素的矩阵乘算法的时间复杂度为  。在很长的时间内，人们认为矩阵乘在算法层面是无法优化的，而自 Strassen 算法伊始，复杂度边界便被不断降低，如图一。目前最快的方法是 Coppersmith–Winograd 算法。

图二：矩阵乘算法复杂度边界的演变

这些算法一般要求三个矩阵符合约束  。

Strassen 算法

Volker Strassen 在 1969 年提出了复杂度为  的矩阵乘算法。这是历史上第一次将矩阵乘的计算复杂度降低到  以下。

基于分治（Divide and Conquer）的思想，Starssen 算法将矩阵  分别拆分为更小的矩阵

其中  。

根据矩阵基本的运算法则，拆分后朴素算法的计算如下所示，共需要八次小矩阵乘法和四次小矩阵加法计算。

显然，分治本身并不能改进矩阵乘的计算效率。在很长的时间内，人们认为矩阵乘没有什么优化算法，直到 Strassen 引入了七个如下所示的用于辅助计算的中间矩阵

在得到这些中间矩阵的记过后，再将其组合得到最后的矩阵：

通过七次乘法和十八次加法，Strassen 算法将矩阵乘的算法复杂度降低到了  （递归地运行该算法）。如图二所示，该算法突破性地将矩阵乘计算复杂度从  拉了下来，后续的算法都是对该算法的某种程度上的改进。而 Strassen 算法也成为了个大算法教材讲解复杂度分析的重要示例。

完全应用 Strassen 算法的一个局限是其要求矩阵乘的规模为  ，这在现实情况中不容易满足。一种解决方法是将规模分解为  其中  无法被 2 整除，那么可以应用 Strassen 算法不断递归拆分计算直到小矩阵规模为  。此时可以用朴素算法计算小矩阵；或者将  补零为  再继续应用 Strassen 算法（亦可直接对大矩阵补零）。最终的性能取决于实现方法和运行的硬件平台。

Coppersmith–Winograd 算法

Strassen 算法提出之后，学者们不断尝试继续降低复杂度，因为  的代价还是太高了。另一方面，Strassen 算法虽然学术意义很大，但实际应用却有限。矩阵乘算法的复杂度边界终于在 Don Coppersmith 和 Shmuel Winograd 的合作下在 1990 年突破性地降低到了  。

Coppersmith–Winograd 算法的思想和 Strassen 算法类似。其证明过程比较复杂，使用的定理太多，这里就不再介绍（实际上是没看懂…），有兴趣的可以参考下面几篇文章：

Matrix multiplication via arithmetic progressions (原始论文)

The Coppersmith-Winograd Matrix Multiplication Algorithm

On the Coppersmith–Winograd method

在 Coppersmith–Winograd 算法之后，学者们依然在不断尝试降低矩阵乘算法的复杂度，但 Strassen 算法和 Coppersmith–Winograd 算法目前依然是矩阵乘算法优化的两个里程碑。

04

基于软件优化的方法

除了从算法分析的角度优化通用矩阵乘，在实际的计算机系统中应用很多的还有软件优化的方法。软件优化方法基于对计算机体系机构和软件系统的特征分析，结合具体计算的特性，设计出针对性的优化方法。对矩阵乘而言，比较重要的软件优化方向包括：改进访存局部性、利用向量指令等。

我们回顾一下矩阵乘的伪代码，其计算操作总数为  （其中  分别指代三层循环执行的次数，2 指代循环最内层的一次乘法和加法） ，内存访问操作总数为  （其中  是累加求和的循环，  指代对  三者的内存访问，  需要先读取内存、累加完毕再存储，且忽略对  初始化时的操作）。GEMM 基于软件优化的性能改进以此为基点。

How to optimize gemm 介绍了如何采用各种优化方法，将最基础的计算改进了约七倍（如图三）。其基本方法是将输出划分为若干个 子块，以提高对输入数据的重用。同时大量使用寄存器，减少访存；向量化访存和计算；消除指针计算；重新组织内存以地址连续等。详细的可以参考原文。

图三：How to optimize gemm 的优化效果

计算拆分展示

本节主要以图形化的方式介绍计算拆分。

图四 将输出的计算拆分为  的小块，即将  维度拆分为两部分。计算该块输出时，需要使用  矩阵的 1 行，和  矩阵的 4 列。

图四：矩阵乘计算  输出

下面是该计算的伪代码表示，这里已经将  中  维度的内部拆分进行了展开。这里的计算操作数仍然是  ，这一点在本文中不会有变化。这里的内存访问操作数尚未出现变化，仍然是  ，但接下来会逐步改进。

for (int m = 0; m < M; m++) {
  for (int n = 0; n < N; n += 4) {
    C[m][n + 0] = 0;
    C[m][n + 1] = 0;
    C[m][n + 2] = 0;
    C[m][n + 3] = 0;
    for (int k = 0; k < K; k++) {
      C[m][n + 0] += A[m][k] * B[k][n + 0];
      C[m][n + 1] += A[m][k] * B[k][n + 1];
      C[m][n + 2] += A[m][k] * B[k][n + 2];
      C[m][n + 3] += A[m][k] * B[k][n + 3];
    }
  }
}

简单的观察即可发现，上述伪代码的最内侧计算使用的矩阵  的元素是一致的。因此可以将  读取到寄存器中，从而实现 4 次数据复用（这里不再给出示例）。一般将最内侧循环称作计算核（micro kernel）。进行这样的优化后，内存访问操作数量变为   ，其中  是对  优化的效果。

类似地，我们可以继续拆分输出的  维度，从而在内侧循环中计算  输出，如图五。

图五：矩阵乘计算  输出

同样地，将计算核心展开，可以得到下面的伪代码。这里我们将    中展示过的 N 维度的计算简化表示。这种拆分可看成是   ，这样    和    的访存均可复用四次。由于乘数效应，   的拆分可以将对输入数据的访存缩减到    。这相对于最开始的    已经得到了 1.6X 的改进，这些改进都是通过展开循环后利用寄存器存储数据减少访存得到的。

for (int m = 0; m < M; m += 4) {
  for (int n = 0; n < N; n += 4) {
    C[m + 0][n + 0..3] = 0;
    C[m + 1][n + 0..3] = 0;
    C[m + 2][n + 0..3] = 0;
    C[m + 3][n + 0..3] = 0;
    for (int k = 0; k < K; k++) {
      C[m + 0][n + 0..3] += A[m + 0][k] * B[k][n + 0..3];
      C[m + 1][n + 0..3] += A[m + 1][k] * B[k][n + 0..3];
      C[m + 2][n + 0..3] += A[m + 2][k] * B[k][n + 0..3];
      C[m + 3][n + 0..3] += A[m + 3][k] * B[k][n + 0..3];
    }
  }
}

到目前为止，我们都是在输出的两个维度上展开，而整个计算还包含一个削减（Reduction）维度   。图六展示了在计算    输出时，将维度    拆分，从而每次最内侧循环计算出输出矩阵    的    部分和。

图六：矩阵乘计算  输出对 K 维度的拆分

下面展示的是这部分计算的展开伪代码，其中维度    和    已经被简写。在这里，最内侧循环发生的计算次数已经从最朴素版本的    发展到了  。

for (int m = 0; m < M; m += 4) {
  for (int n = 0; n < N; n += 4) {
    C[m + 0..3][n + 0..3] = 0;
    C[m + 0..3][n + 0..3] = 0;
    C[m + 0..3][n + 0..3] = 0;
    C[m + 0..3][n + 0..3] = 0;
    for (int k = 0; k < K; k += 4) {
      C[m + 0..3][n + 0..3] += A[m + 0..3][k + 0] * B[k + 0][n + 0..3];
      C[m + 0..3][n + 0..3] += A[m + 0..3][k + 1] * B[k + 1][n + 0..3];
      C[m + 0..3][n + 0..3] += A[m + 0..3][k + 2] * B[k + 2][n + 0..3];
      C[m + 0..3][n + 0..3] += A[m + 0..3][k + 3] * B[k + 3][n + 0..3];
    }
  }
}

在对    和    展开时，我们可以分别复用    和    的数据；在对    展开时，其局部使用的 C 的内存是一致的，那么在    迭代时可以将部分和累加在寄存器中——最内层循环整个迭代一次写到    的内存中。那么内存访问次数为如下，可以看到最后得到的加速比为惊人的 8X。其中，   和    循环共执行  次，每次存储    个    的输出元素，因此    的访存共    次，相比于   可忽略。

实际上，现代处理器的 SIMD 能力远超本文示例中展开成 4 的情况，因此可以使用更为激进的分块计算，从而得到更为优化的性能。但究其原理仍然和本例相似。

到目前为止，我们已经对计算进行了三次维度拆分、展开并寄存器化（代码未展示）和访存复用。这对最基础版本计算似乎已经足够了，因为数百条稳定的顺序计算指令对处理器流水线已经很友好，且编译器可以帮助我们做好软件流水这样的指令调度。

然而一条计算指令只能完成一次乘加操作（MLA）效率还是比较低。实际上即使是最低端的移动手机处理器都会带有 SIMD 支持，访存和计算都可以向量化。因此我们可以再进一步，利用向量操作提高计算的性能。

在介绍向量化计算的细节时，伪代码是很难理解的，下面依据图七介绍量化计算的具体过程。图七左侧部分的三幅小图分别展示了两个    矩阵相乘向量化的要素：首先是计算一个输出元素使用到的输入元素；然后是对各个矩阵内存的编码，均以行优先的形式编号；最后是向量化的具体计算方法。

图七：两个    矩阵相乘的向量化

这里的向量编号方式假定输入输出的内存布局都是行优先，那么两个输入各自的 16 个元素通过 4 次向量访存即可加载到寄存器中。由矩阵乘的规则可知，输入    中的行可一次性用作输出的计算，而输入    的行则要拆分使用。这也是向量计算最容易出错的地方。

图七右侧列出了三份伪代码。第一份 C0 in detail 是计算 C0 中四个输出元素的朴素方法的展开，连续四次计算得到一个 C0 元素的结果。将计算过程稍作重排，即可得到 C0 scheduled 展示的计算，这里连续的四次计算分别处理了 C0 中四个元素的  结果。在连续的四次计算中，重排前只有对 A 的访存是连续的，重排后 B 和 C 的访存都是连续的。那么向量化这些访存和计算，即可得到第三列伪代码中红色 C0 部分。而第三列是通过对 C1、C2、C3 进行类似的处理得到的。

施行向量化操作后，原本需要 64 条计算指令的计算过程所需指令减少到 16 条，访存也有类似效果。而向量化对处理器资源的高效使用，又带来了进一步优化空间，例如可以一次计算   个局部输出。

处理内存布局

上一小节列出的是在输入输出原有内存布局上所做的优化。在最后向量化时，每次内存访问都是四个元素。当这些元素为单精度浮点数时，内存大小为 16 字节，这远小于现代处理器高速缓存行大小（Cache line size）——后者一般为 64 字节。在这种情况下，内存布局对计算性能的影响开始显现。

图八展示的是不同内存组织方式对的影响。图中两者都是行优先的内存排列，区别在于左侧小方块内部是不连续的，右侧小方块内部是连续的。图中用几个数字标记了各个元素在整个内存中的编号。

图八：不同的局部内存组织

想象一下在这两者不同内存组织方式的输入输出中访存，每次向量化内存加载仍是 4 个元素。对于一个局部计算使用到的小方块，左侧四次访存的内存都是不连续的，而右侧则是连续的。当数据规模稍大（一般情况肯定足够大了），左侧的连续四次向量化内存加载都会发生高速缓存缺失（cache miss），而右侧只会有一次缺失。

在常规的数据规模中，由于左侧会发生太多的高速缓存缺失，又由于矩阵乘这样的计算对数据的访问具有很高的重复性，将它重排成右侧的内存布局减少高速缓存缺失，可显著地改进性能。另一方面，矩阵乘中两个输入矩阵往往有一个是固定的参数，在多次计算中保持不变。那么可以在计算开始前将其组织成特定的形状，这种优化甚至可以将性能提高 2x。

到这里为止，对    计算已经有了足够的优化，可以开始考虑视野更广一些的全局优化。图九是一个关于全局优化的小示例。

图九：矩阵乘的全局优化一瞥

图中字母标记的是全局性的工作顺序，即输出数据中外层循环迭代方式。左侧小图是常规的行优先遍历方式，中间小图是列优先的遍历方式。这两者的区别是    和    两个维度的循环哪个在最外层。

上文已经对    和    两个维度分别进行了一次拆分，这里可以继续这种拆分。右侧的图例中是将    和   两个维度分别拆分为    三部分，将外层拆分都交换到外层循环。下面是相应的伪代码。

for (int mo = 0; mo < M; mo += 8) {
  for (int no = 0; no < N; no += 8) {
    for (int mi = 0; mi < 2;mi ++) {
      for (int ni = 0; ni < 2; ni++) {
        int m = mo + mi * 4;
        int n = no + ni * 4;
        C[m + 0..3][n + 0..3] = 0;
        C[m + 0..3][n + 0..3] = 0;
        C[m + 0..3][n + 0..3] = 0;
        C[m + 0..3][n + 0..3] = 0;
        for (int k = 0; k < K; k += 4) {
          C[m + 0..3][n + 0..3] += A[m + 0..3][k + 0] * B[k + 0][n + 0..3];
          C[m + 0..3][n + 0..3] += A[m + 0..3][k + 1] * B[k + 1][n + 0..3];
          C[m + 0..3][n + 0..3] += A[m + 0..3][k + 2] * B[k + 2][n + 0..3];
          C[m + 0..3][n + 0..3] += A[m + 0..3][k + 3] * B[k + 3][n + 0..3];
        }
      }
    }
  }
}

经过这样的调度，从整体计算来看，可看作是将    计算拓展成了    ，其实是同一种思路。

05

神经网络量化中的矩阵乘优化

前两节探讨的是传统优化矩阵乘的两者思想，近年来兴起的神经网络中广泛应用的矩阵乘也可用类似方法优化（实际上深度学习软件可以直接使用已有的数学加速库，如 BLAS）。随着技术的演进，神经网络技术出现了一个重要的方向——神经网络量化。量化技术的出现使得我们可以在深度学习领域使用一些特别的方法来优化矩阵乘，例如 QNNPACK (Quantized Neural Network PACKage) 和 Gemmlowp (GEMM Low Precision) 加速库。本节以 QNNPACK 为例介绍相关的技术。

QNNPACK 是 Facebook 开源的专门用于量化神经网络的计算加速库。QNNPACK 和 NNPACK (Neural Network PACKage) 的作者都是 Marat Dukhan 。到目前为止（2019 年中），QNNPACK 是已公开的，用于移动端（手机）的，性能最优的量化神经网络加速库。

QNNPACK 开源时附带了一份技术报告性质的博客。本节将结合上节的内容简要地从博客原作中抽取一些关于 GEMM 的内容。

量化神经网络

神经网络计算一般都是以单精度浮点（Floating-point 32, FP32）为基础。而网络算法的发展使得神经网络对计算和内存的要求越来越大，以至于移动设备根本无法承受。为了提升计算速度，量化（Quantization）被引入到神经网络中，主流的方法是将神经网络算法中的权重参数和计算都从 FP32 转换为 INT8 。

两种数值表示方法的方程如上。如果对量化技术的基本原理感兴趣，可以参考神经网络量化简介。

应用量化技术后，计算方面显现了若干个新的问题。首先是 NNPACK 这样用于 FP32 的计算加速库无法用于 INT8 ，这导致我们需要新的加速计算方法。再者是输入输出都转化成 INT8 后，内存带宽需求直接下降为    。随之而来的内存容量需求变化出现了一些新的优化机会。而 QNNPACK 充分利用了这些优化方法，并结合神经网络领域的特点，大幅改进了计算性能。

计算划分与削减维度

和上节所述类似，QNNPACK 的计算也是基于对输出的划分，拆分成如图十的    小块。这里需要注意的一点是，原文图例中对   的标记有笔误，   的列高度应该是    而非   。（我们向 QNNPACK 报告了这一问题，目前尚未得到修正。）

图十：QNNPACK 矩阵乘划分示例

QNNPACK 实现了    的和    两种计算核（micro kernel），分别用于支持 armv7 和 arm64 指令集的处理器。这两种计算核在原理上区别不大，后者主要利用了更多的寄存器和双发射（Dual Issue）以提高计算的并行度。

拆分后    计算块使用的内存为    。由于常规的神经网络计算    ， 那么这里的内存消耗一般不超过 16KB，可以容纳在一级高速缓存（L1 Cache）中。QNNPACK 的这一发现是其矩阵乘优化的基础。

如图十一所示，在计算    小块时，传统的方法（即上一节的方法）是在    维度上拆分，在一次计算核的处理中，仅计算    维的局部。那么在每次计算核的处理中，都会发生对输出的加载和存储——要将本次计算产生的部分和累加到输出中。

图十一：QNNPACK 和传统方法计算削减维度的对比

而 QNNPACK 的做法是将整个   维全部在计算核中处理完，这样就几乎完全消除了输出部分和的访存。两种的差异的细节可以参考图十一两侧的伪代码。这里所说的「将整个   维全部」并不是指   维不拆分——在实际计算中 K 维还是会以 8 为基础拆分——而是指拆分后不和其他维度交换（interchange）。

内存组织的特点

上节中曾提到，对内存的重新组织（Repacking）可以改进高速缓存命中率，从而提高性能。但是这种重新组织也是有开销的。

计算核中最小的计算单元处理的是两个    矩阵相乘。传统的方法由于   可能很大，需要对输入内存进行重新组织，防止相邻的访存引起高速缓存冲突，如图十二。

图十二：QNNPACK 和传统矩阵乘对局部计算的处理

而在量化神经网络中，由于    比较小，计算核处理中使用到的内存完全可以容纳在一级高速缓存中，即使不重新组织内存，高速缓存的重用率也足够高。

参考图八左侧部分，QNNPACK 计算核一次会使用 8 行输入（假定图中绘制以 8 为基础分块）。尽管对第一个    矩阵块的向量化加载可能全部是高速缓存缺失（Cache miss），第二个    则全部命中——因为它们已经作为同在一个高速缓存行的内容随第一个矩阵块加载到了高速缓存中。其他矩阵块也是类似情况。

采用了这些基于神经网络领域先验知识的优化方法后，QNNPACK 击败了所有神经网络量化领域的用于移动端加速库。不过，QNNPACK 的拆分着眼于削减维度，没有在输出维度上做全局调度。我们在 QNNPACK 基础上实现了    和    维度的外层循环拆分调度，简单的实验获得了相对于 QNNPACK 1.1x 的性能表现。这项工作目前尚未进入公开流程。

06

总结

至此，本文介绍了 GEMM 优化的基本方法概念，以及神经网络量化领域中 QNNPACK 的优化。GEMM 优化实质上是个非常重要，且和特定领域绑定很强的话题，更进一步的内容需要进入到特定领域深入研究。对相关主题感兴趣的话可以进一步参考下列资料：

Matrix Multiplication (Wikipedia)

Strassen Algorithm (Wikipedia)

The Coppersmith-Winograd Matrix Multiplication Algorithm

How to optimize gemm

QNNPACK

Anatomy of High-Performance Matrix Multiplication

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“强基固本”历史文章

更多强基固本专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！