如何优化一个 CUDA Matmul Kernel 以达到类 cuBLAS 的 Performance：一份 Worklog

搬运自：siboehm.com/articles/22/CUDA-MMM

2022年12月

在这篇文章中，我将迭代地优化一个用 CUDA 编写的 matrix multiplication 实现。
我的目标不是构建一个 cuBLAS 的替代品，而是深入理解用于现代 deep learning 的 GPUs 最重要的 performance 特性。
这包括 coalescing global memory 访问、shared memory caching 和 occupancy 优化等。

你可以从 Github 下载所有 kernels 的代码。也可以查看 wangzyon 的 repo，我从那里复制了 benchmarking 设置。

这篇文章比我通常上传的内容少了一些润色，包含了更多的 sidenotes。我在编写 kernels 时将其用作 ideas 和草图的笔记本。这就是为什么我称之为 worklog :)

GPU 上的 Matrix multiplication 目前可能是最重要的 algorithm，考虑到它几乎占据了 large deep-learning models 在 training 和 inference 过程中的所有 FLOPs。
那么，从头开始编写一个 performant 的 CUDA SGEMM 需要多少工作量？

SGEMM 在 single (=32b) precision 下执行 C=αAB+βC。

我将从一个 naive kernel 开始，逐步应用 optimizations，直到我们的 performance 达到 cuBLAS（NVIDIA 官方的 matrix 库）的 95%（在状态好的时候）：

这里指的是 FP32 下的 cuBLAS。在我的设置中，使用 TF32 或 BF16 precision 进行 matmul 允许 cuBLAS 使用 tensor cores，这将增加 2.5 倍或 3.5 倍的 FLOPS。我可能会在未来的文章中研究 tensor cores / warp matrix 函数。

Kernel 1: Naive 实现

在 CUDA 编程模型中，computation 被组织在一个三层的 hierarchy 中。
每次调用 CUDA kernel 都会创建一个新的 grid，它由多个 blocks 组成。
每个 block 由多达 1024 个单独的 threads 组成。

这些 constants 可以在 CUDA Programming guide 中查找。

处于同一个 block 中的 threads 可以访问同一个 shared memory 区域 (SMEM)。

block 中的 threads 数量可以使用通常称为 blockDim 的变量进行配置，它是一个由三个 ints 组成的 vector。
该 vector 的条目指定了 blockDim.x、blockDim.y 和 blockDim.z 的大小，如下所示：

类似地，grid 中的 blocks 数量可以使用 gridDim 变量进行配置。
当我们从 host 启动一个新 kernel 时

在 accelerator 术语中，*host* 指的是 CPU，*device* 是加速器，这里是 GPU。

，它会创建一个单个 grid，包含指定的 blocks 和 threads。

从这里开始，我将只讨论 2D grids 和 blocks，部分原因是 3D-structure 很少使用，而且绘制 3D 太困难了。

重要的是要记住，我们刚才谈到的 thread hierarchy 主要关乎程序 correctness。
对于程序 performance，正如我们稍后将看到的，将同一个 block 中的所有 threads 视为平等并非一个好主意。

对于我们的第一个 kernel，我们将使用 grid、block 和 thread hierarchy 为每个 thread 分配 result matrix C 中的一个唯一条目。
然后该 thread 将计算 A 的相应 row 和 B 的相应 column 的 dot product，并将结果写入 C。
由于 C 的每个位置仅由一个 thread 写入，我们无需进行 synchronization。
我们将这样启动 kernel：

// 创建足够多的 blocks 以映射 C 的所有条目
dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32), 1);
// 32 * 32 = 1024 thread per block
dim3 blockDim(32, 32, 1);
// 在 device 上启动 kernel 的异步执行
// 函数调用在 host 上立即返回
sgemm_naive<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);

CUDA 代码是从单个 thread 视角编写的。
在 kernel 代码中，我们访问 blockIdx 和 threadIdx 内置变量。
这些变量会根据访问它们的 thread 返回不同的值。

在我们的示例中，threadIdx.x 和 threadIdx.y 将根据 thread 在 grid 中的位置在 0 到 31 之间变化。blockIdx.x 和 blockIdx.y 也是如此，它们将根据 thread 的 block 在 grid 中的位置在 0 到 CEIL_DIV(N, 32) 或 CEIL_DIV(M, 32) 之间变化。

我们将进行大量针对 matrix 在内存中 strided 表示的 indexing 操作。Edward Yang 关于 PyTorch Internals 的文章对 strided tensors 有很好的解释。

