[博客翻译]在单精度通用矩阵乘法中击败cuBLAS

由 发布于 矩阵乘法CUDA优化GPU编程并行计算高性能计算
 收藏

原文地址：https://salykova.github.io/sgemm-gpu


超越cuBLAS的单精度通用矩阵乘法

1. 引言

我清楚地记得Andrej关于现有CUDA学习材料与高性能库中使用的CUDA代码之间差距的帖子：
ak_post
确实，当涉及到SGEMM实现时，有一些优秀的教育博客文章，例如：

  1. 如何优化CUDA矩阵乘法内核以达到类似cuBLAS的性能（由Andrej提到）
  2. CUDA矩阵乘法优化

这些文章逐步分解了如何优化CUDA矩阵乘法内核。然而，就实现的性能而言，它们都无法接近cuBLAS或CUTLASS的速度，尤其是在使用最新的CUDA版本并进行适当基准测试时。根据我的实验，这些实现最多只能达到cuBLAS性能的50-70%。此外，我发现这两篇博客文章在最终优化步骤中的解释有些过于复杂。尽管如此，我仍然认为这些资源对于任何开始学习CUDA编程的人来说都是很好的，因为它们提供了良好的基础知识。
另一方面，我看到一些非常快的SGEMM实现，具有cuBLAS级别的性能：

  1. YHs GEMM
  2. how-to-optimize-gemm

问题是它们没有文档，难以找到和理解，尤其是对于CUDA初学者。CUTLASS也存在类似的问题。虽然它性能卓越，但缺乏解释其内部设计和如何在高效CUDA/PTX中实现的入门或教育材料。另一个值得注意的项目是MaxAs，这是Scott Gray十多年前为Maxwell架构开发的汇编器。该工具允许直接在SASS（NVIDIA GPU的汇编语言）中编程，直接与硬件通信，而不是依赖与硬件无关的CUDA/PTX。使用MaxAs，Scott编写了一个SGEMM实现，达到了GM204芯片理论最大FLOPS的约98%，平均比cuBLAS高出5%。虽然结果令人印象深刻，但用SASS编程缺乏灵活性，并且需要对底层硬件有深入的理解。此外，随着编译器的显著进步，直接使用SASS编程仅在特殊情况下才有优势（例如，如果您构建tinygrad）。CUTLASS仅使用CUDA/PTX代码就在各种GPU架构和矩阵大小上实现了与cuBLAS相当的性能。
但我们真的能超越cuBLAS的障碍吗？在接下来的章节中，我们将简要回顾CUTLASS中使用的高级SGEMM设计，并讨论如何将此设计转化为高效的CUDA/PTX。本指南仅假设您对CUDA编程模型和线性代数有基本了解。如果您是CUDA编程的新手，我强烈建议从这些简短的入门文章开始：

  1. CUDA C和C++的简单介绍
  2. 如何在CUDA C/C++内核中高效访问全局内存
  3. 在CUDA C/C++中使用共享内存
  4. 通过向量化内存访问提高性能

在继续实现之前，让我们谈谈如何在NVIDIA GPU上进行基准测试——这是一个经常被忽视的话题。正确地进行基准测试与代码本身同样重要，特别是在比较不同实现时。

2. 如何在CUDA设备上进行基准测试？

测量内核持续时间的最可靠方法是使用NVIDIA Nsight Compute进行性能分析，并手动提取性能数据。为了获得确定性和可重复的结果，Nsight Compute会自动应用以下设置：

  1. 时钟控制：将GPU时钟频率锁定为其基准值
  2. 缓存控制：在每次重放之前刷新所有GPU缓存
  3. 持久模式

或者，您可以手动应用这些设置，并在运行时测量内核持续时间，而不依赖外部性能分析工具。在Ubuntu上，您可以使用以下命令检索基准核心时钟频率：

nvidia-smi base-clocks

例如，在RTX 3090上，基准核心时钟频率为1395 MHz。接下来，您需要内存时钟频率，它与基准核心时钟频率配合使用：

nvidia-smi -q -d supported_clocks

从支持的频率列表中，选择与基准核心频率兼容的最快内存时钟。内存时钟速度通常比核心时钟速度更稳定。要锁定时钟频率并启用持久模式，请运行以下命令：

sudo nvidia-smi --persistence-mode=1
# NVIDIA RTX 3090
sudo nvidia-smi --lock-gpu-clocks=1395
sudo nvidia-smi --lock-memory-clocks=9501

要重置核心和内存时钟频率，您可以使用：

sudo nvidia-smi --reset-gpu-clocks
sudo nvidia-smi --reset-memory-clocks
sudo nvidia-smi --persistence-mode=0

GPU时钟频率可能会因GPU的热状态而下降，但对于高性能应用，节流通常是由功率限制引起的。硬件故障也可能导致节流。至少在测试运行期间监控GPU的状态是一个好主意。使用以下命令实时跟踪功耗、时钟速度和节流原因：

watch -n 0.1 nvidia-smi --query-gpu=power.draw,clocks.sm,clocks.mem,clocks_throttle_reasons.active --format=csv

