CUDA 算子优化：GEMM

Note

高性能算子系列博客

GEMM（General Matrix Multiplication）是深度学习中最核心的计算模式之一。 Transformer 中的 Linear、Attention 投影以及 MLP 层，本质上都可以归结为矩阵乘法。因此，GPU 对 GEMM 的优化能力直接决定了 AI 训练和推理性能。

本文将从一个最简单的 GEMM kernel 出发，逐步介绍 GPU GEMM 优化中的几个核心思想：

如何将矩阵计算映射到 GPU 线程；
如何利用 GPU memory hierarchy 减少数据访问；
如何提高单个线程的计算密度。

矩阵乘法#

对于矩阵乘法

C = A \times B

其中：

A 的形状为 $M \times K$
B 的形状为 $K \times N$
C 的形状为 $M \times N$

输出矩阵 C 中的每个元素通过下面的公式计算得到：

C_{ij} = \sum_{k = 0}^{K - 1}A_{ik}B_{kj}

也就是说，每个输出元素都需要计算 A 的一行和 B 的一列之间的点积。

对于整个矩阵：

C\in \mathbb{R}^{M\times N}

总计算量约为：

2MNK

而输入矩阵大小为：

MK+KN

因此 GEMM 的核心特点是：

Important

计算量大，并且存在大量数据复用。

如何充分利用这种数据复用，是 GPU GEMM 优化的关键。

Naive GEMM：让 GPU 并行计算#

基于 GEMM 的计算过程，我们可以很容易写出下面的 kernel 代码，每一个线程去处理一个输出元素，读取对应的行和列，进行点积运算，然后将结果写回 memory 中。

例如：

1
C[row][col]

对应一个 thread。

该 thread 需要读取：

A 的第 row 行
B 的第 col 列

然后完成点积运算。

1
__global__ void gemm_naive (const float* A, const float* B, float* C, int M, int N, int K)
2
{
3
    int col = blockDim.x * blockIdx.x + threadIdx.x;
4
    int row = blockDim.y * blockIdx.y + threadIdx.y;
5
    if (row >= M || col >= N) return;
6

7
    float accum = 0.0f;
8
    for (int i = 0; i < K; i++)
9
    {
10
        accum += A[row * K + i] * B[i * N + col];
11
    }
12
    C[row * N + col] = accum;
13
}

这种实现能够利用 GPU 的大量线程并行计算，但是存在一个严重的问题：线程之间没有共享数据。

例如计算 $C_{0,0}, C_{0,1}, C_{0,2}$ ，三个线程都需要读取

A_{0,0}, A_{0,1}, ..., A_{0,K-1}

也就是说，同一行 A 数据会被多个线程重复从 Global Memory 中读取。而 Global Memory 是 GPU 中访问速度最慢的存储层级之一。因此，Naive GEMM 的瓶颈并不是计算能力不足，而是：

Important

没有利用线程之间的数据复用。

这引出了 GEMM 优化的第一个方向：利用 Shared Memory 缓存公共数据。

Shared Memory Tiling：减少 Global Memory 访问#

GPU 中存在多级 memory hierarchy：

1
Register
2
    |
3
Shared Memory
4
    |
5
Global Memory

其中：

Register 速度最快，但是空间有限；
Shared Memory 位于线程块内部，可以被 block 中所有线程共享；
Global Memory 容量最大，但是访问延迟较高。

因此，高性能 CUDA kernel 通常会尽量：

Important

将频繁访问的数据从 Global Memory 搬运到 Shared Memory。

在 naive GEMM 中，一个 thread block 已经覆盖了 C 矩阵中的一个区域，但是 block 内线程之间彼此独立，每个线程都会从 Global Memory 中读取自己需要的数据。为了减少重复访问，我们让 block 内线程协同加载 A 和 B 的 tile 到 Shared Memory，并利用这些共享数据完成计算。

整体的计算流程大致如下：

对应代码：

1
#define BLOCK_SIZE 32
2
__global__ void gemm_v1 (const float* A, const float* B, float* C, int M, int N, int K)
3
{
4
    int col = blockDim.x * blockIdx.x + threadIdx.x;
5
    int row = blockDim.y * blockIdx.y + threadIdx.y;
6
    if (row >= M || col >= N) return;
7

8
    int bx = blockIdx.x;
9
    int by = blockIdx.y;
10
    int tx = threadIdx.x;
11
    int ty = threadIdx.y;
12

13
    const int BM = BLOCK_SIZE;
14
    const int BN = BLOCK_SIZE;
15
    const int BK = BLOCK_SIZE;
16

17
    __shared__ float As[BM * BK];
18
    __shared__ float Bs[BK * BN];
19

20
    A = &A[(by * BM) * K];
21
    B = &B[bx * BN];
22
    C = &C[(by * BM) * N + bx * BN];
23

24
    float accum = 0.0f;
25
    for (int k = 0; k < K; k += BK)
26
    {
27
        // global ==> shared
28
        As[ty * BK + tx] = A[ty * K + tx];
29
        Bs[ty * BN + tx] = B[ty * N + tx];
30

31
        __syncthreads ();
32

33
        A = A + BK;
34
        B = B + BK * N;
35
        for (int i = 0; i < BK; i++)
36
        {
37
            accum += As[ty * BK + i] * Bs[i * BN + tx];
38
        }
39
        __syncthreads ();
40
    }
41
    C[ty * N + tx] = accum;
42
}