__global__ void sgemm_naive(int M, int N, int K, float alpha, const float *A,
                            const float *B, float beta, float *C) {
  // 计算该 thread 负责的 C 中的位置
  const uint x = blockIdx.x * blockDim.x + threadIdx.x;
  const uint y = blockIdx.y * blockDim.y + threadIdx.y;

  // 当 M 或 N 不是 32 的倍数时，需要 `if` 条件。
  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：

如果 matrix 的大小不能被 block 的大小整除，我们将不得不启动额外的 blocks 来处理 remainder。例如，在下图中，我们将创建 9 个具有相同 threadsize 的 blocks，但其中只有 4 个充分利用了它们的 1024 个 threads。这种现象被称为 tile quantization，每当我们要将固定大小的 volume 映射到可变大小的 input 时就会出现。

这个 kernel 在我的 A6000 GPU 上处理三个 4092² fp32 matrices 需要大约 0.5 秒。
让我们进行一些非实现特定的 calculations：

下限估算最快可能的 Runtime

对于两个 4092² matrices 的 matrix multiplication，接着是一个 4092² matrix 的加法（以构成 GEMM）：

1. 总 FLOPS：

   对于 C 的 4092² 个条目中的每一个，我们必须执行两个大小为 4092 的 vectors 的 dot product，每一步涉及一次 multiply 和一次 add。 “Multiply then add” 通常被映射到称为 FMA (fused multiply-add) 的单个 assembly 指令，但仍计为两个 FLOPs。

   2*4092³ + 4092² = 137 GFLOPS

2. 总读取数据（最小！）：3 * 4092² * 4B = 201MB

3. 总存储数据：4092² * 4B = 67MB

因此 268MB 是任何实现在 global GPU memory 之间传输的绝对最小 memory 量，

Global memory 是 GPU 的主内存区域。如果 Nvidia 向你推销一款标榜 80GB 内存和 1TB/s 带宽的 GPU，他们谈论的是 global memory 的 capacity 和 bandwidth。稍后我们将讨论 GPU 上的其他 memory 区域，例如 shared memory，它在物理上是独立的，并且具有非常不同的 performance 特性。

假设它有足够大的 cache。

cuBLAS kernel 在整个计算过程中总共加载了 500MB 的 GMEM。稍后我们将看到增加 arithmetic intensity 如何允许我们实现如此低的访问量。

让我们计算一些 kernel performance 的上限。
该 GPU 标称具有 30TFLOPs/s 的 fp32 compute throughput 和 768GB/s 的 global memory bandwidth。
如果我们达到这些数值，

提醒一下，peak FLOPs 是一个还原主义指标，因为它取决于指令组合 (instruction mix)。如果你选择的 FLOP 是 DIV，你不可能达到 30TFLOPs/s。然而，由于 matmul 主要使用 FMA 指令（这往往是最快的 FLOPs），我们很有可能真正接近那个 peak FLOP 值。

带宽的情况类似：只有当访问模式适合 hardware 时，才能达到 Peak bandwidth。

我们需要 4.5ms 进行计算，0.34ms 进行 memory transfers。
所以在我们的估算中，计算花费的时间大约是 memory accesses 的 10 倍。
这意味着我们最终优化的 kernel 将是 compute-bound 的，只要我们最终传输的 memory volume 小于 278MB 绝对最小值的 10 倍。

A6000 标称具有 309TFLOPs/s 的 tensor core 性能。如果我们可以将 tensor cores 用于我们的 fp32 matmul，计算将仅需 0.44ms，而进行 4092^2 matrix mul 的优化 kernel 几乎肯定仍将是 memory bound 的。这展示了 tensor cores 到底有多快。

现在我们已经计算了 fp32 GEMM 计算的一些下限，让我们回到手头的 kernel，找出为什么它比可能的慢得多。

Naive Kernel 的 Memory Access Pattern

在我们的 kernel 中，同一个 block 中具有 ThreadIds (0, 0) 和 (0, 1) 的两个 threads 将加载 B 的同一 column，但加载 A 的不同 rows。
如果我们假设 zero caching 的最坏情况，那么每个 thread 必须从 global memory 加载 2*4092+1 个 floats。
由于我们总共有 4092² 个 threads，这将导致 548GB 的 memory traffic。

下面是 naive kernel memory access pattern 的可视化，以两个 threads A（红色）和 B（绿色）为例：

总而言之，当我在 A6000 GPU 上运行此 kernel 时，在乘以两个 4092x4092 float32 matrices 时达到约 300GFLOPs。
相当糟糕，考虑到 A6000 被标榜能够达到接近 30 TFLOPs。

仅供对比，300 GFLOPs 大约也是我在早先文章中关于 CPU matmul 使用的 2015 Haswell CPU 上优化的 BLAS 库达到的性能。

那么我们该如何开始让它变得更快呢？
一种方法是优化 kernel 的 memory access pattern，使得 global memory 访问可以被 coalesced（=合并）为更少的访问。

Kernel 2: Global Memory Coalescing

在深入探讨 global memory coalescing 之前，我们需要了解 warp 的概念。
为了执行，一个 block 的 threads 被分组为所谓的 warps，每个 warp 由 32 个 threads 组成。
然后将一个 warp 分配给一个 warp scheduler，这是执行 instructions 的物理核心。

在 Volta 架构之前，情况曾经是 warp 的所有 threads 都由同一个指令流 (instruction stream) 供给。在一个分支上，不走该分支的 threads 使用所谓的 active mask 被禁用。然而，自 Volta 以来，依靠这种 “warp-synchronous” 行为不再是一个好主意，因为即使对于 warp 内的相同 threads，来自不同分支的指令也可能交错执行。

每个 multiprocessor 有四个 warp schedulers。
分组为 warps 是基于连续的 threadId。
如果我们设置 blockDim 为多维，则 threadId 计算如下：

threadId = threadIdx.x+blockDim.x*(threadIdx.y+blockDim.y*threadIdx.z)

然后，具有相邻 threadId 的 threads 成为同一个 warp 的一部分。
下面我尝试说明这一点，使用较小的 8 个 threads 的 “warpsize”（实际的 warps 始终包含 32 个 threads）：

我喜欢将 threadId 的三个维度 x,y,z 视为 “column-major”，因为第一个维度 x 是在 “warpspace” 中连续的那个。我不知道其他人是否使用这个词，但它使概念对我来说更清晰。

warp 的概念与第二个 kernel 相关，因为由属于同一个 warp 的 threads 进行的顺序 memory accesses 可以被分组并作为一个整体执行。
这被称为 global memory coalescing。
在优化 kernel 的 GMEM 内存访问以达到 peak bandwidth 时，这是最重要的注意事项。

下面是一个例子，其中同一个 warp 中 threads 的连续 memory accesses 被分组，允许每个 warp 仅使用 2 个 32B 的 loads 即可执行 8 次 memory accesses：

实际上，GPU 支持 32B、64B 和 128B 的 memory accesses。
因此，如果每个 thread 从 global memory 加载一个 32bit float，warp scheduler（可能是 MIO）可以将这 32*4B=128B 的 load 合并为一个 transaction。
这只有在加载的 floats 在 memory 中是连续的且访问是对齐的情况下才可能实现。

在这种方式下，GPU 上优化 global memory coalescing 与 CPU 上优化 cache line 利用率有很多相似之处。有趣的是，为了允许合并，warp 内的 threads 必须访问连续地址，但访问在 warp 内不必是顺序的。如下图所示：

如果不是，或者由于某些其他原因无法被 coalesced，那么 GPU 将根据需要执行尽可能多的 32B loads 以获取所有 floats，从而导致大量带宽浪费。
Profiling 我们的 naive kernel，我们可以观察到非合并访问的有害影响，因为我们仅达到了 15GB/s 的 GMEM throughput。

回顾上一个 kernel，我们这样为 threads 分配 C 的条目：

const uint x = blockIdx.x * blockDim.x + threadIdx.x;
const uint y = blockIdx.y * blockDim.y + threadIdx.y;

因此，同一个 warp 的 threads（具有连续 threadIdx.x 的那些）从 memory 中非连续地加载 A 的 rows。
Naive kernel 访问 A 的内存模式看起来更像这样：

为了启用 coalescing，我们可以改变我们为 threads 分配结果 matrix C 位置的方式。
这种 global memory access pattern 的变化如下图所示：

要实现这一点，我们只需要更改前两行：

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];
}

