CUDA矩阵乘法优化

前言

纸上的来终觉浅，绝知此事要躬行。

naive写法

一个矩阵的乘法简单如下：C=A*B, 一般用gemm(A,B,C,M,N,K)来表示，其中的m,n,k代表的位置如下，默认是k表示消失的纬度。

上图的红色虚线围起来的是一个block要负责的数据区域，具体的代码如下：

__global__
void matrixMul(const float *A, const float *B, float *C, int M, int N, int K)
{int col = blockIdx.x * blockDim.x + threadIdx.x;int row = blockIdx.y * blockDim.y + threadIdx.y;if(row < M && col < N) {float c = 0;for(int i = 0; i < K; ++i){c += A[row * K + i] * B[i * N + col];}C[row * N + col] = c;}
}

第一步先写一个简单的矩阵乘法，不要觉得简单就不写，后面不管怎么耍花招，都是基于当下的核心逻辑和计算流程的。

share memory优化

代码我们可以看到，每一个线程要k次乘法和K次加法，每做计算一次 FMA（乘累加）之前需要读一次 A 和读一次 B，读取 Global Memory 的代价很大，通常都需要几百个 cycle（时钟周期），而计算一次 FMA 通常只需要几个 cycle，大量的时间被花费在了访存上，那么如何减少访问global memory呢？我们知道share_memory是片上内存，访问速度比较快，我们考虑把A和B中的虚线内的数据放在share_memory上，然后计算的时候从share_memeory上取，这样的话其实会多一次数据从global转移到share上的操作，但是每次做乘加计算取数的时候，省的时间会远远多于这一次操作。

上面我们假设是把A、B中的虚线部分数据导入到share中，但其实share的大小有限，所以还是继续分片，原理如下：

具体的伪代码是：

//分为j块
for(int i = 0; i < j; i+1)
{load_gmem_to_smem(A, i, smemA);load_gmem_to_smem(B, i, smemB);__syncthreads();//computeC_i=gemm(a_i,b_i);//累加C += C_i;}

可以运行的代码：

template <int BLOCK_SIZE>
__global__
void MatrixMulCUDA(const float * __restrict__ A,const float * __restrict__ B,float * __restrict__ C,const int M,const int K,const int N) {// Block indexint bx = blockIdx.x;int by = blockIdx.y;// Thread indexint tx = threadIdx.x;int ty = threadIdx.y;//对应上图a2的左上角位置int aBegin = K * BLOCK_SIZE * by;//对应上图aj的右上角位置（注意上图有溢出到A矩阵外，这里代码默认是刚刚在A矩阵里的）int aEnd   = aBegin + K - 1;//每一个a矩阵的左上角位置距离差int aStep  = BLOCK_SIZE;// 对应上图b2的左上角位置int bBegin = BLOCK_SIZE * bx;// 每一个b矩阵的左上角位置距离差int bStep  = BLOCK_SIZE * N;float Csub = 0;// 存储aj的share memory__shared__ float As[BLOCK_SIZE][BLOCK_SIZE];// 存储bj的share memory__shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];//对应上图就是循环3次for (int a = aBegin, b = bBegin; a < aEnd; a += aStep, b += bStep) {//每个线程下载一个元素到share memoryAs[ty][tx] = A[a + K * ty + tx];Bs[ty][tx] = B[b + N * ty + tx];// Synchronize 为了确保share memory需要的数据都下载到位__syncthreads();// 每个线程计算cj中的一个位置的结果#pragma unrollfor (int k = 0; k < BLOCK_SIZE; ++k) {Csub = fma(As[ty][k], Bs[k][tx], Csub);}// 确保cj中的所有位置的结果都算好了__syncthreads();}//最后把每个线程负责的算的结果从寄存器导入到device memoryC[N * ( BLOCK_SIZE * by + ty ) + BLOCK_SIZE * bx + tx ] = Csub;
}

实验结果如下，当M=160，K=160, N=160的时候，实现效果显示还不错，比naive好但是比cublas还是有差距。

第二次优化

第一次优化后，访存代价从几百 cycle 降低到几十 cycle，并不改变问题的本质。问题的关键在于主体循环由两条 Load 指令与一条 FMA 指令构成，计算指令只占总体的 1/3，计算访存比过低，最终导致了访存延迟不能被隐藏，从而性能不理想.

这里解释一下，线程的指令会发给调度器，调度器分配对应的执行单元。这里比如有20个线程，机器上一个调度器，一个计算单元和访存单元，此刻线程1告诉调度器要执行加法计算，计算单元需要10s可以得到结果, 等待期间调度器就会问问其他线程有需要访存的吗？毕竟有一个访存单元在闲着，这时候线程8说他需要访存操作，不过一次访存需要200s, 因为计算单元速度很快，所以当20个线程的计算任务都完成时，只有一个线程的访存任务完成，所以后面还要200s*19这么长的时间。这里可以发现一共需要200S*20的时间，我们计算时间10s*20被完成隐藏了用户感知不到，这就是所谓的隐藏延迟，知道了原理，我们的任务说白了就是别让计算单元闲着，如果一个线程的计算时间如果和访存时间一模一样，那么完全就可以隐藏计算或者是访存的时间了，这不是美滋滋？但是但是，这里的前提是只有一个访存单元和计算单元，实际上底层硬件还是差距很大的，不同架构和型号的显卡也不一样，这也是为啥同样的的代码在不同机器上的性能不一样，此外调度器的规则，计算与访存任务的依赖等等都有可能导致性能差异。

float c[4][4] = {{0}};
float a_reg[4];
float b_reg[4];
for(int i = 0; i < K; i += tile_with)
{load_gmem_tile_to_smem(A, i, smemA);load_gmem_tile_to_smem(B, i, smemB);__syncthreads();//computefor(int j = 0; j < tile_with; ++j) {// load tile from shared mem to register load_smem_tile_to_reg(smemA, j, a_reg);load_smem_tile_to_reg(smemB, j, b_reg);// compute matrix multiply accumulate 4x4mma4x4(a_reg, b_reg, c)；}// 累加C += C_i;
}

未完待续。。。
参考：https://zhuanlan.zhihu.com/p/410278370