GEMM-intrinsic优化
本文最后更新于:1 年前
写在前面:
对于一个想自行实现的GEMM算子,需要结合芯片的体系结构考虑,主要分为两步:①外层循环;②微内核
- 外层循环主要做的是:分块+移动数据,这里分块要仔细考虑片上缓存和片外缓存的大小,如CPU就考虑L3 cache,L2 cache,L1 cache的大小来进行大小合适的分块,并且分块怎么分也有讲究,以提高数据重用率;然后就是移动数据到缓存上;这里的计算和访存是可以overlap的,下文统一叫存算重叠(当然数据预取,双缓冲其实意思差不大多)
- 微内核则主要考虑存算重叠(取指,计算,写回),指令重排,指令打包,以发挥硬件的最大性能
而我们这里的微内核没有精细到ILP(Instruction Level Parallel),而是采用向量intrinsic编程,来实现汇编上一层的优化,因此最终并不能完全将硬件的性能发挥到巅峰
但是intrinsic function是在如AVX2指令集上封装了一层的类C的可无缝兼容于C/C++函数的function,编写起来更加顺手一些,因此先在这之上进行算子编写,之后再转用内联汇编~
内积和外积
在矩阵乘中的内积和外积,与我们数学物理中所认识的内积外积不同,如下图所示:

在矩阵中所说的内积外积取决于我们对矩阵乘的操作。在,这两个矩阵上,我们做矩阵乘可以得到。
拿出A的第一行与B的第一列相乘,可以得到C阵的第一行第一列的元素,如C[0][0] = A[0][:] * B[:][0]
。这种矩阵乘的方式叫内积,按如此方式(如图式)计算矩阵有下列伪码:
1 |
|
拿出A的第一列与B的第一行相乘,此时可以得到一个跟C阵规模一样的子阵C1:C1=A[M][0]*B[0][N]
。这种矩阵乘依次滑动k即可,得到K个子阵,累加起来便是最终的C阵。这种矩阵乘的方式叫做外积,伪码如下:
1 |
|
外积的方式显然更适合把两个矩阵按K这个维度对应拆分给多个核心计算,然后最后再AllReduce;
而两者从单一核心上计算来看,其实都没有很好的利用内存局部性这一性质,不过对于外积,它算完了就切到下一个维度,并没有反复的cache miss,而对于内积,对于B矩阵的来说,cache miss在反复的发生着,因为计算一行C阵的数据,就要过一遍B阵的数据,而每次B阵的数据还得重新拿,且矩阵按行存储,并不能一下就直接拿到一整行的数据,因此对于这种内积方式,需要进行循环重排以优化内存局部性,提升性能(不然大量时间耗费在访存上了),重排也很简单,就是由原先的M->N->K的遍历方式换成M->K->N,如下图式:

CPU算力
单颗CPU{单精度/双精度}flops = 核心数 * 单个核心flops
单个核心flops = 主频 * 单周期{双精度/单精度}计算性能
单周期单精度的计算性能 =( FMA计算单元数目 * 2(一次两条指令)* 寄存器位宽 )/ 32(单精度)
主频可以通过lscpu
里面查看
FMA数目应该要翻阅处理器的手册或去官网查详细的参数
寄存器位宽需要配合FMA支持情况查看,有YMM
256位的寄存器,也有ZMM
512位的寄存器
FLAME的矩阵优化
在学习的时候,跟着参考资料三的项目进行了一些基础的矩阵优化[3],主要包括的优化思想是循环展开(loop unroll),内联函数(inline function),矩阵分块(blocking),矩阵打包(packing),intrinsic编程。因为原项目的优化有点点啰嗦,我给它重写并合并了一些优化手段
主要思想:
- 通过朴素的矩阵乘,将内积的函数单独
add_dot
提取出来,然后由1x1,即一行乘一列的1次内积变换为1x4,即一次做4次内积,取A阵的一行,与B阵的四列分别进行相乘,则一次可以完成四次内积; - 在此之上对
add_dot
进行内联,以减少函数调用次数,和减少循环次数以及利用空间局部性; - 利用寄存器存储频繁使用的值
- 循环展开
- 在1x4思想上变为4x4,即一次完成十六次内积
- 采用intrinsic编程(SSE指令集,由于只有128位长的向量寄存器,其实对于64位的双精度浮点数有点不是很够用…)
- 分块,打包
相关代码细节:OperatorDev
矩阵分块
针对行主序的矩阵进行计算,其中,我们以BLAS(基本的线性代数函数接口规范)有:
接口声明如下:
1 |
|
在进行矩阵计算的时候,我们一个很大的矩阵无法放进寄存器里,也可能无法放进L1,L2cache中,因此需要对矩阵分块放置在这些
micro-kernel
参考文章
本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!