SGEMM CUDA 算子初探

SGEMM_CUDA 是一个 SGEMM 算子从 0 到 1 优化的教程，配套文字版教程在「How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog」。这里会按照一个新手的视角，结合上面的教程重新写一篇入门教程。

SGEMM 的全称是 Single precision GEneral Matrix Multiply，计算 $C = \alpha AB + \beta C$，精度是 fp32。NVIDIA cuBLAS (Basic Linear Algebra Subprograms) 提供了极致的优化算法，SGEMM_CUDA 的目标是一步步优化实现其 95% 的性能。

如果你不熟悉 GPU 架构和 CUDA 编程，强烈建议阅读「CUDA Programming Guide」。

理论极限

计算 $C = AB + C$，其中矩阵 A、B 和 C 的尺寸是 $4092 \times 4092$。

计算

总计算量 FLOPs (Floating Point Opertions) 可以被拆解为 $AB$ 计算量加 $+C$ 计算量。

P.S. 英文原文中错误使用了 FLOPs 和 FLOPS，前者表示的是总计算量，比如一个矩阵运算需要多少次运算，后者（FLoating-point OPerations Per Second）表示的是计算速度，比如 A100 显卡的理论速度为 312 TFLOPS。

$AB$ 矩阵乘，计算每个元素需要用到矩阵 A 的一行，和矩阵 B 的一列，每个元素相加再相乘，一个结果需要的计算量是 $2 \times 4096$。一共 $4092^2$ 个元素需要相乘，最终的计算量是 $2 \times 4096^3$。

$X + C$ 每个元素计算 1 次加法，总计算量为 $4092^2$。

总计算量是 $2 \times 4092^3 + 4092^2 = 137,053,437,840 \approx 137GFLOPs$。

我用的实验卡是 NVIDIA T4，FP32 的理论算力是 8.141 TFLOPS，理论极限运算速度是 137 GFLOPs / 8.141 TFLOPS 约等于 16.828ms。

存储

加载矩阵 A、B 和 C 共需要 $3 \times 4092^2 \times 4B \approx 201MB$，写入矩阵 C 共需要 $4092^2 \times 4B \approx 67MB$。

总存储量是 268 MB，T4 理论带宽速度 320 GB/s，极限存储传输时间为 0.838ms。

Compute-bound

计算时间是传输时间的 20x，可以认为这是一个计算密集型任务（compute-bound）。

Kernel 1: Naive

实现

Kernel

__global__ void sgemm_naive(int M, int N, int K, float alpha, const float *A,
                            const float *B, float beta, float *C) {
  // compute position in C that this thread is responsible for
  const uint x = blockIdx.x * blockDim.x + threadIdx.x;
  const uint y = blockIdx.y * blockDim.y + threadIdx.y;

  // `if` condition is necessary for when M or N aren't multiples of 32.
  if (x < M && y < N) {
    float tmp = 0.0;
    for (int i = 0; i < K; ++i) {
      tmp += A[x * K + i] * B[i * N + y];
    }
    // C = α*(A@B)+β*C
    C[x * N + y] = alpha * tmp + beta * C[x * N + y];
  }
}

启动 kernel

dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32), 1);
dim3 blockDim(32, 32, 1);
sgemm_naive<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);

CUDA 代码显示的是一个 thread 的视角。

这里 blockIdx.x * blockDim.x 找到当前 thread 的 block 在 x 轴的起点，blockIdx.y * blockDim.y 找到 block 在 y 轴的起点。

找到 block 的第一个元素的坐标后，就可以通过 threadIdx.x 和 threadIdx.y 定位到矩阵 C 的坐标。Kernel 1 的每个 thread 计算一个矩阵 C 的元素。

2D 数组在这篇教程中会被展开为 1D 数组，左边是逻辑上的样子（2D），右边是物理上的（1D）。

