GEMM是General Matrix Multiply的缩写，即通用矩阵乘法。矩阵乘法的应用场景很多，优化方法也很多，但具体优化应该根据具体的硬件资源来进行优化，本篇笔记用来记录学习在arm架构上使用neon intrinsic函数实现矩阵乘优化。
1、按行存储和按列存储
虽然我们定义数组时，直观上认为在内存中是按照行存储的，但是不同的硬件架构、编程规范以及编译器设计都会影响数组在内存中的实际存储方式。定义的数组在内存中实际的存储方式由按行存储（row-major）和按列存储（column-major）。在Fortran编程当中数组是默认按列存储的，而在C/C++当中通常是按行存储的。
假设在定义一个数据array[M][N]后，这个数组在按行存储和按列存储是不一样的。但都是将高维数组映射为一维，然后按行或按列进行存储的。按行存储的架构中，在读取数据时，是按行读取的，按列存储的架构中是按照列存取的，这样做是为了让加载的数据能够保存在缓存当中，提升下次访存的命中。
在程序中定义一个M行N列的一维的数组array后，如果需要访问第i行第j列的元素，那么不同的存储方式访问的方式也是不同的。按行存储的应该使用访问语句array[i*N+j]。按列存储的使用访问语句array[j*M+i]。
下面开始实现通用矩阵乘法的练习。使用的矩阵的按行存储的。平台是jetson nano，搭载armv8-A架构cpu，和一块Maxwell架构gpu。数组大小统一使用640*640进行实验。
2.naive
使用C语言很容易实现两个矩阵的乘法。使用一个循环遍历矩阵A的行，一个循环遍历B的列，然后一个循环遍历当前A的行和当前B的列，相乘后相加，最后保存到矩阵C的对应位置中。
void naive(const float *A,const float *B,float *C,const int N){
for(int i=0;i<N;i++){
for(int j=0;j<N;j++){
float Cij=0;
for(int k=0;k<N;k++){
Cij+=A[i*N+k]*B[k*N+j];
}
C[i*N+j]=Cij;
}
}
}
在不打开任何编译器优化设置的时候，运行时间主要是5秒左右。
在打开-O3优化以后，运行时间将至2秒左右。
3.neon intrinsic函数
由于我们使用的数组默认是按行存储的，在对矩阵A进行读取时，能够很好的利用带宽，并且能够将取到的一行数据长时间的保存在缓存上，在进行一次读取矩阵A的一行以后，后面读取该行就会很高效。但是，在对矩阵B读取时，由于我们需要读取的是每一列的数据，但是B的存储也是按照行存储的，所以不能够很好的利用好带宽和缓存，因此读取B的效率是较慢的。
但是可以使用arm架构扩展的neon指令集，能够并行读取B的多个数据，那么就可以一次读取B中一行的多个数据。如果是float类型数据，使用neon 128位寄存器可以存放4个数据，也就是说读取一次可以读4列，那么在读取矩阵B的数据时，会比naive方法快4倍以上。
以下是使用neon intrinsic函数实现通用矩阵乘法，一个循环遍历矩阵A的行，一个循环按照4个float元素为一组遍历，然后循环当前A的行和B的列组，进行乘加。
void GEMM_neon(const float *A,const float *B,float *C,const int N){
for(int i=0;i<N;i++){
int len=N/4;
for(int j=0;j<N;j+=4){
float32x4_t buffer=vdupq_n_f32(0);
for(int k=0;k<N;k++){
float32x4_t vb=vld1q_f32(&B[k*N+j]);
register float va=A[i*N+k];
buffer=vmlaq_n_f32(buffer,vb,va);
}
vst1q_f32(&C[i*N+j],buffer);
}
}
}
不打开-O3的运行时间如下：
打开-O3后的运行时间为：
上述代码在读取矩阵B的时候使用了neon寄存器，即一次性能够计算四列的数据。但是在矩阵A的读取上，我们一次只读取了一个数据，其实这样在矩阵A上面的效率是较低的，下面实现一种使用neon寄存器读取矩阵A中一行数据的四个元素，这样在读取矩阵A时的并行性也增大了。但是还不够，其实我们可以使用4个寄存器来同时保存一个矩阵A中的小矩阵（4*4），这样可以同时处理目标矩阵C中的多行/多列的数据。
首先需要实现的是一个4*4的矩阵乘法，因为大矩阵都可以划分为小矩阵的乘法来迭代实现，而为什么选取4作为基本尺度呢？那是因为neon寄存器的大小是128bit，只能够存储4个float数值，另外就是4*4的矩阵乘法大量使用在仿真计算等领域（我也不了解。。）。
下面是实现的两个4*4矩阵的乘法：
void matrix_multiply_4x4_neon(float *A, float *B, float *C) {
// these are the rows A
float32x4_t A0;
float32x4_t A1;
float32x4_t A2;
float32x4_t A3;
// these are the rows B
float32x4_t B0;
float32x4_t B1;
float32x4_t B2;
float32x4_t B3;
// these are the rows C
float32x4_t C0;
float32x4_t C1;
float32x4_t C2;
float32x4_t C3;
// Load B matrix by rows
B0 = vld1q_f32(B);
B1 = vld1q_f32(B + 4);
B2 = vld1q_f32(B + 8);
B3 = vld1q_f32(B + 12);
// Zero accumulators for C values
C0 = vmovq_n_f32(0);
C1 = vmovq_n_f32(0);
C2 = vmovq_n_f32(0);
C3 = vmovq_n_f32(0);
// Multiply accumulate in 4x1 blocks
A0 = vld1q_f32(A);
C0 = vfmaq_laneq_f32(C0, B0, A0, 0);
C0 = vfmaq_laneq_f32(C0, B1, A0, 1);
C0 = vfmaq_laneq_f32(C0, B2, A0, 2);
C0 = vfmaq_laneq_f32(C0, B3, A0, 3);
vst1q_f32(C, C0);
A1 = vld1q_f32(A+4);
C1 = vfmaq_laneq_f32(C1, B0, A1, 0);
C1 = vfmaq_laneq_f32(C1, B1, A1, 1);
C1 = vfmaq_laneq_f32(C1, B2, A1, 2);
C1 = vfmaq_laneq_f32(C1, B3, A1, 3);
vst1q_f32(C + 4, C1);
A2 = vld1q_f32(A + 8);
C2 = vfmaq_laneq_f32(C2, B0, A2, 0);
C2 = vfmaq_laneq_f32(C2, B1, A2, 1);
C2 = vfmaq_laneq_f32(C2, B2, A2, 2);
C2 = vfmaq_laneq_f32(C2, B3, A2, 3);
vst1q_f32(C + 8, C2);
A3 = vld1q_f32(A + 12);
C3 = vfmaq_laneq_f32(C3, B0, A3, 0);
C3 = vfmaq_laneq_f32(C3, B1, A3, 1);
C3 = vfmaq_laneq_f32(C3, B2, A3, 2);
C3 = vfmaq_laneq_f32(C3, B3, A3, 3);
vst1q_f32(C + 12, C3);
}
首先使用寄存器来分别保存4*4的矩阵A和矩阵B，然后使用一个intrinsic函数来实现这个矩阵的加法。vfmaq_laneq_f32()；是一个乘法跟加法的融合计算函数，将某个 float32x4_t 与另外一个 float32x4_t 中的单个数据进行 broadcast 乘法，将结果跟第三个 float32x4_t 相加。这个函数的理解其实优点难的。
A0 = vld1q_f32(A);
C0 = vfmaq_laneq_f32(C0, B0, A0, 0);
C0 = vfmaq_laneq_f32(C0, B1, A0, 1);
C0 = vfmaq_laneq_f32(C0, B2, A0, 2);
C0 = vfmaq_laneq_f32(C0, B3, A0, 3);
vst1q_f32(C, C0);
这段代码是使用矩阵A的第一行来和整个矩阵矩阵相乘。A的第一行先和B的第一行按照A的通道为1的元素来进行广播相乘，因为B0就是矩阵B的第一行，按照矩阵的乘法，只能和A的第一行的第一个元素相乘，然后累加到寄存器c0中，这样将B的四个行遍历结束后，就能计算出C矩阵中的一行了。
然后将上面的思想用到一张640*640的数组中计算，其实就是对矩阵A的行进行4个为一组的分行块，对矩阵B的列进行分列块，每一次分好的块进行4*4的矩阵计算，就能够迭代出整个矩阵C的结果了。具体代码如下：
void matrix_multiply_640x640_neon(float *A, float *B, float *C,const int N){
for(int i=0;i<N;i+=4){
for(int j=0;j<N;j+=4){
float32x4_t c0=vmovq_n_f32(0);
float32x4_t c1=vmovq_n_f32(0);
float32x4_t c2=vmovq_n_f32(0);
float32x4_t c3=vmovq_n_f32(0);
for(int k=0;k<N;k+=4){
float32x4_t b0=vld1q_f32(B);
float32x4_t b1=vld1q_f32(B+N);
float32x4_t b2=vld1q_f32(B+2*N);
float32x4_t b3=vld1q_f32(B+3*N);
float32x4_t a0=vld1q_f32(A);
c0=vfmaq_laneq_f32(c0, b0, a0, 0);
c0=vfmaq_laneq_f32(c0, b1, a0, 1);
c0=vfmaq_laneq_f32(c0, b2, a0, 2);
c0=vfmaq_laneq_f32(c0, b3, a0, 3);
vst1q_f32(&C[i*N+k],c0);
float32x4_t a1=vld1q_f32(A+N);
c1=vfmaq_laneq_f32(c1, b0, a1, 0);
c1=vfmaq_laneq_f32(c1, b1, a1, 1);
c1=vfmaq_laneq_f32(c1, b2, a1, 2);
c1=vfmaq_laneq_f32(c1, b3, a1, 3);
vst1q_f32(&C[(i+1)*N+k],c1);
float32x4_t a2=vld1q_f32(A+2*N);
c2=vfmaq_laneq_f32(c2, b0, a2, 0);
c2=vfmaq_laneq_f32(c2, b1, a2, 1);
c2=vfmaq_laneq_f32(c2, b2, a2, 2);
c2=vfmaq_laneq_f32(c2, b3, a2, 3);
vst1q_f32(&C[(i+2)*N+k],c2);
float32x4_t a3=vld1q_f32(A+3*N);
c3=vfmaq_laneq_f32(c3, b0, a3, 0);
c3=vfmaq_laneq_f32(c3, b1, a3, 1);
c3=vfmaq_laneq_f32(c3, b2, a3, 2);
c3=vfmaq_laneq_f32(c3, b3, a3, 3);
vst1q_f32(&C[(i+3)*N+k],c3);
}
}
}
}
实现结果：
在不打开-O3优化的情况下，实现640*640矩阵乘法的时间不到1.5s。
在打开-O3优化的时候，时间大幅度降低，不到0.1s。
目前只能优化到这里，但感觉还可以在空间局部性上面进行优化，后面有思路了再来更新。
ps：能力有限，如果不理解可参考链接。
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
4.cuda实现
更新一下使用GPU端进行矩阵乘的计算。
与CPU计算的不同点主要是参与计算的核心变多了，不需要使用仅有的几个计算核心来顺序的执行。可以使用一个核心负责计算结果矩阵中的一个值，也就是计算矩阵A B的某一行和某一列的乘加。
下面是代码：
__global__
void matrix_mul(float *A,float *B,float *C,int const M,const int K,int const N){
int row=blockIdx.y*blockDim.y+threadIdx.y;
int col=blockIdx.x*blockDim.x+threadIdx.x;
__shared__ float tile_A[BLOCK_SIZE][BLOCK_SIZE];
__shared__ float tile_B[BLOCK_SIZE][BLOCK_SIZE];
float sum=0;
for(int stride=0;stride<=K/BLOCK_SIZE;++stride){
int id_m=row*K+stride*BLOCK_SIZE+threadIdx.x;
if(row<M&&stride*BLOCK_SIZE+threadIdx.x<K){
tile_A[threadIdx.y][threadIdx.x]=A[id_m];
}
__syncthreads();
int id_n=(stride*BLOCK_SIZE+threadIdx.y)*N+col;
if(col<N&&stride*BLOCK_SIZE+threadIdx.y<K){
tile_B[threadIdx.y][threadIdx.x]=B[id_n];
}
__syncthreads();
for(int i=0;i<BLOCK_SIZE;++i){
sum+=tile_A[threadIdx.y][i]*tile_B[i][threadIdx.x];
}
__syncthreads();
}
if(row<M&&col<N){
C[row*N+col]=sum;
}
}
使用的计算数组A和B和之前cpu端使用的数组大小一样，都是640*640，共享内存大小和线程组织的块大小一样，都是32*32，方便进行数据传输。计算平台是jetson nano上的一块Maxwell架构的gpu。
使用共享内存需要注意的是bank冲突。也就是当一个warp内的多个线程访问同一个bank的不同地址时，就会bank冲突，那么这部分冲突的线程就会串行执行了。会影响计算性能。在下面代码中，我们知道，在一个block中，每个warp内的线程一般有一样的threadIdx.y和连续的threadIdx.x，那么我们可以利用这样一个性质，使用不同的threadIdx.x来访问共享内存，这样可以有效缓解bank冲突。
tile_A[threadIdx.y][threadIdx.x]
使用nvprof进行性能分析如下：
可以看到matrix_mul这个核函数的运行时间是23.787ms，与之前使用neon intrinsic函数进行处理所消耗的81ms减少了接近60ms。
3280
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
