在GPU上的矩阵乘法

GPU上的矩阵乘法技巧' (GPU上的矩阵乘法技巧)

如何在CUDA中实现最先进的矩阵乘法性能。

“从矩阵乘法中汲取灵感,以vaporwave风格创作的极简主义艺术” by DALLE-2

这个博客来源于我突然意识到我对GPU上的矩阵乘法知之甚少。在做了这么多机器学习项目之后,我觉得我应该了解机器学习中最重要的操作是如何工作的:什么是“张量核心”? 为什么每个人都说“数据移动是瓶颈”? GPU实际上有多快?

为了回答这些问题,我决定必须走出我的PyTorch泡泡,冒险进入CUDA的深渊。我写下这篇博客来记录我学到的一切,希望读者不必像我一样费力地翻阅CUDA文档/代码来寻找答案。

在这个旅程中,我学到的一件事就是并行矩阵乘法很难。高效的矩阵乘法严重依赖于您使用的特定硬件和尝试解决的问题规模。没有一种适用于所有情况的解决方案。

别再罗嗦了,让我们开始吧!

回顾GPU体系结构

让我们回想一下(NVIDIA)GPU是如何工作的。GPU通过运行许多线程来实现并行处理。每个线程在一个CUDA核心上执行,尽管在给定时间内,只有一部分线程是活动的,因此可能有比CUDA核心更多的线程。每个线程,无论其是否活动,都有自己的一组寄存器

32个线程组成一个线程束。一个线程束中的所有线程必须一起执行(或者一起处于非活动状态)。在大多数情况下,非活动线程束要远多于活动线程束,而线程束调度器负责在给定时间内选择要执行的线程束。这使得GPU可以通过在等待数据时调度其他线程束来隐藏内存访问的延迟。

一组线程束称为线程块。线程块中的所有线程在同一个流多处理器(SM)中执行。每个线程块都有自己的共享内存,所有线程都可以访问。

注意: 较新的架构

从Volta架构开始,每个线程还具有自己的程序计数器、调用堆栈等。这意味着线程束中的每个线程可以同时执行不同的指令。

Volta架构还引入了张量核心,它专门用于解决特定大小的矩阵乘法。每个活动线程束都可以访问一个张量核心。

在最新的Hopper架构中,引入了线程块群集的概念,表示一组线程块。它可以更精细地控制线程块的调度,并允许同一群集中的其他线程块访问一个线程块的共享内存。

并行化矩阵乘法

假设我们要计算:

在这种情况下,我们说问题规模是(M, N, K)。为了并行化这个操作,我们可以将A和B分成较小的矩阵,分别进行矩阵乘法运算,然后将结果连接起来形成C。

具体来说,我们可以按行将A分割(即将M分割为大小为M’的块),按列将B分割(即将N分割为大小为N’的块),得到:

我们可以看到C中的每个子矩阵都是相互独立的,所以我们可以很容易地并行计算每个子矩阵。

实际应用中,K可能太大无法直接加载到内存并进行计算。相反,典型的实现还会将K分割为大小为K’的块,遍历每个块,并累加(通过求和)部分结果。这被称为串行-K约简(与并行-K约简相对,见下面的部分)。数学上,它看起来像这样:

注意:填充

在问题大小不可被分区大小整除的任何时刻,我们需要添加填充。当我们将分区输入(𝐴ᵢ,ₖ和𝐵ₖ,ⱼ)加载到较低级别的内存中时,我们通常会在加载的分区(对于𝐴ᵢ,ₖ为M’×K’,对于𝐵ₖ,ⱼ为K’×N’)始终添加零,以确保其始终“完整”。在将结果写回全局内存时需要特别小心,以避免越界错误。