我们这样调用它：

这一点对我来说并不是显而易见的，但启用 GMEM coalescing 在 assembly 中没有任何改变，请参阅 Godbolt 上的 SASS 输出。

Global memory coalescing 将 memory throughput 从 15GB/s 提高到 110GB/s。
Performance 达到 2000 GFLOPS，与第一个 naive kernel 的 300 GFLOPS 相比有了巨大提升。
对于下一个 kernel，我们将使用 GPU 的快速 on-chip memory（称为 shared memory）来 cache 将被重复使用的数据。

Kernel 3: Shared Memory Cache-Blocking

除了巨大的 global memory 之外，GPU 还有一个小得多的物理上位于芯片上的 memory 区域，称为 shared memory (SMEM)。
在物理上，每个 SM 有一个 shared memory。

这是 A100 GPU memory hierarchy 的一个有用示意图（来源）：

逻辑上，这个 shared memory 在 blocks 之间进行 partitioned。
这意味着一个 thread 可以通过 shared memory 块与其 block 中的其他 threads 进行通信。
在我的 A6000 GPU 上，每个 block 最多可以访问 48KB 的 shared memory。

SMEM 的量是 configurable，通过权衡较大的 shared memory 和较小的 L1 cache。具体细节请参见 compute capability documentation。此外，通过使用 dynamic shared memory，每个 thread 可以使用超过 48KB 的 SMEM。

由于 shared memory 位于 on-chip，它的 latency 比 global memory 低得多，而 bandwidth 高得多。
我没能找到 Ampere 架构的良好 benchmark 结果，但对于 Volta（2017 年发布），这篇论文中进行的 benchmarks 报告了 750GiB/s 的 global memory bandwidth 和 12,080GiB/s 的 shared memory bandwidth。

看起来这些数字自 Volta 以来变化不大。Nvidia 报告我的 A6000 (Ampere) 的最大 GMEM bandwidth 约为 750GB。

因此，对于下一个 kernel，我们将把 A 的一块和 B 的一块从 global memory 加载到 shared memory 中。
然后我们将对这两个 chunks 执行尽可能多的工作，每个 thread 仍被分配一个 C 的条目。
我们将沿着 A 的 columns 和 B 的 rows 移动这些 chunks，在 C 上执行 partial sums，直到计算出结果。

如下图所示：

代码的重要部分如下，变量名与上图对应：

通常，我没有编写支持任意大小 M、N 和 K 的代码，因为条件检查 (condition checking) 引入了大量杂乱且不太有趣。为了确保 kernel 正确工作，我通过与 cuBLAS 进行比较，使用随机数据和几种不同的 matrix 大小对其进行测试。

// 将指针移动到起始位置
A += cRow * BLOCKSIZE * K;                    // row=cRow, col=0
B += cCol * BLOCKSIZE;                        // row=0, col=cCol
C += cRow * BLOCKSIZE * N + cCol * BLOCKSIZE; // row=cRow, col=cCol

float tmp = 0.0;
// 外层循环沿着 A 的 columns 和 B 的 rows 推进，
// 直到我们在 C 中完整计算出结果。
for (int bkIdx = 0; bkIdx < K; bkIdx += BLOCKSIZE) {
  // 让每个 thread 将 A 和 B 中的一个元素
  // 从 global memory 加载到 shared memory。
  // 使 threadCol (=threadIdx.x) 成为连续索引
  // 以允许 global memory 访问合并 (coalescing)
  As[threadRow * BLOCKSIZE + threadCol] = A[threadRow * K + threadCol];
  Bs[threadRow * BLOCKSIZE + threadCol] = B[threadRow * N + threadCol];

  // 阻塞该 block 中的 threads，直到 cache 完全填充
  __syncthreads();

  // 指针移动到下一个 chunk
  A += BLOCKSIZE;
  B += BLOCKSIZE * N;

  // 在当前 cached block 上执行 dotproduct
  for (int dotIdx = 0; dotIdx < BLOCKSIZE; ++dotIdx) {
    tmp += As[threadRow * BLOCKSIZE + dotIdx] *
            Bs[dotIdx * BLOCKSIZE + threadCol];
  }
  // 结束时需要再次 sync，以避免较快的 threads 
  // 在较慢的 threads 完成之前将下一个 block 获取到 cache 中
  __syncthreads();
}
C[threadRow * N + threadCol] =
    alpha * tmp + beta * C[threadRow * N + threadCol];

此 kernel 达到了约 2200 GFLOPS，比之前的版本提高了 50%。

仅提高 50% 的部分原因是我们的前一个 kernel 已经具有相当好的 L1 cache hit rates。

我们离达到 GPU 可以提供的约 30 TFLOPs 还很远。
从下面的 roofline 图中可以明显看出这一点：

请注意我们是如何实现比 cuBLAS 更高的 memory bandwidth 的。但由于我们对从 memory 加载的每个 byte 执行的工作量少得多（=较低的 arithmetic intensity），整体 performance 较差。

在 CHUNKSIZE 为 32 时，这使用了 2*32*32*4B=8KB 的 shared memory 空间。

此信息也可以通过使用 --ptxas-options=-v 编译获得，输出为：Used 37 registers, 8192 bytes smem, 400 bytes cmem[0]。

我的 A6000 GPU 为每个 block 提供的 shared memory 空间最大为 48KB，所以我们离达到那个 limit 还很远。
这不一定是个问题，因为增加每个 block 的 shared-memory 使用量是有副作用的。
每个 multiprocessor (SM) 总共有 100KB 的 SMEM 可用。
这意味着如果我们修改我们的 kernel 以使用全部 48KB 的可用 SMEM，每个 SM 同时只能保持加载两个 blocks。
在 CUDA 术语中，增加每个 block 的 SMEM 利用率会降低 occupancy。
Occupancy 定义为每个 SM 的 active warps 数量与每个 SM 最大可能的 active warps 数量之比。

