gemm优化思想（已完结）

Polaris北极星少女

已于 2025-01-14 17:28:53 修改

阅读量883

点赞数 22

分类专栏： GPU编程 # gemm 文章标签：算法人工智能

于 2024-11-03 23:30:01 首次发布

本文链接：https://blog.csdn.net/qq_62704693/article/details/143420950

版权

GPU编程同时被 2 个专栏收录

23 篇文章 2 订阅

订阅专栏

gemm

4 篇文章 0 订阅

订阅专栏

来源参考：深入浅出GPU优化系列：GEMM优化（一） - 知乎https://zhuanlan.zhihu.com/p/435908830

深入浅出GPU优化系列：GEMM优化（二） - 知乎https://zhuanlan.zhihu.com/p/442930482 深入浅出GPU优化系列：GEMM优化（三） - 知乎https://zhuanlan.zhihu.com/p/481600052

核心思想

HPC优化的核心思想：怎么样让数据放在更近的存储上来掩盖计算的延时，从而减少存储墙的影响。

global memory层面(全局内存)

假设有矩阵A,B，需要计算矩阵A和B的乘，即矩阵C。

前提：A、B、C三个矩阵的维度分别为， $m\cdot k$ ， $k\cdot n$ ， $m\cdot n$ ，三个矩阵中的数据是单精度浮点数。

思考过程：对于C中每一个元素， $C[i][j]$ ，可以看作是A的一行和B的一列进行一次归约操作。采用最naive的GEMM算法，在GPU中，一共开启 $m\cdot n$ 个线程，即矩阵C的元素数，也就是 $m\cdot n$ 个每个线程需要读取矩阵A的一行与矩阵B的一列进行计算，而后将计算结果写回至矩阵C中。因而，完成计算一共需要从global memory（全局内存）中进行 $2mnk$ 次读操作和 $m*n$ 次写操作。

读操作数计算：读矩阵A的一行与读矩阵B的一列计算，计算出矩阵C的数据，因为矩阵A有m行，矩阵B有n列，所以是 $m\cdot n$ 次读，又因为每个线程都分别要读矩阵A与矩阵B的k个元素进行计算即K次迭代，所以每次读 $2k$ 个元素，又因为上述知道 $m\cdot n$ 次读，所以总共是 $2k*m*n$ 次读操作。

写操作数计算：因为矩阵C共有 $m*n$ 个元素，所以计算出 $m*n$ 个数据放入矩阵C就是写操作的过程，共有 $m*n$ 次写操作。

问题：大量的访存操作（减少访存）使得GEMM效率难以提高，因而考虑global memory中进行分块，并将矩阵块放置到shared memory（共享内存）中

从global memory(全局内存)到shared memory(共享内存)

共享内存：3.2.4. 共享内存-CSDN博客https://blog.csdn.net/qq_62704693/article/details/143350766?spm=1001.2014.3001.5502

对global memory进行分块的GEMM算法示意图见下图。

计算过程：在GPU中开启 $M\cdot N$ 个block(线程块)，代表被划分后的C的分块数量，每个block负责C中一个维度为 $bm\cdot bn$ 的小矩阵块或者分块的计算。计算中一共有K次迭代，每一次迭代都需要读取A中一个维度为 $bm\cdot bk$ 的小矩阵块和B中一个维度为 $bk\cdot bn$ 的小矩阵块，并将其放置在shared memory中。因而，完成C中所有元素的计算一共需要从global memory中读取 $M\cdot N\cdot K\cdot (bm\cdot bk\dotplus bk\cdot bn)$ ，即 $m\cdot n\cdot k(1/bm+1/bn)$ 个单精度浮点数。相比于naive的GEMM算法，访存量减少为原来的 $1/2\cdot (1/bm+1/bn)$ 。