在 kernel 中：

1
__shared__ float As[BM * BK];
2
__shared__ float Bs[BK * BN];

定义两个 shared memory tile：

As 保留 A 的局部块；
Bs 保留 B 的局部块。

每轮循环：

1
for (int k = 0; k < K; k += BK)

完成一次：

1
Global Memory
2
        |
3
        v
4
Shared Memory
5
        |
6
        v
7
Matrix Multiply

因此，相比 naive GEMM：

Global Memory 访问次数降低；
数据复用增加；
算术强度（Arithmetic Intensity）提高。

经过 Shared Memory 优化后，Global Memory 访问已经大幅减少，但是计算过程仍然存在进一步优化空间。

在 GPU 中，数据访问速度存在明显差异：

1
Register
2
    ↑
3
Shared Memory
4
    ↑
5
Global Memory

其中 register 是线程私有的高速存储空间。CUDA 中的局部变量通常会被编译器放入寄存器中，因此如果一个线程能够同时计算多个输出元素，这些中间结果可以保存在寄存器中，减少重复的数据访问。

Thread Tile：提高线程计算密度#

Shared Memory Tiling 解决了线程块之间的数据复用问题，但是此时每个线程仍然只负责计算一个输出元素：

1
accum += As[...] * Bs[...];

一个线程完成一次矩阵乘累加后立即结束，导致线程内部计算量较低。

为了进一步提高计算资源利用率，可以让一个线程负责多个输出元素。这样，一个线程加载的数据可以参与更多计算，同时多个输出的累加结果可以保存在 register 中。

例如，一个 thread 不再计算：

1
C[0][0]

而是计算：

1
C[0][0]   C[0][1]
2
C[1][0]   C[1][1]

这就是 thread tile。

Note

不要说 Thread Tile 是为了利用寄存器，这种说法会稍微本末倒置。更准确的因果关系应该是：

1
想提高线程计算量
2
        ↓
3
一个线程计算多个输出
4
        ↓
5
需要保存多个累加结果
6
        ↓
7
这些结果放在寄存器
8
        ↓
9
获得更高计算密度

也就是说 Thread Tile 是优化目标，Register 是实现手段。

代码如下：

1
template <const int BM, const int BN, const int BK, const int TM, const int TN>
2
__global__ void sgemm(float* A, float* B, float* C, int M, int N, int K)
3
{
4
    int bx = blockIdx.x;
5
    int by = blockIdx.y;
6

7
    int block_row_thread = BN / TN; // block中一行的thread数量
8
    int block_col_thread = BM / TM; // block中一列的thread数量
9
    int thread_num = block_row_thread * block_col_thread; // block中thread总量
10

11
    int tx = (threadIdx.x % block_row_thread) * TN; // threadtile左上角x坐标
12
    int ty = (threadIdx.x / block_row_thread) * TM; // threadtile左上角y坐标
13

14
    __shared__ float As[BM * BK];
15
    __shared__ float Bs[BK * BN];
16

17
    A = &A[by * BM * K];
18
    B = &B[bx * BN];
19
    C = &C[by * BM * N + bx * BN];
20

21
    int a_tile_row = threadIdx.x / BK;
22
    int a_tile_col = threadIdx.x % BK;
23
    int a_tile_stride = thread_num / BK; // BM/(BM/(thread_num/BK)) = thread_num/BK = stride
24

25
    int b_tile_row    = threadIdx.x / BN;
26
    int b_tile_col    = threadIdx.x % BN;
27
    int b_tile_stride = thread_num / BN;
28

29
    float accum[TM][TN] = { 0.0f };
30
    for (int k = 0; k < K; k += BK)
31
    {
32
        for (int i = 0; i < BM; i += a_tile_stride)
33
        {
34
            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
35
        }
36
        for (int i = 0; i < BK; i += b_tile_stride)
37
        {
38
            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
39
        }
40
        __syncthreads();
41

42
        A += BK;
43
        B += BK * N;
44

45
        for (int row = 0; row < TM; row++)
46
        {
47
            for (int col = 0; col < TN; col++)
48
            {
49
                for (int i = 0; i < BK; i++)
50
                {
51
                    accum[row][col] += As[(ty + row) * BK + i] * Bs[i * BN + (tx + col)];
52
                }
53
            }
54
        }
55
        __syncthreads ();
56
    }
57
    for (int row = 0; row < TM; row++)
58
    {
59
        for (int col = 0; col < TN; col++)
60
        {
61
            C[(ty + row) * N + (tx + col)] = accum[row][col];
62
        }
63
    }
64
}

矩阵乘法#

Naive GEMM：让 GPU 并行计算#

Shared Memory Tiling：减少 Global Memory 访问#

Thread Tile：提高线程计算密度#

支持与分享

目录