高 occupancy 是有用的，因为它允许我们通过拥有更大的可发射 instructions 池来隐藏操作的高 latency。

在 GPU 上，像 FMA 这样的 math operations 具有 4 个 cycles 的 latency，在 1.5GHz clock 频率下等于 2.6ns。相比之下，在最近的 x86 CPU 上，FMA 具有 6 个 cycles 的 latency，在 3.5GHz clock 频率下为 1.8ns。

在 SM 上保持更多 active blocks 加载有三个主要 limits：register 计数、warp 计数和 SMEM capacity。
让我们为当前 kernel 做一个示例 calculation。

Kernel 3 的 Occupancy 计算

以下是从 cudaGetDeviceProperties API 获得的我的 GPU 的相关 hardware stats（Multiprocessors 是我们之前谈到的 SMs）：

Shared memory 的量是通过使用名为 SharedMemoryCarveout 的特性 configurable。所谓的 unified data cache 被 partitioned 为 L1 cache 和 shared memory，所以我们可以权衡较少的 shared-memory 以换取更多的 L1 cache。

以下是我们 kernel 的 resource demands：

工作以 block granularity 调度到 SMs 上。
只要 SM 有足够的 resources 容纳它们，它就会加载更多的 blocks。
Calculation：

我发现了很多官方和非官方的 occupancy calculators，但没有官方公式说明如何计算 occupancy。结果是正确的（我使用 NVIDIA 官方工具进行了检查），但可能在 rounding 应用等方面存在微小误差。

• Shared memory: 8192B/Block + 1024B/Block 用于 CUDA runtime 使用 = 9216B/Block。(102400B per SM) / (9216B per Block) = 11.11 ⇒ 11 Blocks 上限。
• Threads: 1024 Threads per Block, 每个 SM 最大 1536 threads ⇒ 上限 1 block。
• Registers: 37 regs per thread * 32 threads per warp = 1184 regs per warp。Register allocation granularity 在 warp 级别是 256 regs，因此向上舍入到 1280 regs per warp。我们有 (1024 threads / 32) = 32 warps per block，因此 1280 regs per warp * 32 warps per block = 40960 regs per block。每个 SM 最大 65536 regs ⇒ 上限 1 block。

  出人意料的是，register file 中的空间比 shared memory 还要多！每个 block 最多可以使用 48 KB 的 SMEM，但可以使用 65536*4B = 262 KB 的 register 空间。

所以这个 kernel 受限于每个 block 的 threads 数量以及每个 thread 的 registers 数量。
我们每个 SM 最多只能加载一个 block，从而得出最终 occupancy 为 32 active warps / 48 max active warps = 66%。

66% 的 occupancy 并不算太糟，所以这不能解释为什么我们的 kernel 运行得这么慢。

我们知道通过观察 cuBLAS 达到约 245 FLOPs/Byte，可以将我们的 kernel 优化到高 arithmetic intensity (AI)。在非常高和非常低的 AI 下，都不需要高 occupancy 来达到 peak throughput。有关这方面的更多细节，请参阅 V. Volkov 的博士论文 及其对 “cusp behaviour” 的介绍：

查看 profiler 给我们提供了一些线索。首先，如果我们查看执行 instructions 的 mix，其中大部分是 memory loads：

LDS 是 shared memory loads。FMA 是我们的 fused multiply add。IADD3 是 “3 输入整数加法”，我们需要它来沿着 K 维度移动 pointers。

我们的内层循环在 PTX 中看起来像这样（Godbolt 链接 ：

ld.shared.f32   %f91, [%r8+3456];
ld.shared.f32   %f92, [%r7+108];
fma.rn.f32      %f93, %f92, %f91, %f90;

这不太好，因为 memory load 肯定比简单的 FMA 具有更高的 latency，而且我们知道我们的 kernel 应该是 compute bound 的。
当我们查看 profiler 对 warp states 的采样时，我们可以看到这种影响。这量化了每条执行指令在每种状态下花费的 cycles 数量：

Stall Not Selected 意味着 warp 有资格被调度，但 scheduler 选择了另一个有资格的 warp。这为我们之前的假设提供了证据，即 occupancy 目前不是问题。

这些状态的含义在 Kernel Profiling Guide 中有记录。
对于 Stall MIO Throttle，它的解释是：

Warp 正处于等待 MIO (memory input/output) 指令队列不再满的状态。在 MIO pipelines 利用率极高的情况下，这种 stall 原因很高，这些情况包括特殊的 math instructions、dynamic branches 以及 shared memory 指令。

我们没有使用特殊的 math instructions，也没有使用 dynamic branches，所以很明显我们是在等待 SMEM 访问返回时发生了 stalling。
那么我们如何让我们的 kernel 发射更少的 SMEM 指令呢？
一种方法是让每个 thread 计算多个 output 元素，这允许我们在 registers 中执行更多工作，并减少对 SMEM 的依赖。

Kernel 4: 通过 1D Blocktiling 为每个 Thread 计算多个结果

接下来的这个 kernel 的工作方式与上一个类似，但增加了一个新的 inner loop，用于为每个 thread 计算多个 C 条目。
我们现在使用大小为 BM*BK + BN*BK = 64*8 + 64*8 = 1024 个 floats 的 SMEM cache，每个 block 总共 4KB。
下面是一个可视化图。
我用橙色和红色高亮了两个 threads 以及它们在 inner loop 中访问的数值。

该 kernel 的所有重要更改都发生在 inner loop 中。
从 GMEM 到 SMEM 的加载基本上保持不变。
让我们看一看：

Godbolt 链接

// 在 registerfile 中为结果分配 thread-local 缓存
float threadResults[TM] = {0.0};

// 外层循环处理 block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
  // 填充 SMEM 缓存（与之前相同）
  As[innerRowA * BK + innerColA] = A[innerRowA * K + innerColA];
  Bs[innerRowB * BN + innerColB] = B[innerRowB * N + innerColB];
  __syncthreads();

  // 为外层循环移动 blocktile
  A += BK;
  B += BK * N;

  // 计算 per-thread 结果
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    // 我们将 dotproduct 循环作为外层循环，
    // 这样可以方便地复用 Bs 的条目，我们可以将其缓存在 tmp 变量中。
    float Btmp = Bs[dotIdx * BN + threadCol];
    for (uint resIdx = 0; resIdx < TM; ++resIdx) {
      threadResults[resIdx] +=
          As[(threadRow * TM + resIdx) * BK + dotIdx] * Btmp;
    }
  }
  __syncthreads();
}