示例输出可能如下所示：

308.50 W, 1395 MHz, 9501 MHz, 0x0000000000000000

位掩码0x0000000000000000表示没有节流，时钟以其最大速度运行。值0x0000000000000001表示空闲状态。任何其他值都表明正在发生节流。有关位掩码值及其含义的完整列表，请参阅NvmlClocksThrottleReasons文档。
一旦您锁定了时钟频率，您可以直接在CUDA中使用CUDA事件测量内核持续时间。以下是一个示例：

cudaEvent_t start, stop;
cudaEventCreate(&start); cudaEventCreate(&stop);
float elapsed_time_ms = 0.0;
cudaEventRecord(start);
kernel<<<...>>>(...);
cudaEventRecord(stop);
cudaEventSynchronize(stop);
cudaEventElapsedTime(&elapsed_time_ms, start, stop);

为了获得可靠的测量结果，通常使用多次重放。在这种情况下，应在每次内核重放之前刷新GPU缓存。这可以使用cudaMemsetAsync完成，如nvbench所示：

// 刷新L2缓存
int dev_id{};
int m_l2_size{};
void* buffer;
checkCudaErrors(cudaGetDevice(&dev_id));
checkCudaErrors(cudaDeviceGetAttribute(&m_l2_size, cudaDevAttrL2CacheSize, dev_id));
if (m_l2_size > 0) {
  checkCudaErrors(cudaMalloc(&buffer, static_cast<std::size_t>(m_l2_size)));
  int* m_l2_buffer = reinterpret_cast<int*>(buffer);
  checkCudaErrors(cudaMemsetAsync(m_l2_buffer, 0, static_cast<std::size_t>(m_l2_size)));
  checkCudaErrors(cudaFree(m_l2_buffer));
}

将时钟频率锁定为其基准值是测量内核速度的可靠方法。然而，在现实场景中，算法通常不会在锁定时钟的情况下运行。为了实现最佳性能，您的算法需要既快速又节能。您的算法消耗的功率越少，硬件可以维持的时钟速度就越高。NVIDIA GPU通常在达到功率限制之前就积极降低时钟频率，这可能会显著降低应用程序性能。为了考虑到这一点，我们在锁定和解锁时钟条件下对实现进行基准测试，测试速度和功率效率。在我们的基准测试中，我们评估矩阵大小从1024到12,800，步长为128。对于每个矩阵大小，我们启动1000*exp((-matsize+1024)/3100.0))次内核重放，并将运行时间计算为后半部分重放的平均值。例如，给定矩阵大小问题m=n=k=4096，我们运行sgemm 1000*exp((-4096 + 1024)/3100.0))=371次，并测量最后185次运行的平均持续时间，确保时钟已经稳定。这种性能分析策略即使在GPU时钟解锁的情况下也能产生一致且可重复的结果。

避免使用WSL进行性能测量。为了确保准确和可靠的结果，请使用原生Linux环境。

3. 内存布局

在本实现中，我们假设矩阵以行主序存储。一个维度为M x N的矩阵A存储为长度为M*N的连续数组。元素A[row][col]通过一维原始C指针ptr[row*N + col]访问，其中0<=col<=N-1和0<=row<=M-1。矩阵乘法表示为C=AB，其中矩阵A, B, C的形状分别为M \times K, K \times N, 和 M \times N。mem_layout
要将此实现调整为以列主序存储的矩阵，只需交换操作数A和B，因为：
\[C^\text{T} = (A B)^\text{T} = B^\text{T} A^\text{T},\]
这里，A, B, C是以行主序存储的矩阵，而A^\text{T}, B^\text{T}, C^\text{T}是对应的转置矩阵（即以列主序存储）。
cuBLAS提供了一个API来计算SGEMM：

cublasSgemm(m, n, k, A, lda, B, ldb, C, ldc); // 简化形式

其中m, n, k表示矩阵大小M, N, K。参数lda, ldb, ldc分别是矩阵A, B, C的_前导维度_。前导维度是迭代矩阵元素时变化最快的维度的长度（即第一维度的长度）。对于以行主序存储的矩阵，前导维度通常是列数，因此通常lda=k, ldb=n, ldc=n。然而，情况并非总是如此。在需要计算较大矩阵的子矩阵的情况下，前导维度可能大于列数。
矩阵也可能用零填充以支持向量化内存加载或张量核心。向量化加载指令允许使用一条指令一次加载多个元素。尽管向量化加载减少了指令总数并提高了带宽利用率，但它们也对输入数据施加了对齐约束，因此前导维度必须能被2（对于64位加载）或4（对于128位加载）整除。下图说明了128位（=4个浮点数）加载的情况。
mem_align
请注意，如果前导维度不能被4整除，则无法在不触及下一行元素的情况下加载第一行的元素。用零填充有助于解决这个问题，但需要额外的内存。另一种解决方案是在运行时检查前导维度是否能被4整除。如果是，则使用向量化加载，否则使用标量加载。此外，过去常用零填充来启用张量核心计算。例如，在cuBLAS版本<11中，Tensor Core FP16操作要求m, n, k是8的倍数。