读操作数计算：在GPU中开启 $M\cdot N$ 个block，每个block负责C中一个维度为 $bm\cdot bn$ 的小矩阵块的计算，又因为每个block都分别要读矩阵A与矩阵B $K$ 个矩阵块进行计算即 $K$ 次迭代，又因为每个block在矩阵A中每个小矩阵块读取 $(bm\cdot bk\dot)$ 个操作数，矩阵B中每个小矩阵块读取 $(bk\cdot bn\dot)$ 个操作数，所以每个block在矩阵A每次读取 $M\cdot N\cdot K\cdot (bm\cdot bk)$ ，每个block在矩阵B每次读取 $M\cdot N\cdot K\cdot (bk\cdot bn)$ ，两者相加为总读操作数， $M\cdot N\cdot K\cdot (bm\cdot bk\dotplus bk\cdot bn)$ 。

优点：通过global memory中分块算法极大地减少了对global memory的访存量。并且，相比于naive算法，对global进行分块可以更充分地利用数据局部性。在naive算法中，每一个线程都需要直接从global memory中取数，其时延非常长，计算性能非常差。而进行分块后，将维度为， $bm\cdot bk$ ， $bk\cdot bn$ 的小矩阵块先存储到shared memory之中。而后计算单元进行计算时可以直接从shared memory中取数，大大减少了访存所需要的时延。

从shared memory(共享内存)到register(寄存器)

前提：设 $b\in(0,1]$ 将A、B、C三个矩阵划分为多个维度为， $bm\cdot bk$ ， $bk\cdot bn$ ， $bm\cdot bn$ 的小矩阵块。其中设 $M=\frac{m}{bm}=\frac{1}{b}$ ， $N=\frac{n}{bn}=\frac{1}{b}$ ， $K=\frac{k}{bk}=\frac{1}{b}$ 代表每行每列的小矩阵块数。三个矩阵形成， $M\cdot K$ ， $K\cdot N$ ， $M\cdot N$ 的小矩阵网格。进一步从shared memory到register的过程。在这里，只分析一个block中的计算。当进行K轮迭代中某一轮迭代时，GPU将维度为， $bm\cdot bk$ ， $bk\cdot bn$ 的小矩阵块存储到shared memory中，而后各个线程将shared memory中的数据存入register中进行计算。

计算过程：

不对shared memory分块（未改进）：一个block中含有 $bm\cdot bn$ 个线程，每一个线程负责C中一个元素的计算。则一个block一共需要对shared memory进行 $2\cdot bm\cdot bn\cdot bk$ 次读操作。

对shared memory进行分块（改进）：对 $bm\cdot bn$ 的小矩阵进行再一次划分，将其划分为多个维度为 $rm\cdot rn$ 的子矩阵。设 $X=\frac{bm}{rm}$ ， $Y=\frac{bn}{rn}$ 。则一个block需要负责 $X\cdot Y$ 个子矩阵。随后，在一个block中开启 $X\cdot Y$ 个线程，每个线程负责一个维度为 $rm\cdot rn$ 的子矩阵的计算。在计算中，一个block一共需要从shared memory读取 $X\cdot Y\cdot (rm+rn)\cdot bk$ ，即 $bm\cdot bn\cdot bk\cdot (\frac{1}{rm}+\frac{1}{rn})$ 个单精度浮点数。相比于未分块的算法，对于shared memory中的访存量减少为原来的 $\frac{1}{2}\cdot (\frac{1}{rm}+\frac{1}{rn})$ 。并且，由于将数据放入register中，可以直接对数据进行运算，减少了从shared memory中取数的时延。

读操作数计算：在一个block中开启 $X\cdot Y$ 个线程，一个block负责C中一个维度为 $bm\cdot bn$ 的小矩阵块的计算，又因为一个block需要负责 $X\cdot Y$ 个子矩阵，又因为每个线程负责一个维度为 $rm\cdot rn$ 的子矩阵的计算，在共享内存A矩阵块中读取 $rm$ 行操作数，在共享内存B矩阵块中读取 $rn$ 列操作数，迭代 $bk$ 次，所以一个block一共需要从shared memory读取 $X\cdot Y\cdot (rm+rn)\cdot bk$ ，即 $bm\cdot bn\cdot bk\cdot (\frac{1}{rm}+\frac{1}{rn})$ 个单精度浮点数。

