cuda编程-矩阵乘法优化

上一篇blog实现了一个基本的cuda程序，本篇记录如何实现矩阵乘法并加速。主要内容是对blog内容的学习和复现，目的是逐步逼近cuBLAS实现的单精度矩阵乘法：

$$\mathbf{C} = \alpha \mathbf{A}\mathbf{B} + \beta \mathbf{C}$$

其中$\mathbf{A}\in \mathbb{R}^{M \times K},\mathbf{B}\in \mathbb{R}^{K \times N},\mathbf{C}\in \mathbb{R}^{M \times N}$。源代码在这里。

Naive vs cuBLAS

因为$\mathbf{C}{i,j} = \alpha \sum_{k=0}^{K-1} \mathbf{A}_{i,k} \mathbf{B}_{k,j} + \beta \mathbf{C}{i,j}$，很容易我们可以想到为$\mathbf{C}$的每个元素分配一个thread计算值并输出：

__global__ void sgemm_naive(
    int M, int N, int K, const float alpha, const float *A, const float *B, const float beta, float *C)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    const int j = blockIdx.y * blockDim.y + threadIdx.y;

    // avoids memory access error if threads are more than elements
    if (i < M && j < N)
    {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k)
        {
            sum += A[i * K + k] * B[k * N + j];
        }
        C[i * N + j] = alpha * sum + beta * C[i * N + j];
    }
}

dim3 gridDim((M - 1) / 32 + 1, (N - 1) / 32 + 1);
dim3 blockDim(32, 32, 1);
sgemm_naive<<<gridDim, blockDim>>>(M, N, K, alpha, d_A, d_B, beta, d_C);

block的数目要确保了能覆盖$\mathbf{C}$的所有元素，当$M,N$不能被32（这里简单的让x和y方向上的thread数目相等，当然也可以分别选择不同的值）整除时，边缘处总有一些threads在空跑，这是无法避免的。sgemm_naive的各项指标如下：

• 总浮点数操作：$MN(2K+3)$
• 总显存访问次数：$MN(2K+1)$
• 总显存占用：$(MN+MK+KM)*4B$

当$M=N=K=4096$时，总浮点数操作为137.5GFLOPS，GTX1080的FP32运算能力是8.873TFLOPS，最好情况下可以在15.5ms左右完成sgemm。要存储矩阵并完成计算，理想情况下总显存占用为201.3MB，GTX1080显存8GB，因此显存大小是没有问题的。不过，总显存访问次数计算出的显存通量是549.8GB，而GTX1080的带宽峰值是320GB/s，因此完成这么多次显存访问总共需要1.718s，远远大于实际计算时间。sgemm_naive的运行时间完全由显存访问主导，大部分时间GPU核心都在等待数据。实际运行中只会更慢，在我的电脑上，sgemm_naive大概需要2900ms。

而cublasSgemm呢，只需要大约20ms，接近15ms的理论极限，快的惊人！

Memory Coalescing

在基本概念中我们提到GPU会把block中每32个thread组成一个warp交由硬件调度，考虑一个32x32的block，那么到底是row上的32个threads组成warp还是column上的32个threads组成warp？事实上，针对多维情况，会首先转化成1维，并取连续的32个threads组成warp。对于三元组（x,y,z）CUDA的规则是x优先，每个block中的全局threadId由下式确定：

$$threadId = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z$$

warp的特点之一是如果这些threads读取的是同样的数据，会触发一种within-warp broadcast的共享机制，实际上不需要每个thread都重新读取数据，因此会显著减小开销。另一个特点是这些threads读取显存时可以一次性读取大量连续数据（比如32B、64B、128B）到更快的内存上（例如L2 cache）。例如一个warp中的32个thread可以选择各自读取1个浮点数（4B），需要操作32次，但如果他们读的是一块连续的内存，实际上可以用一个128B的读取指令等价，只需操作一次，直接加速32倍。满足memory coalescing的程序就会去尽量利用这种机制，让一个warp里的threads去利用一块相近的而内存区域。

那么sgemm_naive是不是memory coalescing的呢？首先我们要明确，在当前的例子中，C++采取row-major的存储顺序，即矩阵的每一行在内存中是连续的。我们当前的block大小是32x32，正好每32个threads组成一个warp。而这个wrap中的threads都是threadIdx.y不变而threadIdx.x递增。根据我们对于每个thread与元素的映射规则：

const int i = blockIdx.x * blockDim.x + threadIdx.x;
const int j = blockIdx.y * blockDim.y + threadIdx.y;

i递增而j不变，因此程序会不断读取$\mathbf{B}$的第j列和$\mathbf{A}$不同i行。根据warp的特点，读取$\mathbf{B}$的列会触发within-warp broadcast，因此尽管列中的元素不是连续存储的，threads读取这些元素的开销很小；但是对于$\mathbf{A}$，尽管每一行是连续的内存，但每个thread用的是来自不同行的数据，同一时刻读取的元素（每一列中的32个元素）在内存里是不连续的，自然无法满足memory coalescing的要求。