对于一个线程来说，x 是固定的，矩阵 A 的 shape 是 (M, K)，所以 x * K 就固定在某一行，x * K + i 遍历一行的全部数据。同理，+ y 是固定的，可以理解为列是固定的，每次加 i 个 矩阵 B 的行长（N），i * N + y 表示遍历一列的全部数据。

实验

T4 测试结果如下表，理论算力 8.141 TFLOPS，性能达成率 ~1.95%。

Kernel 2: Global Memory Coalescing

Warp 和 Warp Scheduler

NVIDIA GPU 定义相邻的 32 个 threads 为一组 warp，它们按照 SIMD 方式运作，即同一个指令，不同的数据。

每个 SM 有 4 个 warp schedulers，warp redisent 则更多，同一时间最多有 4 个 warps 并行执行。Warp 的状态分为 waiting 和 ready，前者可能在等 global memory 等。只有处在 ready 状态的 warp 可以被执行。

相邻的 32 个 threadId 被认为是一个 warp，threadId 的计算方式是 threadId = threadId.x + blockDim.x * threadId.y + blockDim.y * blockDim.x * threadId.z。

其中 blockDim.y * blockDim.x 表示一个 z 轴长度，* threadId.z 表示跨越了几个 z 轴长度（绿色部分），blockDim.x * threadId.y 表示跨越了几个 y 轴长度（蓝色），最后红色表示 x 内的长度。

Warp ID 的计算方法是 warp = threadId / 32。

Coalescing

当一个 warp 内的 threads 读取一块连续的空间，访存可以被组合为一次 LOAD，这称之为 coalescing。GPU 通常支持 32B、64B 和 128B 内存访问。比如，连续读取 32bits float 时，warp scheduler 可以用一次 LOAD 读取 32 个 floats（32 * 4B = 128B）。

• Coalescing 必须要内存对齐。
• Warp 内的读取顺序无要求，整体是连续的即可。

改进思路

这张图解释了 Naive Kernel 性能不佳的原因。

假设 block(0,1) 计算矩阵 C 的右上角，第一个 warp（warp0）负责计算矩阵 C 的一列。右半部分以时间维度展示了 warp0 的线程在同一时间加载矩阵 A 的不同列，time0 t0、t1 分别读取 A[i][0]、A[i+1][0]，无法通过 coalescing 连续加载 floats。

根因是 warp 的 threads 计算矩阵 C 不同行元素，改进思路是使其计算矩阵 C 的同一行。

实现

Kernel 实现

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

if (x < M && y < N) {
  float tmp = 0.0;
  for (int i = 0; i < K; ++i) {
    tmp += A[x * K + i] * B[i * N + y];
  }
  C[x * N + y] = alpha * tmp + beta * C[x * N + y];
}

启动 kernel，blockDim 变为 1D，共计 1024 个 threads。

dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32));
dim3 blockDim(32 * 32);
sgemm_coalescing<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);

P.S. 图中 time0 和 item1 应该被修正为 iter0 和 iter1。

Block 从 2D 被摊平为了 1D，thread 数量依然保持 1024。核心是 warp 在加载矩阵 B 时，是按行顺序加载，warp scheduler 可以将几个小的 LOAD 请求合并为一个大的请求。另外，矩阵 C 的写入也是顺序的，也可以吃到 coalescing 的红利。

同样的，右半边展示了 iter0 和 iter1 时刻 warp 的工作细节：

iter0

iter1

实验

4092 size 可以跑出 513.3 GFLOPS 的性能，性能达成率 ~6.31%，比 Naive Kernel 快了 ~3.24x。

Kernel 3: Shared Memory Cache-Blocking

Shared Memory (SMEM)

每个 SM 物理上有一个 SMEM。每个 block 享有一个 SMEM 的分块（chunk），block 内的线程都可以访问这块区域。速度上 shared memory 远远快于 global memory。比如 A6000 GPU，每个 block 允许访问最多 48KB SMEM。