cuda矩阵乘法优化之路

吃肉夹馍不要夹馍

已于 2024-10-17 11:26:32 修改

阅读量1.1k

点赞数 11

分类专栏：读书笔记算法文章标签：矩阵线性代数 cuda c++

于 2024-10-12 11:14:34 首次发布

本文链接：https://blog.youkuaiyun.com/sinat_41053216/article/details/142873141

版权

读书笔记同时被 2 个专栏收录

12 篇文章

订阅专栏

算法

7 篇文章

订阅专栏

cuda矩阵乘法优化之路

参考：

矩阵乘法的 CUDA 实现、优化及性能分析

CUDA 矩阵乘法终极优化指南

CUDA 实践：矩阵乘法 | CUDA

看了网上一些矩阵乘法，很多尽管给了思路，也给了一部分代码，但是要么没有启动配置，要么无法运行，依然需要自己深入理解。
于是自己记录一下，cuda矩阵乘法的优化方法。会给出详细的原版kernel代码，以及对应的启动配置。

矩阵乘法，首先最朴素的方法，当然是类似CPU的直接计算，不过相对于CPU而言，GPU可以直接进行并行计算，对于矩阵乘 $P = MN$ ,其中 $\in R^{m\times n}$ , $\in R^{m\times k}, N \in R^{k \times n}$

这里为了方便，直接使用方阵 m = n

朴素矩阵乘法的kernel实现：

__global__ void matrix_mul_gpu1(double *M, double* N, double* P, int width)
{
    int i = threadIdx.x + blockDim.x * blockIdx.x;
    int j = threadIdx.y + blockDim.y * blockIdx.y;
                
    double sum = 0;
    for(int k=0;k<width;k++)
    {
        int a = M[j*width+k];
        int b = N[k*width+i];
        sum += a*b;
    }
    P[j*width+i] = sum;
}

每一个线程计算P中的一个元素。在我的机器上(Intel® Core™ i7-8700K CPU @ 3.70GHz + nvidia 1080) 这种实现性能已经和CPUopenblas差不多了

在这里插入图片描述

基准对比

分别使用openblas和cublas的矩阵乘作为计算结果和计算性能的基准对比

openblas：

double* matrix_mul_cpu(double *M, double *N){
    /*
    使用CPU openblas 实现矩阵乘法
    编译时添加 -lopenblas
    */
    double *P = new double[Row*Col];
    clock_t start=clock();
    cblas_dgemm(
        CblasRowMajor,
        CblasNoTrans,
        CblasNoTrans,
        Row,
        Col,
        Col,
        1.0,
        M,
        Row,
        N,
        Col,
        0.0,
        P,
        Col
    );
    std::cout<<" cpu openblas cost: "<<static_cast<double>((clock() - start)/1000)<<" ms"<<std::endl;
    return P;
}

cublas:
cublas是列优先存储，但是CPU以行有限相当于进行了转置，输入的A，B cublas里实际是 $A^T$ ,$ B^T$,
$B^T * A^T = (A*B)^T = C^T$ 而输出也是列优先，这时候再按照行优先读取，正好得到 $C^T)^T = C$
（注意编译时需要添加 -lcublas）

double* matrix_mul_0(double *M, double*N){
    double* P = new double[Row*Col];
    double *dev_m, *dev_n, *dev_p;
    cudaMalloc((void**)&dev_m, sizeof(double)*Row*Col);
    cudaMalloc((void**)&dev_n, sizeof(double)*Row*Col);
    cudaMalloc((void**)&dev_p, sizeof(double)*Row*Col);
    cudaMemcpy(dev_m, M, sizeof(double)*Row*Col, cudaMemcpyHostToDevice);
    cudaMemcpy(dev_n, N, sizeof(double)*Row*Col, cudaMemcpyHostToDevice);

    double alpha=1.0;
    double beta=0.0;
    cublasHandle_t handle;
    cublasCreate(&handle);
    cublasDgemm(
        handle,
        CUBLAS_OP_N,CUBLAS_OP_N,
        Col,
        Row,
        Col,
        &alpha,
        dev_n,
        Col,
        dev_m,
        Row,
        &beta,
        dev_p,
        Col
    );
    cublasDestroy(handle);

    cudaMemcpy(P, dev_p, sizeof(double) * Row * Col, cudaMemcpyDeviceToHost);

    cudaFree(dev_m);
    cudaFree(dev_n);
    cudaFree(dev_p);
    return P;
}

cuda优化一：共享内存

这个应该也是比较普遍的，在cuda学习书中应该都有介绍过，这里安利一下PMPP。

raw实现中，导致性能低下的原因之一是，每个线程计算时候，都直接从global memory中取值:

for(int k=0;k<width;k++){
    int a = M[j*width+k];
    int b = N[k*width+i];
    sum += a*b;
}

nvidia显卡中，global内存是存储在DRAM内存中，读取效率很慢，可以改用share memory，共享内存的读取速度相对DRAM更快，并且共享内存对线程块中的所有线程可见，在raw gemm中，假设一共有m*n个线程，那么对global memory的读取次数有 $m\times n \times k \times2$ ，使用共享内存之后，由于块内内存对线程共享，所以对global memory读取次数将减少一倍