因此，应当尽量使threadIdx.x的变化同内存排序一致，我们只需要交换xy的方向：

const int j = blockIdx.x * blockDim.x + threadIdx.x;
const int i = blockIdx.y * blockDim.y + threadIdx.y;

此时，同一个warp里j递增而i不变，程序会反复读取$\mathbf{A}$的第i行和$\mathbf{B}$的不同列。读取$\mathbf{A}$的行会触发within-warp broadcast，而对于$\mathbf{B}$而言，同一时刻每个thread读取的元素在内存中是连续的，满足memory coalescing的要求。

__global__ void sgemm_coalesce(
    int M, int N, int K, const float alpha, const float *A, const float *B, const float beta, float *C)
{
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    const int i = blockIdx.y * blockDim.y + threadIdx.y;

    // avoids memory access error if threads are more than elements
    if (i < M && j < N)
    {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k)
        {
            sum += A[i * K + k] * B[k * N + j];
        }
        C[i * N + j] = alpha * sum + beta * C[i * N + j];
    }
}

sgemm_coalesce只需要392ms。

Tiled Matrix Multiply with Shared Memory

sgemm_coalesce的主要开销还是由从全局内存读取数据的次数决定。一种改进思路是用shared memory实现tiled matrix multiply。tiled matrix multiply通过把数据分块，一次性读取分块数据到更快的内存上（认为更快的内存访问时间可以忽略不计），从而减小整体全局内存访问的开销。假设缓存的大小是B x B，则总显存访问次数为$MN(2K/B+1)$，因此B越大，访问全局内存的时间开销就越少。

#define BLOCKSIZE 32

__global__ void sgemm_shared(
    int M, int N, int K, const float alpha, const float *A, const float *B, const float beta, float *C)
{
    __shared__ float As[BLOCKSIZE * BLOCKSIZE];
    __shared__ float Bs[BLOCKSIZE * BLOCKSIZE];

    A += blockIdx.y * BLOCKSIZE * K;
    B += blockIdx.x * BLOCKSIZE;
    C += blockIdx.y * BLOCKSIZE * N + blockIdx.x * BLOCKSIZE;

    const int j = blockIdx.x * BLOCKSIZE + threadIdx.x;
    const int i = blockIdx.y * BLOCKSIZE + threadIdx.y;

    // avoids memory access error if threads are more than elements
    if (i < M && j < N)
    {
        float fSum = 0.0f; // stores result of (threadIdx.y, threadIdx.x) on each block
        for (int iBlkIdx = 0; iBlkIdx < K; iBlkIdx += BLOCKSIZE)
        {
            if (iBlkIdx + threadIdx.x < K)
            {
                As[threadIdx.y * BLOCKSIZE + threadIdx.x] = A[threadIdx.y * K + threadIdx.x];
            }
            if (iBlkIdx + threadIdx.y < K)
            {
                Bs[threadIdx.y * BLOCKSIZE + threadIdx.x] = B[threadIdx.y * N + threadIdx.x];
            }
            __syncthreads(); // syncronize until  all caches are fulfilled

            // updates to the next chunk
            A += BLOCKSIZE;
            B += BLOCKSIZE * N;

            // dot product on caches
            for (int iInnerLoop = 0; iInnerLoop < BLOCKSIZE; ++iInnerLoop)
            {
                if (iBlkIdx + iInnerLoop < K)
                {
                    fSum += As[threadIdx.y * BLOCKSIZE + iInnerLoop] * Bs[iInnerLoop * BLOCKSIZE + threadIdx.x];
                }
            }

            __syncthreads();
        }

        C[threadIdx.y * N + threadIdx.x] = alpha * fSum + beta * C[threadIdx.y * N + threadIdx.x];
    }
}

上面的核函数各为$\mathbf{A}$和$\mathbf{B}$分配了32 x 32的shared memory，同一个block中的threads都可以访问。这里我们假设block的大小就是(32, 32, 1)，因此block中的每一个thread可以通过(threadIdx.x, threadIdx.y)对应shared memory中的一个元素，也对应$\mathbf{C}$中的一个输出元素。

block中的thread有三种工作：1、读取$\mathbf{A}$和$\mathbf{B}$中的对应块位置元素到shared memory，等待直到所有threads完成元素读取；2、block中的每个thread使用shared memory执行分块矩阵乘法，并累加计算的结果；3、移动分块在$\mathbf{A}$和$\mathbf{B}$中的位置，重复1和2的工作，直到block所对应的$\mathbf{A}$的行和$\mathbf{B}$的列都遍历结束，最终将值写入$\mathbf{C}$中。

此外，也需要考虑矩阵大小不能被block大小整除的情况，通过边界条件来控制thread的计算。sgemm_shared的执行时间为147ms。