此 kernel 达到了约 8600 GFLOPs，比之前的 kernel 快了 2.2 倍。
让我们计算一下在上一个 kernel（每个 thread 计算一个结果）中，每个 thread 执行了多少次 memory accesses：

• GMEM: 外层循环 K/32 次迭代 * 2 次 loads (实际上是 (BK + BK) * K / BK)
• SMEM: 外层循环 K/32 次迭代 * BLOCKSIZE (=32) * 2 次 loads (实际上是 K * 2)
• 每个结果的 memory accesses：K/16 GMEM, K*2 SMEM
• 此时 BK = 32

而对于我们的新 kernel，每个 thread 计算八个结果：

• GMEM: 外层循环 K/8 次迭代 * 2 次 loads (实际上是 $\frac{K * (BN + BM)}{\frac{BN * BM}{8}}$)
• SMEM: 外层循环 K/8 次迭代 * BK(=8) * (1 + TM(=8))
• 每个结果的 memory accesses：K/32 GMEM, K*9/8 SMEM
• 此时 BN = BM = 64, BK = 8；这里的形状对于 memory access 的量是有影响的

正如预期的那样，我们现在每个 instructions 由于 memory pressure 而 stalling 的 cycles 减少了许多：

注意：坐标轴与之前的图相比已经发生了变化。

关于 Compiler 优化的侧记

上面我们显式地将 B 的条目缓存到了 Btmp 中，并为了效率重新排序了两个 inner loops。
如果我们不这样做，代码看起来像这样：

for (uint resIdx = 0; resIdx < TM; ++resIdx) {
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    threadResults[resIdx] +=
      As[(threadRow * TM + resIdx) * BK + dotIdx] * Bs[dotIdx * BN + threadCol];
  }
}

有趣的是，这并不会对 performance 产生不利影响。
这令人惊讶，因为我们的 inner 两个 loops 现在会产生 BK (=8) * TM (=8) * 2 = 128 次 SMEM 访问，而不是之前的 72 次。
查看汇编 (Godbolt 链接 ：

// 第一个 inner-most loop
ld.shared.f32   %f45, [%r9];
ld.shared.f32   %f46, [%r8];
fma.rn.f32      %f47, %f46, %f45, %f212;
ld.shared.f32   %f48, [%r9+256];
ld.shared.f32   %f49, [%r8+4];
fma.rn.f32      %f50, %f49, %f48, %f47;
ld.shared.f32   %f51, [%r9+512];
ld.shared.f32   %f52, [%r8+8];
fma.rn.f32      %f53, %f52, %f51, %f50;
ld.shared.f32   %f54, [%r9+768];
ld.shared.f32   %f55, [%r8+12];
fma.rn.f32      %f56, %f55, %f54, %f53;
ld.shared.f32   %f57, [%r9+1024];
ld.shared.f32   %f58, [%r8+16];
fma.rn.f32      %f59, %f58, %f57, %f56;
ld.shared.f32   %f60, [%r9+1280];
ld.shared.f32   %f61, [%r8+20];
fma.rn.f32      %f62, %f61, %f60, %f59;
ld.shared.f32   %f63, [%r9+1536];
ld.shared.f32   %f64, [%r8+24];
fma.rn.f32      %f65, %f64, %f63, %f62;
ld.shared.f32   %f66, [%r9+1792];
ld.shared.f32   %f67, [%r8+28];
fma.rn.f32      %f212, %f67, %f66, %f65;
// 第二个 inner-most loop
ld.shared.f32   %f68, [%r8+32];
fma.rn.f32      %f69, %f68, %f45, %f211;
ld.shared.f32   %f70, [%r8+36];
fma.rn.f32      %f71, %f70, %f48, %f69;
ld.shared.f32   %f72, [%r8+40];
fma.rn.f32      %f73, %f72, %f51, %f71;
ld.shared.f32   %f74, [%r8+44];
fma.rn.f32      %f75, %f74, %f54, %f73;
ld.shared.f32   %f76, [%r8+48];
fma.rn.f32      %f77, %f76, %f57, %f75;
ld.shared.f32   %f78, [%r8+52];
fma.rn.f32      %f79, %f78, %f60, %f77;
ld.shared.f32   %f80, [%r8+56];
fma.rn.f32      %f81, %f80, %f63, %f79;
ld.shared.f32   %f82, [%r8+60];
fma.rn.f32      %f211, %f82, %f66, %f81;
// ... 像这样继续 inner-loops 3-8 ...

compiler 展开了两个 loops

compiler 可以展开它们，因为 loop count 在 compile time 是已知的。

然后消除了 Bs 条目的重复 SMEM loads，所以我们最终得到了与我们优化的 CUDA 代码相同数量的 SMEM 访问。

当 PTX 编译为 SASS 时，来自 Bs 的 SMEM loads 被 vectorized 了：

这已经暗示了我们稍后将执行的一项 optimization：转置 As 使得我们也可以 vectorize 那些 loads。

LDS     R26, [R35.X4+0x800] // a 32b load from As
LDS.128 R8,  [R2]           // a 128b load from Bs
LDS.128 R12, [R2+0x20]
LDS     R24, [R35.X4+0x900]
LDS.128 R20, [R2+0x60]
LDS     R36, [R35.X4+0xb00]
LDS.128 R16, [R2+0x40]
LDS.128 R4,  [R2+0x80]
LDS     R38, [R35.X4+0xd00]

改进领域：Arithmetic Intensity

我们目前的 kernel 仍然遭受着与 kernel 3 相同的 stalling-for-memory 问题，只是程度较轻。
所以我们将再次应用相同的 optimization：每个 thread 计算更多的结果。
这让我们的 kernel 运行得更快的主要原因是它增加了 arithmetic intensity。

定义为在 GMEM 和 SMEM 之间传输的每个 byte（load + store！）所执行的 FLOPs 数量。

下面我尝试更直观地说明为什么每个 thread 计算更多结果会提高 arithmetic intensity：

计算一个结果的 square 比计算一个结果的 column 更有效率，因为我们可以 share 更多的 inputs：

总之，我们所有的 kernels 都执行相同数量的 FLOPs，但我们可以通过每个 thread 计算更多结果来减少 GMEM 访问次数。
只要我们仍受 memory bound，我们将继续优化 arithmetic intensity。

Kernel 5: 通过 2D Blocktiling 增加 Arithmetic Intensity

Kernel 5 的基本思想是计算 C 的 8*8 元素 grid per thread。
Kernel 的第一阶段是所有 threads 协同工作以填充 SMEM cache。
我们将让每个 thread 加载多个元素。
代码如下所示：

这是 GMEM 加载的图形表示：

for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
  As[(innerRowA + loadOffset) * BK + innerColA] =
      A[(innerRowA + loadOffset) * K + innerColA];
}
for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
  Bs[(innerRowB + loadOffset) * BN + innerColB] =
      B[(innerRowB + loadOffset) * N + innerColB];
}
__syncthreads();

