深度神经网络模型包含的计算密集型算子一般为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的计算步骤如下:
首先通过
vmovups
指令将A矩阵的第一列移动到两个YMM寄存器中(这里假设为YMM0以及YMM1);随后,对于B矩阵第一行的第一个元素,使用
vbroadcastss
指令进行广播并存储到一个YMM寄存器内(这里假设为YMM2),然后使用fma指令vfmadd231ps
将YMM0和YMM1内的元素与YMM2内元素对应相乘,并将结果累加到C矩阵的两个YMM寄存器内,这里假设为YMM4以及YMM5;沿着B矩阵第一行进行循环,重复步骤2,B矩阵广播当前行内其它数据时重复使用YMM2寄存器,并将计算结果依次累加到YMM6~YMM15寄存器内;
A矩阵前进一列,B矩阵前进一行,并重复步骤1~3,最终完成整个C(