4. 并行线程执行

CUDA编译.cu文件的轨迹如下：
ptxas
在第一阶段，CUDA代码被编译为PTX（并行线程执行）指令——中间高级代码，可以视为虚拟GPU架构的汇编。这样的虚拟GPU完全由它提供给应用程序的功能或特性集定义。PTX不直接在任何真实架构上运行。它必须经过优化并转换为本地目标架构指令（第二阶段）。NVIDIA提供了一种机制，可以将PTX代码插入到您的CUDA程序中，这样您可以在源代码中混合使用CUDA/PTX，并在PTX生成期间仍能享受代码优化的好处。通过将部分代码重写为PTX，您可以1）减少生成的PTX指令总数2）精确指定所需的PTX指令3）通过限定符调整指令4）应用编译器缺乏或C++语言扩展禁止的优化。重要提示！使用内联PTX汇编不会自动使您的代码比用CUDA编写的代码更快。只有当您手写的PTX比编译器生成的更好时，它才会更快。
在本实现中，我们将直接以PTX编程算法的某些部分，因此如果您从未使用过内联PTX汇编，我强烈建议您查看内联PTX汇编的简短概述。PTX指令有详细文档，可在PTX指令集中找到。我们现在将简要回顾本实现中使用的PTX指令。

4.1. 全局内存加载

对于全局内存加载，我们将使用ld.global.f32指令。这里，“ld”表示“加载”，“f32”表示“32位浮点数”。以下CUDA代码

float reg; // 单个浮点寄存器
float* gmem_ptr = data_in_global_memory; // 指向全局内存的指针
reg = *gmem_ptr; // 全局内存 -> 寄存器传输

可以在PTX中实现为：

float reg; // 单个浮点寄存器
float* gmem_ptr = data_in_global_memory; // 指向全局内存的指针
asm volatile("ld.global.f32 %0, [%1];" : "=f"(reg) : "l"(gmem_ptr));

"=f"中的f表示float数据类型，=修饰符指定寄存器被写入。l表示无符号64位整数。我们还使用volatile关键字确保在生成PTX期间不会删除或移动该指令。

4.2. 全局内存存储

对于全局内存存储，有`st.global.f32指令：

float reg; // 单个浮点寄存器
float* gmem_ptr = data_in_global_memory; // 指向全局内存的指针
// *gmem_ptr = reg; 可以在PTX中实现为：
asm volatile("st.global.f32 [%0], %1;" : : "l"(gmem_ptr), "f"(reg));

4.3. 全局到共享内存传输

当您在CUDA中编写如下代码时：

__shared__ float smem_ptr[n]; // 指向共享内存的指针
float* gmem_ptr = data_in_global_memory; // 指向全局内存的指针
*smem_ptr = *gmem_ptr; // 全局到共享内存传输

会发生一个两步过程。首先，数据从全局内存提取到寄存器，然后将该数据从寄存器复制到共享内存。此外，在传输过程中数据会在所有缓存级别中缓存。
standard_ld
全局到共享内存传输
因此，PTX中的全局到共享内存传输由两个数据移动指令ld.global和st.shared组成：

__shared__ float smem_ptr[n]; // 指向共享内存的指针
uint64_t smem_addr;
// 将通用地址转换为共享地址（st.shared指令的存储位置）
asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(smem_addr) : "l"(smem_ptr));
float* gmem_ptr = data_in_global_memory; // 指向全局内存的指针
float buffer;
// 全局内存 -> 寄存器
asm volatile("ld.global.f32 %0, [%1];" : "=f"(buffer) : "l"(gmem_ptr));
// 寄存器 -> 共享内存
asm volatile("st.shared.f32 [%0], %1;" : : "l"(smem_addr), "f"(buffer));

在Ampere架构之前，无法将数据从全局内存直接传输到共享内存，绕过寄存器存储。从Ampere架构开始，有异步拷贝指令允许这样做。这些指令的使用将在后面演示。

4.4. 向量化共享内存加载和存储

在PTX中，您还可以实现向量化内存操作（使用一条指令加载/存储多个元素）。这里，v4表示包含四个元素的向量：

float reg0, reg1, reg2, reg3;
uint64_t addr;
...
// 共享内存128位加载
asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];"
       : "=f"(reg0), "=f"(reg1), "=f"(reg2), "=f"(reg3)
       : "l"(addr));
// 共享内存128位存储
asm volatile("st.shared.v4.f32 [%0], {%1, %2, %3, %4};"
       :
       : "l"(addr), "f"(reg0), "f"(reg1), "f"(reg2), "f"(reg3));

4.5. 谓词执行

在PTX中，条件执行使用可选的保护谓词实现。以下CUDA代码：

float reg;
float* ptr; // 指向全局内存的指针
unsigned guard;
...
if (guard != 0) {
  *ptr = reg;
}

可以转换为PTX：

float reg;
float* ptr;
unsigned guard;
...
asm volatile(".reg .pred p;\n\t" // 声明谓词'p'
       ".setp.ne.u32 p, %2, 0;\n\t" // 如果(guard != 0)则将'p'设置为true；