现在 SMEM cache 已填充，我们让每个 thread 将其相关的 SMEM 条目相乘并将结果累加到 local registers 中。
下面我说明了沿着 input matrices 的（未改变的）outer loop，以及用于 dot product、TN 和 TM 维度的三个 inner loops：

代码中有趣的部分如下：

Godbolt 链接 ：

// 在 registerfile 中为结果分配 thread-local 缓存
float threadResults[TM * TN] = {0.0};
// 为 As 和 Bs 准备的 register 缓存
float regM[TM] = {0.0};
float regN[TN] = {0.0};

// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
  // 填充 SMEM 缓存
  for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
    As[(innerRowA + loadOffset) * BK + innerColA] =
        A[(innerRowA + loadOffset) * K + innerColA];
  }
  for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
    Bs[(innerRowB + loadOffset) * BN + innerColB] =
        B[(innerRowB + loadOffset) * N + innerColB];
  }
  __syncthreads();

  // 推进 blocktile
  A += BK;     // 向右移动 BK 列
  B += BK * N; // 向下移动 BK 行

  // 计算 per-thread 结果
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    // 将相关的 As 和 Bs 条目加载到 registers
    for (uint i = 0; i < TM; ++i) {
      regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
    }
    for (uint i = 0; i < TN; ++i) {
      regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
    }
    // 对 register 缓存执行 outer product，累积
    // 到 threadResults
    for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
      for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
        threadResults[resIdxM * TN + resIdxN] +=
            regM[resIdxM] * regN[resIdxN];
      }
    }
  }
  __syncthreads();
}

在 inner loop 中，我们可以通过将 dotIdx 作为 outer loop，并将两个 inner loops 所需的数值显式加载到 registers 中，来减少 SMEM 访问次数。
下面是随着时间推移的 dotIdx 循环图示，用于可视化在每一步中哪些 SMEM 条目被加载到 thread-local registers：

为了便于绘制，我不得不减少了一些维度。在 kernel 中：BK=TM=TN=8。

最终 performance：16TFLOPs，又一次翻倍提升。
让我们重复 memory access calculation。
我们现在每个 thread 计算 TM*TN = 8*8 = 64 个结果。

• GMEM: K/8 (outer loop iters) * 2 (A+B) * 1024/256 (sizeSMEM/numThreads) loads （实际上是 $\frac{(BM + BN) * K}{\frac{BN * BM}{64}} / 64$，这里除以 64 是 per result）
• SMEM: K/8 (outer loop iters) * 8 (dotIdx) * 2 (A+B) * 8 loads（实际上是 $\frac{8(K + K)}{64}$）
• 每个结果的 memory accesses：K/64 GMEM, K/4 SMEM
• 其中 TM = TN = 8（每个 thread 计算一个 8*8 的 tile）

性能正逐渐达到可接受的水平，然而，由于 memory pipeline 拥塞导致的 warp stalls 仍然过于频繁。
对于 kernel 6，我们将采取两项措施尝试改善这一点：转置 As 以实现 SMEM loads 的 auto-vectorization，并向 compiler 保证 GMEM 访问是对齐的。

Kernel 6: Vectorize SMEM 和 GMEM 访问

我在之前已经暗示过的第一个 optimization 是转置 As。
这将允许我们使用 vectorized SMEM loads（SASS 中的 LDS.128）从 As 加载。

本质上就是通过 transpose 来 enable vectorization。

下面是与 kernel 5 相同的三个 inner loops 的可视化，但现在 As 在 memory 中被转置了：

查看汇编

Godbolt 链接 ：

我们看到将 As 加载到 registers 中（曾经是一个 32b 的 LDS load）现在也是一个 128b 的 LDS.128 load，就像 Bs 已经实现的那样。
这为我们带来了 500GFLOPs 的加速，约 3%。

接下来，我们将使用 vector datatypes（即 float4）向量化所有往返于 GMEM 的 loads 和 stores。

代码看起来像这样：

Godbolt 链接 ：

float4 tmp =
    reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
// 在从 GMEM 到 SMEM 的传输过程中转置 A
As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;

reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
    reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
__syncthreads();

这导致 32b 的 GMEM load 指令（LDG.E 和 STG.E）被 128b 的 counterparts（LDG.E.128 和 STG.E.128）所取代。
最初，我对为什么运行这个：

reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
    reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];

会比手动展开访问（或使用 pragma unroll）更快感到困惑：

Bs[innerRowB * BN + innerColB * 4 + 0] = B[innerRowB * N + innerColB * 4 + 0];
Bs[innerRowB * BN + innerColB * 4 + 1] = B[innerRowB * N + innerColB * 4 + 1];
Bs[innerRowB * BN + innerColB * 4 + 2] = B[innerRowB * N + innerColB * 4 + 2];
Bs[innerRowB * BN + innerColB * 4 + 3] = B[innerRowB * N + innerColB * 4 + 3];

难道 compiler 就不能合并第二个版本并也生成 128b loads 吗？
我认为原因是 compiler 无法验证传递给 kernel 的 float* B 指针是否是 128b 对齐的，而这是使用 LDG.E.128 的一个要求。
因此，reinterpret_cast 的唯一目的就是向 compiler 保证 float* B 指针是对齐的。