在这里插入图片描述

假设一个线程块负责计算结果P中的一个 $BLOCK\_SIZE \times BLOCK\_SIZE$ 的子矩阵块 $P_{ij}$ ，

其实等于将矩阵M中一个子矩阵 $M_i, M_i \in R^{BLOCK\_SIZE \times width}$ 乘以N中的子矩阵 $N_j,N_j \in R^{width \times BLOCK\_SIZE}$ , 由于共享内存大小常常是有限的，与其直接读取 $M_i, N_j$ ，不如以一个小方块，分别在 $M_i, N_j$ 中沿着列增大和行增大的防线滑动读取，存到共享内存中，计算结果之后滑动到下一步，最后将结果累加。

共享内存的gemm：

__global__ void matrix_mul_gpu2(double *M, double* N, double* P, int width){
    /*
    share memory
    */
   int row = blockIdx.y*BLOCK_SIZE + threadIdx.y;
   int col = blockIdx.x*BLOCK_SIZE + threadIdx.x;

   __shared__ double share_M[BLOCK_SIZE][BLOCK_SIZE];
   __shared__ double share_N[BLOCK_SIZE][BLOCK_SIZE];

   uint tiles = (width)/BLOCK_SIZE;
   double value=0;
   for(uint i=0; i<tiles; i++){
    share_M[threadIdx.y][threadIdx.x] = M[row*width + i*BLOCK_SIZE + threadIdx.x];
    share_N[threadIdx.y][threadIdx.x] = N[(i*BLOCK_SIZE + threadIdx.y)*width + col];
    __syncthreads();

    #pragma unroll
    for(uint k =0;k<BLOCK_SIZE; k++){
        value += share_M[threadIdx.y][k] * share_N[k][threadIdx.x];
    }
    __syncthreads();
   }

   P[row*width + col] = value;
}

相比于 raw gemm, 使用共享内存之后性能提升了不少：

在这里插入图片描述

cuda优化二：计算多个值

在上面共享内存中，优化了数据内存读取，然而仍然是一个线程计算一个值，一次读取和计算的比例很低，在PMPP中提到一个启发式的思路，即一个线程计算两个P中元素

在这里插入图片描述

相对于1中的共享内存优化中，相当于每加载一个BLOCK_SIZE* width 的$M_i $ 时候，从N中加载两个 width * BLOCK_SIZE ，重复利用了Mi中读取的值，其实这也是PMPP中的一个课后习题。

一个线程计算2个元素kernel：

注意因为每个线程加载了两个 N元素，所以在X方向上的线程块数量要减少一半

