OneDNN GEMM(AVX FP32)算法浅析

深度神经网络模型包含的计算密集型算子一般为Dense、MatMul以及Conv。对于推理引擎，其执行Dense以及MatMul算子时，一般都会调用该引擎的GEMM或者GEMV（BatchSize等于1）实现；执行Conv算子时，虽然推理引擎一般存在更好的策略，如Winograd、NCHWxC等，但大多数推理引擎也会提供GEMM策略并在一定条件下选择并执行。当卷积核各个维度的尺寸都为1时，Conv算子的执行可直接调用GEMM，当卷积核尺寸不为1时，也可以通过Im2Col将输入进行变换后，再调用GEMM获取算子的输出结果。可以认为，GEMM的执行速度将直接影响以上计算密集型算子，特别是Dense以及MatMul算子的推理效率。

OneDNN作为Intel提供的算子加速库，在x86平台上具有非常优秀的性能。一款推理引擎，如果想要在x86平台有所作为，那么OneDNN一定是其参考并且benchmark比较的对象。对于GEMM，OneDNN会根据当前运行时设备最高支持的指令集，以及算子在MNK三个维度的大小，选择不同的实现，其所设计的情况非常庞大且复杂。本文将会对OneDNN GEMM算法在FP32 AVX2指令集下的一个比较通用的分支进行描述，其代码实现详见代码文件([1])。

在以下描述中，输入A矩阵的两个维度分别为M、K，B矩阵的两个维度分别为K、N，输出C矩阵的两个维度分别为M、N。

Micro Kernel

OneDNN在AVX2指令集下采用了A(16, 8) * B(8, 6) = C(16, 6)的Micro Kernel，其计算思路如下图所示（图有点抽象）：

整个Micro Kernel的计算步骤如下：

1. 首先通过vmovups指令将A矩阵的第一列移动到两个YMM寄存器中（这里假设为YMM0以及YMM1）；

2. 随后，对于B矩阵第一行的第一个元素，使用vbroadcastss指令进行广播并存储到一个YMM寄存器内（这里假设为YMM2），然后使用fma指令vfmadd231ps将YMM0和YMM1内的元素与YMM2内元素对应相乘，并将结果累加到C矩阵的两个YMM寄存器内，这里假设为YMM4以及YMM5；

3. 沿着B矩阵第一行进行循环，重复步骤2，B矩阵广播当前行内其它数据时重复使用YMM2寄存器，并将计算结果依次累加到YMM6~YMM15寄存器内；

4. A矩阵前进一列，B矩阵前进一行，并重复步骤1~3，最终完成整个C