将此与 SMEM loads 进行对比，对于后者，compiler 会自动生成 vectorized loads，因为那块内存不是由用户管理的。

Kernel 6 达到了 19TFLOPs。
Profiler 仍显示出一系列问题区域和优化机会：我们遇到了 shared-memory bank conflicts（cuBLAS 会避免这一点），我们的 occupancy 比必要的要高，而且我们还没有实现任何 double buffering（CUTLASS docs 似乎建议这是非常有用的）。

但在讨论这些之前，让我们先处理一些更低挂的果实：Autotuning kernel 的 parameters。

Kernel 9: Autotuning

我们已经累积了总共五个 template parameters：

• BM、BN 和 BK，指定了我们从 GMEM 缓存到 SMEM 的数据量。
• TM 和 TN，指定了我们从 SMEM 缓存到 registers 的数据量。

对于 kernel 6，这些参数被设置为 BM=BN=128 和 BK=TM=TN=8。
我编写了一个 bash script，搜索所有合理的组合并对它们的运行时间进行 benchmark。
这要求我确保：

1. 我知道哪些 parameter 组合是合理的，并跳过那些不合理的。

   不合理配置的一个例子：我希望 vectorize 所有的 SMEM loads，因此 BM * BK（As 的大小）需要能被 4 * NUM_THREADS 整除，因为每个 thread 在从 GMEM 到 SMEM 的加载循环的每次迭代中都会发射一个 4-wide load。

2. 对于剩余的约 400 个不同的 hyperparameter 设置，kernel 实现都是正确的。

对代码进行必要的修改最终花费了不少时间。

事实证明，最优 parameters 会根据 GPU 型号的不同而有很大差异。

我想这就是为什么像 Triton 这样的 compilers 提供了 routines 用于 autotuning。我很好奇 cuBLAS 是如何运作的，它们可能在 cuBLAS binary 内部存储了一个预先计算好的从 {GPU type, matrix size, dtype, …} 到最优 GEMM 实现的 mapping。

在我的 A6000 上，BM=BN=128 BK=16 TM=TN=8 将性能提高了 5%，从 19 TFLOPs 增加到 20 TFLOPs。
在 A100 SMX4 40GB 上，同样的配置达到了 12 TFLOPs，比 autotuner 找到的最优设置（BM=BN=64 BK=16 TM=TN=4，达到了 12.6 TFLOPs）差了 6%。

A100 的 fp32 性能比 A6000 差，这就是为什么 FLOPs 数值较低的原因（cuBLAS 在 A100 上达到 14.7 TFLOPs）。Nvidia 对 A100 的标定值为 19.5 TFLOPs，A6000 为 38.7 TFLOPs。

我无法解释为什么这些特定的 parameters 最终产生了最优性能。
Autotuning 是有效的，每个高性能库都在使用它，但它也让人觉得很不圆满。

我确信，如果有足够的时间，有足够的机会接触底层 performance counters 并与 Nvidia 工程师交流，我最终会弄明白的。坚信计算机是可以被理解的是一件好事。

Kernel 10: Warptiling

目前，我们的循环结构看起来像这样：

我们现在将在 blocktiling 和 threadtiling 循环之间增加另一层 tiling：warptiling。
Warptiling 最初有些令人困惑，因为与 blocks 和 threads 不同，warps 不会显式出现在 CUDA 代码中的任何地方。
它们是一种 hardware feature，在 scalar CUDA-software 世界中没有直接对应的概念。
我们可以通过 warpId=threadIdx.x % warpSize 计算给定 thread 的 warpId，其中 warpSize 是一个内置变量，在我工作过的任何 CUDA GPU 上都等于 32。

实际上我们在优化 L0 cache (register cache)，这个 cache 非常小，由 reuse bit 控制（大约只有六个槽位），通过让 warp 处理的数据更紧密，使 register 的 reuse 程度上升。

Warps 与性能相关，因为（除了其他原因外）：

• Warps 是映射到属于 SM 的 warp-schedulers 的调度单位。

  在我的 A6000 上，每个 SM 中有四个 warp schedulers。我设想它看起来像这样：
  

• Shared-memory bank conflicts（我将在未来的文章中讨论）仅发生在同一个 warp 的 threads 之间。
• 最近的 GPU 上存在 register cache，更紧密的 threadtiling 为我们提供了更多的 register cache locality。

Warptiling 在 CUDA 代码中看起来像这样：

Godbolt 链接

// dotIdx 在 SMEM 内容上循环
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
  // 为该 thread 的 warptile 部分填充 registers
  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
    for (uint i = 0; i < TM; ++i) {
      regM[wSubRowIdx * TM + i] =
          As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
             threadRowInWarp * TM + i];
    }
  }
  for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
    for (uint i = 0; i < TN; ++i) {
      regN[wSubColIdx * TN + i] =
          Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
             threadColInWarp * TN + i];
    }
  }

  // 执行 warptile matmul。稍后这将很好地映射到
  // 在 tensor cores 上执行的 warp-wide 矩阵指令。
  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
      // 计算具有 register-cache locality 的 per-thread 结果
      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
          threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
                        (wSubColIdx * TN) + resIdxN] +=
              regM[wSubRowIdx * TM + resIdxM] *
              regN[wSubColIdx * TN + resIdxN];
        }
      }
    }
  }
}

我尽力在下面可视化了所有三个层级的 tiling，虽然结构正变得相当复杂。

关于高效 GEMMs 的 CUTLASS 文档 对 warptiling 进行了更深入的研究，他们的可视化图非常具有启发性。

每个 warp 将计算大小为 (WSUBN * WNITER) x (WSUBM * WMITER) 的 chunks。
每个 thread 计算 WNITER * WMITER 个大小为 TM*TN 的 chunks。

在 autotuning 参数后，性能在 A100 上从 19.7 TFLOPs 提高到 21.7 TFLOPs。

这里有一个图表，将我们的 warptiling kernel 与 cuBLAS 在不断增加的 matrix 大小下进行了比较：

我是在 A100 上生成的这张图，这就是为什么绝对 FLOPs 数值不同。

在维度 2048 和 4096 处，我们测得的 FLOPs 仅比 cuBLAS 慢几个百分点。
然而，对于较小的 matrices，与 Nvidia 的库相比，我们的表现很差！
这是因为 cuBLAS 包含的不是一个单一的 SGEMM 实现，而是数百个。