优点：相比于未分块的算法，对于shared memory中的访存量减少为原来的 $\frac{1}{2}\cdot (\frac{1}{rm}+\frac{1}{rn})$ 。并且，由于将数据放入register中，可以直接对数据进行运算，减少了从shared memory中取数的时延。

register分块

寄存器分块的背景

在NVIDIA的GPU架构中，寄存器文件（Register File）被分为多个bank。例如，在Maxwell架构中，寄存器文件被分为4个bank，寄存器ID的低2位（即寄存器ID % 4）决定了它属于哪个bank。如果一条指令的源寄存器中有两个或更多来自同一bank，就会发生bank冲突，导致指令重发射，浪费计算周期。

前提

考虑最后一步，即register中的计算，并且只分析一个thread。在完成以上的过程后，对于一个线程而言，它现在拥有： $rm$ 个A矩阵的寄存器值， $rn$ 个B矩阵的寄存器值，以及 $rm\cdot rn$ 个C矩阵的寄存器值。通过这些寄存器的值，需要计算 $rm\cdot rn$ 个数。这需要 $rm\cdot rn$ 条FFMA指令。对于maxwell架构的GPU而言，bank数为4，寄存器id%4即所属bank。

原理

这个时候会涉及到寄存器的bank conflict。在NV的GPU中，每个SM不仅会产生shared memroy之间的bank 冲突，也会产生寄存器之间的bank冲突。这一点对于计算密集型的算子十分重要。像shared memory一样，寄存器的Register File也会被分为几个bank，如果一条指令的的源寄存器有2个以上来自同一bank，就会产生冲突。指令会重发射，浪费一个cycle。

数据的prefetch

最后，我们来讲讲如何通过对数据进行prefetch来减少访存的latency。我们再来回顾GEMM的过程，并且仔细地看看这个访存的latency到底是怎么导致的。对于一个block而言，需要计算一个 $bm\cdot bn$ 的矩阵块，这个时候需要进行K次迭代，每次迭代都需要先将来自A和B的两个小块送到shared memory中再进行计算。而从global中访存实际上是非常慢的，所以导致了latency。虽然GPU中可以通过block的切换来掩盖这种latency，但是由于分配的shared memory比较多，活跃的block并不太多，这种延时很难被掩盖。对于一个thread，需要计算一个 $rm\cdot rn$ 的小矩阵，但是必须先将数据从shared memory传到寄存器上，才能开始进行计算。所以导致了每进行一次迭代，计算单元就需要停下来等待，计算单元不能被喂饱。

为此，需要进行数据的Prefetch来尽可能地掩盖这种latency。思想也比较简单，需要多开一个buffer，进行读写分离。示意图如下。当block进行第2轮迭代时，需要对A2和B2进行计算，在计算单元进行计算的同时，我们将A3和B3提前放置到shared memory。而后，在进行第3轮迭代时，就可以直接对shared memory中的A3和B3进行计算，而不需要等待从global memory搬运到shared memory的时间。寄存器上的Prefetch也是同理。

背景

数据的Prefetch（预取）是一种常用的优化技术，旨在通过提前将数据从慢速存储（如全局内存）加载到快速存储（如共享内存或寄存器）中，从而掩盖内存访问的延迟。在GPU编程中，尤其是在计算密集型任务如GEMM（通用矩阵乘法）中，Prefetch可以显著提高计算单元的利用率，减少等待时间。

GEMM中的访存延迟问题

在GEMM的计算过程中，一个block需要计算一个大小为 $bm \cdot bn$ 的矩阵块，这通常需要 $K$ 次迭代。每次迭代都需要从全局内存（global memory）中加载A和B矩阵的小块到共享内存（shared memory），然后再进行计算。由于全局内存的访问速度较慢，这种加载操作会引入较大的延迟。

