通用矩阵乘（GEMM）优化与卷积计算

本文简要介绍通用矩阵乘（GEMM，General Matrix Multiplication）优化的基本概念和方法、神经网络量化中矩阵乘的优化方法。

通用矩阵乘优化

基本概念

通用矩阵乘（下文简称 GEMM）的一般形式是 = C=AB， 其中 A 和 B 涵盖了各自转置的含义。图一是矩阵乘计算中为计算一个输出点所要使用的输入数据。三个矩阵的形状也如图所示。

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

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

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

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

 

 How to optimize gemm 的优化效果

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];
    }
  }
}

 

 

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];
    }
  }
}

 

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];
    }
  }
}

 

 

 

 

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];
        }
      }
    }
  }
}

 

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

前两节探讨的是传统优化矩阵乘的两者思想，近年来兴起的神经网络中广泛应用的矩阵乘也可用类似方法优化（实际上深度学习软件可以直接使用已有的数学加速库，如 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 后，内存带宽需求直接下降为 14 。随之而来的内存容量需求变化出现了一些新的优化机会。而 QNNPACK 充分利用了这些优化方法，并结合神经网络领域的特点，大幅改进了计算性能。

 

 

内存布局与卷积性能

神经网络中卷积的内存布局主要有 NCHW 和 NHWC 两种。最后重点分析一下 im2col 1×1 卷积性能和内存布局的关系。

对于不需要额外调整输入的 1×1 卷积，将 NCHW 内存布局的卷积对应到矩阵乘 = 时， 是卷积核（filter）， 是输入（input）。各个矩阵的维度如图十二所示。

图十二：NCHW 内存布局卷积转换成的矩阵乘

对该矩阵施行划分后，将计算核的访存局部性表现标记在图十二中。其中 Inside 表示小块矩阵乘内部的局部性，Outside 表示在削减维度方向的局部性。

对输出而言，小块内向量化访存局部性较差，外部表现取决于全局计算方向——行优先则局部性较好，列优先则较差。由于卷积核可以事先重排内存，因此视其局部性都较好。输入则小块内外都较差，因为削减维度是列优先的，几乎每次加载输入都会发生高速缓存缺失。

图十三是与之相对的 NHWC 内存布局的示例。值得注意的是，NHWC 和 NCHW 中 、 矩阵所代表的张量发生了调换—— = × 。具体的拆分方式仍然一样。

图十三：NHWC 内存布局卷积转换成的矩阵乘

在 NHWC 中，输出的局部性表现和 NCHW 一样。同样的，卷积核也视作局部性表现较好。对于输入，小方块的内部局部性表现不是很好，因为几次向量加载的地址不连续；而外部局部性表现则较好，因为在削减维度滑动使用的内存是连续的——这一点在「处理内存布局」小节中已有阐述。

可以看到，对于 1×1如果采用 im2col 方法计算，且不对输入输出进行额外的内存重排，那么 NHWC 的访存特征是显著优于 NCHW 的。

总结

至此，本文介绍了 GEMM 优化的基本方法概念，在神经网络领域中 QNNPACK 基于量化对 GEMM 的优化，和 im2col 方法对卷积计算及其内存布局的意义。GEMM 优化实质上是个非常重要，且和特定领域绑定很强的话题，更进一步的内容需要进入到特定领域深入研究。如果对 GEMM 各种优化技巧所带来的性能收益感兴趣，可以参考 How to optimize gemm。如果对 GEMM 优化和体系结构结合的理论感兴趣，可以参考 Anatomy of High-Performance Matrix Multiplication 。

参考

• How to optimize gemm
• QNNPACK
• Anatomy of High-Performance Matrix Multiplication
• Convolution in Caffe