GEMM-intrinsic优化

本文最后更新于:1 年前

写在前面:

对于一个想自行实现的GEMM算子,需要结合芯片的体系结构考虑,主要分为两步:①外层循环;②微内核

  • 外层循环主要做的是:分块+移动数据,这里分块要仔细考虑片上缓存和片外缓存的大小,如CPU就考虑L3 cache,L2 cache,L1 cache的大小来进行大小合适的分块,并且分块怎么分也有讲究,以提高数据重用率;然后就是移动数据到缓存上;这里的计算和访存是可以overlap的,下文统一叫存算重叠(当然数据预取,双缓冲其实意思差不大多)
  • 微内核则主要考虑存算重叠(取指,计算,写回),指令重排,指令打包,以发挥硬件的最大性能

而我们这里的微内核没有精细到ILP(Instruction Level Parallel),而是采用向量intrinsic编程,来实现汇编上一层的优化,因此最终并不能完全将硬件的性能发挥到巅峰

但是intrinsic function是在如AVX2指令集上封装了一层的类C的可无缝兼容于C/C++函数的function,编写起来更加顺手一些,因此先在这之上进行算子编写,之后再转用内联汇编~

内积和外积

在矩阵乘中的内积和外积,与我们数学物理中所认识的内积外积不同,如下图所示:

在矩阵中所说的内积外积取决于我们对矩阵乘的操作。在AMKA_{M*K},BKNB_{K*N}这两个矩阵上,我们做矩阵乘可以得到CMNC_{M*N}

拿出A的第一行与B的第一列相乘,可以得到C阵的第一行第一列的元素,如C[0][0] = A[0][:] * B[:][0]。这种矩阵乘的方式叫内积,按如此方式(如图式)计算矩阵有下列伪码:

1
2
3
4
5
6
7
8
// 假设矩阵的存储是行优先/行主序的 
for(int i=0;i<M;i++){ // a_row -> M
for(int j=0;j<N;j++){ // b_col -> N
for(int k=0;k<K;k++){ // a_col == b_row -> K
C[i*N + j] += A[i*K + k]*B[k*N + j]; //这里采用一维数组存储矩阵的值,可以结合行主序列主序一起理解
}
}
}

拿出A的第一列与B的第一行相乘,此时可以得到一个跟C阵规模一样的子阵C1:C1=A[M][0]*B[0][N]。这种矩阵乘依次滑动k即可,得到K个子阵,累加起来便是最终的C阵。这种矩阵乘的方式叫做外积,伪码如下:

1
2
3
4
5
6
7
8
// 假设矩阵的存储是按行优先存储的
for(int k=0;k<K;k++){
for(int i=0;i<M;i++){
for(int j=0;j<N;j++){
C[i*N + j] += A[i*K + k] * B[k*N + j];
}
}
}

外积的方式显然更适合把两个矩阵按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支持情况查看,有YMM256位的寄存器,也有ZMM512位的寄存器

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

矩阵分块

针对行主序的矩阵进行计算,其中AM×KA_{M×K}BK×NB_{K×N},CM×NC_{M×N}我们以BLAS(基本的线性代数函数接口规范)有:

C=αAB+βCC = \alpha AB + \beta C

接口声明如下:

1

在进行矩阵计算的时候,我们一个很大的矩阵无法放进寄存器里,也可能无法放进L1,L2cache中,因此需要对矩阵分块放置在这些

micro-kernel

参考文章


本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!