我想这就是该库有 500MB 编译代码的原因。打印所有 kernels：cuobjdump --list-text <cublas location>。

在运行时，根据维度，cuBLAS 会选择运行哪个 kernel。

我在 4096 以内的所有维度上启动了方阵的 matmuls，发现了 16 种不同的 SGEMM kernels。这里有一个脚本用于查找由 cuBLAS 启动的 kernel（感谢 Horace He）。

我追踪了 cuBLAS 调用，这些是它在每个大小时调用的 kernels：

我使用了 Nsight Systems CLI 来完成此操作。

在维度 256 处，它调用了两个 kernels：一个 matmul kernel 接着一个 reduction kernel。

Split-K 指的是跨多个 threadblocks 划分 K 维度。这意味着每个 block 仅计算 C 的 chunks 的一部分，cuBLAS 随后使用一个 reduce kernel 来累加最终结果。这需要一些额外的 memory 空间来在 reduction 之前存储中间结果。我设想它看起来像这样（但我对此并不确定）：

所以，如果我们试图编写一个高性能库，使其适用于所有 shapes 和 sizes，我们将针对不同的 shapes 进行特殊化处理，并在运行时 dispatch 到最合适的那个。

我还想报告一个负面结果：对于此 kernel，我还实现了一项名为 thread swizzling 的优化。
该技术假设 threadblocks 是按 blockIdx 递增的顺序启动的，并优化 blockIdx 到 C chunks 的映射方式，以增加 L2 locality。

请记住 L2 是为整个 GPU 存在一次的 global memory 的 cache。

这篇 Nvidia post 提供了更多信息和可视化。
它并没有提高性能，大概是因为 L2 hit rate 已经相当高，达到 80%，所以我最终删除了 swizzling 代码。

如果有人感兴趣，commit 在 这里。

将 BK 上的循环移向外层是有意义的，因为它遵循了我们的准则：“加载一些数据，然后对该数据执行尽可能多的工作”。
这进一步意味着在 BK 循环内部发生的所有 计算 都将是独立的，并且可以被 parallelized（例如使用 ILP）。

我们现在也可以开始 prefetching 下一次循环迭代所需的数据，这种技术称为 double buffering。

工作进展中：Kernel 11

如果我以后继续研究这篇文章，我会接着看以下内容：

1. Double buffering，为了 computation 和 memory 加载更好的交错 (interleaving)。目前，请参阅 CUTLASS Pipelining。在 CUTLASS 中，double buffering 在两个层面上完成：GMEM ⇒ SMEM 和 SMEM ⇒ Registerfile。
   • 在 Hopper 中，为 warp specialization 引入了新指令，例如用于让某些 warp 比其他 warp 使用更少的 registers。这与直接从 GMEM 加载到 SMEM 而不先经过 registers 的特殊指令相结合，可以用来减轻 register pressure。
2. 摆脱 SMEM bank conflicts。这可以通过优化 SMEM 中的数据布局来完成。
3. 通过查看生成的 PTX，更好地理解在 Triton 中实现的 GEMM kernels。

结论

撰写这篇文章的经历与我之前的在 CPU 上优化 SGEMM 的文章类似：迭代地优化 SGEMM 是深入理解 hardware 性能特性的最佳方式之一。
在编写 CUDA 程序时，我感到惊讶的是，一旦我对 kernel 的工作方式有了很好的可视化，实现代码竟然如此简单。

此外：Powerlaws 无处不在。
我花了两个周末编写了前 6 个 kernels，达到了 peak FLOPs 的 80%，然后又花了 4 个周末进行 autotuning 和 warptiling 才达到 94%。
我在编写这些代码时所学到的东西也看到了收益递减，因此我把追寻最后 6% 的工作推迟到了未来某个时间。

我所有的代码都可以在 Github 上找到。

最后，非常感谢 Godbolt.org 的创建者（用于查看 PTX 和 SASS assembly）以及 Excalidraw 的创建者（用于绘制 kernels）！
这两个工具使用起来都非常愉快，并帮助我学得更快。

---

来 Anthropic 从事 kernels 工作吧！

如果你喜欢这样的 kernel 工作，你很可能非常适合 Anthropic 的 Performance team。来和我一起工作吧！团队由 Tristan Hume 领导，他是我遇到过的最有能力且最体贴的 manager。我们为 GPUs、TPUs 和 AWS Trainium 优化 Anthropic 的模型。

---

进一步的资源和参考文献

• 我开始写这篇文章是因为我偶然发现了 wangzyon 的 Github repository，先是实验了他的 kernels，然后从头开始重写了一切。同样相关的还有这篇关于 CUTLASS 库的 Nvidia Blogpost。
• 必读参考文献：官方的 CUDA Toolkit Programming Guide 和 CUDA Best Practices Guide。Kernel Profiling Guide 包含了更多关于底层 hardware 细节（如 caches 和 pipelines）以及可以收集的各种 metrics 的信息。
• Onur Mutlu 是 ETH 的教授，他将他的讲座上传到 Youtube。与本文特别相关的有 Computer Architecture 和 Acceleration on Heterogeneuous Systems。
• Understanding Latency Hiding on GPUs，这是一篇深入探讨如何设计 workloads 以充分利用 memory bandwidth 和 computation 的博士论文。它出自 2016 年，因此仅涵盖了较旧的 GPU 架构。关于 warp-synchronous 编程的章节已经过时，请参阅 using CUDA warp-level primitives。
• Lei Mao（Nvidia 的一名工程师）在他的 blog 上有很好的 CUDA 内容，包括关于 proper CUDA error handling 的内容。
• 似乎没有任何好的官方资源来理解 SASS。有 Nvidia 关于 CUDA binary 实用程序的文档。更有用的可能是查看开源 SASS assemblers，如 Da Yan 的 turingas。
• 我正在收集可读且经过优化的 CUDA 代码示例供学习：
  • ONNX Runtime 的 CUDA provider，例如他们的 softmax 实现。
  • NVIDIA 的开源 CUTLASS 库，例如他们的 GEMM 实现，该实现使用 double-buffering 来 prefetch 最内层维度，这在我的 kernels 中仍然缺失。