从高层次上看,GPU上的矩阵乘法并行计算发生了三次嵌套分区:1. 第一个分区发生在线程块级别。每个线程块负责计算Cᵢ,ⱼ = Aᵢ Bⱼ。2. 第二个分区发生在线程束级别。线程块级别问题Cᵢ,ⱼ进一步分割,使每个线程束负责计算Cᵢ,ⱼ⁽ᵐⁿ⁾ = Aᵢ⁽ᵐ⁾ Bⱼ⁽ⁿ⁾。3. 第三个分区发生在指令级别。某些指令需要特定大小的输入。例如,第二代张量核心对于fp16的问题大小为(16,8,8),而通过纯数量乘法在CUDA核心上的直接实现仅操作大小为(1,1,1)。因此,线程束级别问题进一步分割,使每个块具有适合指令大小的适当大小:Cᵢ,ⱼ⁽ᵐⁿ⁾⁽ᵃᵇ⁾ = Aᵢ⁽ᵐ⁾⁽ᵃ⁾ Bⱼ⁽ⁿ⁾⁽ᵇ⁾。

我们需要三个分区级别的原因在下一节中会看到。

数据冗余

如果我们每次执行计算时都从全局内存重新获取数据到寄存器,那么矩阵乘法很容易成为内存限制。关键观察是,许多子输入Aᵢ和Bⱼ在不同的子矩阵乘法中被重复使用。例如,Aᵢ对于Cᵢ,₁,Cᵢ,₂,…是必需的,而Bⱼ对于C₁,ⱼ,C₂,ⱼ,…是必需的。如果我们可以最小化冗余的数据移动并尽可能地重用加载的数据,那么我们可以获得最佳吞吐量。

在CUDA中,有三种类型的用户可访问的内存:

下面是如何利用每种内存类型的高级视图:

  1. 每个线程块首先将其所需的输入从全局内存加载到共享内存中。对这些数据的后续访问将由共享内存而不是较慢的全局内存提供。
  2. 在每个线程块内,每个线程束首先将其所需的输入从共享内存加载到寄存器中。对这些数据的后续访问将由快速的寄存器直接提供。

深入细节

线程块级别

在线程块级别上,问题被划分为大小为(M’,N’,K’)的子问题。因此,每个线程块负责计算C的一部分,表示为:

通过将子输入Aᵢ,ₖ和Bₖ,ⱼ加载到共享内存中,减少了冗余数据的移动。当我们完成计算Aᵢ,ₖ Bₖ,ⱼ后,将加载下一个块(Aᵢ,ₖ₊₁和Bₖ₊₁,ⱼ)。

线程束级别

在线程束级别上,子问题进一步划分为大小为(M’‘,N’‘,K’’)的子子问题。因此,每个线程束负责计算Cᵢ,ⱼ的片段,表示为Cᵢ,ⱼ⁽ᵐ ⁿ⁾:

通过将子子输入Aᵢ,ₖ⁽ᵐ ˡ⁾和Bₖ,ⱼ⁽ˡ ⁿ⁾加载到寄存器中,减少了冗余数据的移动。然后,线程束内对Aᵢ,ₖ⁽ᵐ ˡ⁾和Bₖ,ⱼ⁽ˡ ⁿ⁾的任何访问都将由快速寄存器提供。

注意:在寄存器间分配数据

值得注意的是,寄存器仅限于线程级别。这意味着寄存器中的输入不能被线程束中的其他线程访问。具体将Aᵢ,ₖ⁽ᵐ ˡ⁾和Bₖ,ⱼ⁽ˡ ⁿ⁾分割到每个线程的寄存器上的方式取决于所使用的具体指令。关于每个指令的详细说明,请参考NVIDIA关于线程束级别矩阵乘加指令的文档。

张量核心级别

为了执行矩阵乘法,我们使用GPU上的张量核心(Tensor Cores)。我的GPU(RTX 2060)具有第二代张量核心,专门解决大小为(M’’‘, N’’‘, K’’’)=(16, 8, 8)的fp16问题。因此,我们进一步将Cᵢ,ⱼ⁽ᵐ ⁿ⁾分割以适应指令期望的大小:

在这里,所有输入已经在寄存器中,因此数据移动开销很小。

注意:张量核心

张量核心操作是线程束级别指令,这意味着线程束中的所有线程需要同时执行张量核心指令,共同准备数据以被一个张量核心消耗。

选择分区大小

因此,鉴于我们希望最大程度地减少数据移动,我们应选择尽可能大的分区大小以利用所有的共享内存和寄存器,对吗?嗯,并不完全是这样。