虽然GPU可以通过切换不同的block来掩盖部分延迟，但由于每个block需要分配较多的共享内存，活跃的block数量有限，因此延迟很难被完全掩盖。此外，对于每个线程而言，它需要从共享内存中加载数据到寄存器才能进行计算，这也会导致计算单元在每次迭代中停下来等待数据加载，从而无法充分利用计算资源。

Prefetch的基本思想

Prefetch的核心思想是通过读写分离和多开buffer来提前加载数据，从而掩盖内存访问的延迟。具体来说，Prefetch可以分为两个层次：
1. 共享内存的Prefetch：在计算当前迭代的同时，提前将下一轮迭代需要的数据从全局内存加载到共享内存。
2. 寄存器的Prefetch：在计算当前迭代的同时，提前将下一轮迭代需要的数据从共享内存加载到寄存器。

通过这种方式，计算单元可以在当前迭代中计算数据的同时，提前准备好下一轮迭代所需的数据，从而减少等待时间。

共享内存的Prefetch

1. 第1轮迭代：

从全局内存加载A1和B1到共享内存。
计算单元开始计算A1和B1。

2. 第2轮迭代：

在计算A1和B1的同时，从全局内存加载A2和B2到共享内存的另一个buffer中。
计算单元完成A1和B1的计算后，直接开始计算A2和B2。

3. 第3轮迭代：

在计算A2和B2的同时，从全局内存加载A3和B3到共享内存的第一个buffer中。
计算单元完成A2和B2的计算后，直接开始计算A3和B3。

通过这种方式，每次迭代的计算和数据的加载可以重叠，从而掩盖全局内存访问的延迟。

寄存器的Prefetch

寄存器的Prefetch与共享内存的Prefetch类似，但发生在更细粒度的层次上。以下是一个寄存器Prefetch的流程：

1. 第1轮迭代：

从共享内存加载A1和B1的数据到寄存器。
计算单元开始计算A1和B1。

2. 第2轮迭代：

在计算A1和B1的同时，从共享内存加载A2和B2的数据到另一组寄存器。
计算单元完成A1和B1的计算后，直接开始计算A2和B2。

3. 第3轮迭代：

在计算A2和B2的同时，从共享内存加载A3和B3的数据到第一组寄存器。
计算单元完成A2和B2的计算后，直接开始计算A3和B3。

通过这种方式，寄存器的加载和计算可以重叠，从而减少计算单元的等待时间。

Prefetch的实现要点

1. 双缓冲（Double Buffering）：

使用两个buffer来实现读写分离。一个buffer用于当前迭代的计算，另一个buffer用于下一轮迭代的数据加载。
在每次迭代结束时，交换两个buffer的角色。

2. 同步机制：

在Prefetch过程中，需要确保数据加载和计算之间的正确同步。通常使用`__syncthreads()`来同步线程块内的线程。

3. 资源分配：

Prefetch需要额外的存储空间（如共享内存和寄存器），因此需要合理分配资源，避免资源不足。

示例代码

以下是一个简单的共享内存Prefetch的伪代码示例：

__global__ void gemm_kernel(float *A, float *B, float *C, int M, int N, int K) {
    extern __shared__ float shared_mem[];
    float *A_shared = shared_mem;
    float *B_shared = shared_mem + blockDim.x * blockDim.y;

    int tid = threadIdx.x + threadIdx.y * blockDim.x;
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;

    float C_value = 0.0f;

    for (int k = 0; k < K; k += blockDim.x) {
        // Prefetch next tile
        if (k + blockDim.x < K) {
            A_shared[tid] = A[row * K + (k + blockDim.x)];
            B_shared[tid] = B[(k + blockDim.x) * N + col];
        }
        __syncthreads();

        // Compute current tile
        for (int i = 0; i < blockDim.x; ++i) {
            C_value += A_shared[threadIdx.y * blockDim.x + i] * B_shared[i * blockDim.y + threadIdx.x];
        }
        __syncthreads();
    }

    C[row * N + col] = C_value;
}