GEMM-intrinsic优化

写在前面:

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

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

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

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

内积和外积

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

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

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

// 假设矩阵的存储是行优先/行主序的 
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阵。这种矩阵乘的方式叫做外积，伪码如下：

// 假设矩阵的存储是按行优先存储的
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

矩阵分块

针对行主序的矩阵进行计算,其中$A_{M×K}$$B_{K×N}$,$C_{M×N}$我们以BLAS(基本的线性代数函数接口规范)有:

$$C = \alpha AB + \beta C$$

接口声明如下:

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

micro-kernel

参考文章

1. intel-intrinsic ↩
2. x64 CPU GEMM优化 ↩
3. CPU流水线并行 ↩
4. (简易的GEMM优化)[https://github.com/flame/how-to-optimize-gemm] ↩