线程块分区大小

从渐近的角度来看,随着问题规模的增加,确实,我们希望尽可能多地使用共享内存和寄存器。然而,对于小规模的问题,我们可能会遇到两个问题:

  1. 大分区大小意味着较少的线程块。由于每个线程块只能在一个SM上执行,这可能意味着我们无法利用所有的SM。
  2. 对于不可被分区大小整除的问题规模,我们将不得不对输入添加更多的填充。这降低了效率,因为将在有意义的输入上进行的计算较少。

一个典型的实现可能使用大小为(M’,N’,K’) = (128,256,32)的分区大小。

Warp分区大小

一般来说,较大的Warp分区大小意味着会有较少的冗余数据移动,但代价是较少的Warp数量。Warp数量过少意味着我们将无法隐藏内存访问的延迟(因为当前的Warp在等待数据时可能会用完其他Warp的调度)。

一个典型的实现可能使用大小为(M”,N”,K”) = (64,64,32)的分区大小。

指令分区大小

这完全取决于您的GPU支持的指令。对于我的RTX 2060,用于fp16 Tensor Core矩阵乘法(带有fp16累积)的ptx指令是mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16,它期望输入的大小为(16,8,8)。

更多优化

上述技术可以使我们在问题规模较大时接近GPU的理论峰值性能。然而,对于较小的问题规模,它们并不高效。有两个常见的技术可以进一步提高矩阵乘法的性能:并行K降维和软件流水线。

并行K降维

在M和N较小的情况下,我们可能只有少数线程块。例如,如果(M’,N’) = (128,256) 并且原始问题的大小满足 M ≤ 128 和 N ≤ 256,我们只会有一个线程块,因此我们只能利用 GPU 的一小部分计算能力! (例如,我的RTX 2060有30个SM,所以为了最大化利用率,我们希望至少有30个线程块。)

在K较大的情况下(即使M和N较小),我们可以通过进行并行K降维来利用更多的并行性。回忆一下,在串行K降维中,每个线程块迭代以下总和:

并将中间结果累加到Cᵢ,ⱼ中。但在并行K降维中,我们只分配每个线程块来计算总和的一个元素(即 Aᵢ,ₖ Bₖ,ⱼ)。这使我们可以将线程块数量增加K/K’倍,从而利用更多的SM。

缺点是现在我们需要分配更多内存来存储每个线程块的结果,并调用第二个内核来对部分结果进行最终降维以获得Cᵢ,ⱼ。

软件流水线

通常,CUDA通过安排其他Warp执行来隐藏内存访问的延迟,而一个Warp在等待数据时。这要求我们有足够的Warp来掩盖延迟。

然而,在执行GEMM时,Warp的数量通常相对较小。这是因为Warp的数量受“每个线程块可用的寄存器数与每个Warp所需的寄存器数之比”限制。对于矩阵乘法,我们使用每个Warp的大量寄存器来最大化数据重用。因此,我们可能没有足够的Warp来掩盖延迟。

“累加器元素通常占用至少半个线程的总寄存器预算。” —— CUTLASS文档

为了缓解这种影响,我们可以使用软件流水线。实质上,我们可以(手动)使用特殊指令异步地预加载下一次迭代的循环输入。在加载输入的同时,我们可以继续在当前迭代上进行计算。如下图所示:

来自CUTLASS的软件流水线

这是可能的,因为GPU与任何现代CPU一样:只要它们之间没有数据依赖性,它就可以对内存访问和算术操作进行流水线处理。这就是所谓的指令级并行性

矩阵乘法的实际应用

如果您想看看所有这些概念如何在实际实现中结合在一起,请查看我在CUDA上从头开始训练MNIST的实现。在那里,我使用CUDA训练了一个多层感知机,对于隐藏层大小为128,实现了相比优化的PyTorch加快了6倍的速度

GitHub – andylolu2/cuda-mnist

通过在GitHub上创建一个帐户,为andylolu2/cuda-mnist的开发做出贡献。

github.com

参考资料

1. CUTLASS文档2. CUDA文档3. CUTLASS示例