__global__ void matrix_mul_gpu3(double *M, double* N, double* P, int width){
    /*
    share memory
    for one row M, two column N
    */
    unsigned int bx = blockIdx.x;
    unsigned int by = blockIdx.y;
    unsigned int tx = threadIdx.x;
    unsigned int ty = threadIdx.y;
    __shared__ double Ms[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ double Ns[BLOCK_SIZE][BLOCK_SIZE * 2];
    // M 的第一个子矩阵左上元素索引
    int mBegin = width * BLOCK_SIZE * by;
    // M 的最后一个子矩阵右上元素索引，用来控制循环何时结束
    int mEnd = mBegin + width - 1;
    // M 子矩阵每次循环需要加上的距离
    int mStep = BLOCK_SIZE;
    // N 的第一个子矩阵左上元素索引
    int nBegin = BLOCK_SIZE * bx;
    // N 子矩阵每次循环需要加上的距离
    int nStep = BLOCK_SIZE * width;
	//Psub 用来存储结果矩阵每块元素的部分和
    int nMid = width/2;

    double Psub1 = 0.0;
    double Psub2 = 0.0;
    // 遍历所有子矩阵
    for (int m = mBegin, n = nBegin; m <= mEnd; m += mStep, n += nStep)
    {
        // 将数组加载到共享存储器，每个线程加载 3 个元素
        Ms[ty][tx] = M[m + width * ty + tx];
        Ns[ty][tx] = N[n + width * ty + tx];

        if((n+width*ty+tx+nMid<width*width)&&(bx*BLOCK_SIZE+tx+nMid<width)){
            Ns[ty][tx + BLOCK_SIZE] = N[n+width*ty + tx + nMid];
        }else{
            printf("bx:%d,n:%d,tx:%d, width: %d, %d\n",bx,n,tx,width, n+tx+nMid);
        }

        __syncthreads();
        // 一次计算两个矩阵
        for (int k = 0; k < BLOCK_SIZE; ++k)
        {
            Psub1 += Ms[ty][k] * Ns[k][tx];
            Psub2 += Ms[ty][k] * Ns[k][tx + blockDim.x];
        }

        __syncthreads();
    }
    // 写回计算结果，每个线程写一个
    int idx = width*(by*BLOCK_SIZE + ty) + BLOCK_SIZE*bx + tx;
    P[idx] = Psub1;
    if(idx + nMid<width*width){
        P[idx + nMid] = Psub2;
    }
}

对比之前共享内存的结果：

在这里插入图片描述

性能还是有所提升的。

但是仍然有优化的空间，考虑一个线程分别从M，N中加载4个值，为什么4是因为有LD.128可以读取4个值。则一个线程也可以计算P中的4个值。

假设一次从M中加载了4个值，则这4个值可以被N中的4列数据共用。

在这里插入图片描述

反过来，以此类推，N中的一列数据，也可以被M中的4行数据共用

在这里插入图片描述

所以，每个线程读取4*4的M和 4*4的N 可以计算 4*4的P。

这样可以大幅提高每个线程计算与io的比值。

kernel

#define T_THREAD 4
__global__ void matrix_mul_gpu5(double *M, double* N, double* P, int width){
    /*
    share memory 
    compute 16 value every thread
    */
   const int share_size_m = BLOCK_SIZE*T_THREAD;
   const int share_size_k = BLOCK_SIZE*T_THREAD;
   const int share_size_n = BLOCK_SIZE*T_THREAD;
    __shared__ double share_M[share_size_m][share_size_k];
    __shared__ double share_N[share_size_k][share_size_n];

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    // printf("bx %d, by:%d,tx:%d, ty:%d, width:%d\n",bx,by,tx,ty,width);

    double C_value[T_THREAD][T_THREAD] = {0.0};
    for(int bk=0; bk<width; bk+=share_size_k){
        // load m, n
        for(int i=0;i<T_THREAD;i++){
            for(int j=0;j<T_THREAD;j++){
                share_M[ty*T_THREAD+i][tx*T_THREAD+j] = M[(share_size_m*by+ty*T_THREAD+i)*width + bk + tx*T_THREAD+j];
                share_N[ty*T_THREAD+i][tx*T_THREAD+j] = N[(bk+ty*T_THREAD+i)*width + share_size_n*bx + tx*T_THREAD+j];
            }
        }
        __syncthreads();
        
        for (int k = 0; k < share_size_k; k++){
            for (int i = 0; i < T_THREAD; i++){
                for (int j = 0; j < T_THREAD; j++){
                    C_value[i][j] += share_M[ty*T_THREAD+i][k] * share_N[k][tx*T_THREAD+j];
                }
            }            
        }
        __syncthreads();
    }

    for (int i = 0; i < T_THREAD; i++)
    {
        for (int j = 0; j < T_THREAD; j++)
        {
            P[(share_size_m*by + ty*T_THREAD+i)*width + share_size_n*bx+tx*T_THREAD+j] = C_value[i][j];
        }
    }
}

最后的结果也比较喜人：

在这里插入图片描述

在我的机器上测试，可能有雨cuda版本不高或者其他什么别的原因，性能基本打平了cublas。

看到其他资料中，还提到了有L128读取，以及pragma unroll，理论上可以优化指令读取，不过在我测试的这机器上，并没有发现能显著提高性能，将之也加入吧，留个坑以后看看是什么原因

#define DOUBLE4(pointer) (reinterpret_cast<double4*>(&(pointer))[0])

__global__ void matrix_mul_gpu6(double *M, double* N, double* P, int width){
    /*
    share memory 
    compute 16 value every thread
    use  LDG.128 load memory
    */
   const int share_size_m = BLOCK_SIZE*T_THREAD;
   const int share_size_k = BLOCK_SIZE*T_THREAD;
   const int share_size_n = BLOCK_SIZE*T_THREAD;
    __shared__ double share_M[share_size_m][share_size_k];
    __shared__ double share_N[share_size_k][share_size_n];

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    // printf("bx %d, by:%d,tx:%d, ty:%d, width:%d\n",bx,by,tx,ty,width);

    double C_value[T_THREAD][T_THREAD] = {0.0};
    for(int bk=0; bk<width; bk+=share_size_k){
        // load m, n
        for(int i=0;i<T_THREAD;i++){
            DOUBLE4(share_M[ty*T_THREAD+i][tx*T_THREAD]) = DOUBLE4(M[(share_size_m*by+ty*T_THREAD+i)*width + bk + tx*T_THREAD]);
            DOUBLE4(share_N[ty*T_THREAD+i][tx*T_THREAD]) = DOUBLE4(N[(bk+ty*T_THREAD+i)*width + share_size_n*bx + tx*T_THREAD]);
        }
        __syncthreads();
        
        for (int k = 0; k < share_size_k; k++){
            #pragma unroll
            for (int i = 0; i < T_THREAD; i++){
                #pragma unroll
                for (int j = 0; j < T_THREAD; j++){
                    C_value[i][j] = fma(share_M[ty*T_THREAD+i][k], share_N[k][tx*T_THREAD+j], C_value[i][j]);
                }
            }            
        }
        __syncthreads();
    }

    for (int i = 0; i < T_THREAD; i++)
    {
        DOUBLE4(P[(share_size_m*by + ty*T_THREAD+i)*width + share_size_n*bx+tx*T_THREAD]) = DOUBLE4(C_value[i][0]);
    }
}