Nvidia Tensor Core-CUDA HGEMM优化进阶

目录

1 背景

2 结果

2.1 测试条件

2.2 设备规格

2.3 RTX3090

2.4 RTX A6000 

3 矩阵分块

3.1 Block分块

3.2 Warp分块

4 访存优化

4.1 宽指令访存

4.2 数据复用

4.3 异步拷贝

4.4 消除Bank Conflict

4.4.1 Padding

4.4.2 Permuted

4.5 提高L2 Cache命中率

4.6 提高寄存器重复利用率

5 Pipeline优化

5.1 Double Buffer

5.2 Stage

6 其他

6.1 优化方法

6.2 源码

---

1 背景

GEMM（General Matrix Multiplication）矩阵乘法是深度学习中最常用且最耗时的算法之一，特别是在CNN、RNN、Transformer等领域中。在这些领域中，大量的矩阵乘法操作需要被快速计算和处理。因此，高效的矩阵乘法实现对于深度学习任务的性能和准确性至关重要。

HGEMM（Half-precision General Matrix Multiplication）半精度矩阵乘法在Nvidia GPU上Tensor Core硬件单元的加持下，可以在保持准确性的同时大幅提高计算速度，由此带来的性能优势可以显著改善深度学习中推理和训练任务的实现速度。

Tensor Core的出现为半精度矩阵乘法的优化带来了突破性的进展，在Nvidia GPU上使用Tensor Core来优化半精度矩阵乘法算法已经成为了近年来GPU加速计算的热点研究领域之一。

2 结果

本文主要采用手写WMMA API和MMA PTX CUDA HGEMM Kernel的方式调用Tensor Core，再进行性能调优，并与Cublas的Tensor Core性能作比较，通过探究各种矩阵分块和优化方法，目前在256 ~ 16384维度之间的性能均不低于Cublas性能的95%，许多场景下性能超越Cublas，代码开源在cuda_hgemm。

2.1 测试条件

• HGEMM：C (M * N, Half, Row Major) = A (M * K, Half, Row Major) * B (K * N, Half, Col Major)

• CUDA：11.3

• GPU：RTX3090、RTX A6000

2.2 设备规格

RTX3090和RTX A6000的设备规格如下。

Graphics Card	RTX3090	RTX A6000
GPU Codename	GA102	GA102
GPU Architecture	Ampere	Ampere
GPCs	7	7
TPCs	41	42
SMs	82	84
CUDA Cores / SM	128	128
CUDA Cores / GPU	10496	10752
Tensor Cores / SM	4 (3rd Gen)	4 (3rd Gen)
Tensor Cores / GPU	328 (3rd Gen)	336 (3rd Gen)
GPU Boost Clock (MHz)	1695	1800
Peak FP32 TFLOPS	35.6	38.7
Peak FP16 TFLOPS	35.6	38.7
Peak FP16 Tensor TFLOPS with FP16 Accumulate	142	154.8
Peak FP16 Tensor TFLOPS with FP32 Accumulate	71	154.8
Memory Interface	384-bit	384-bit
Memory Clock (Data Rate)	19.5 Gbps	16 Gbps
Memory Bandwidth	936 GB/sec	768 GB/sec
L1 Data Cache/Shared Memory	10496 KB	10752 KB
L2 Cache Size	6144 KB	6144 KB
Register File Size	20992 KB	21504 KB

2.3 RTX3090

2.4 RTX A6000 

3 矩阵分块

3.1 Block分块

对于固定尺寸的输入矩阵来说，通常Block分块尺寸越大，意味着单个Block内的计算量越大，但是需要并行计算的Block数量越少，这是Block维度计算量和并行度的权衡。

一般来说针对不同尺寸的输入矩阵，需要采用不同的分块策略，才能更好地实现Block维度计算量和并行度的平衡，输入矩阵的尺寸越小，Block分块尺寸越小，输入矩阵的尺寸越大，Block分块尺寸越大。另一方面结合硬件规格的限制，Block分块尺寸一般为32、64、128、256之间的组合。

3.2 Warp分块

确定好Block分块尺寸之后，需要继续确定Warp分块尺寸，一般来说，Warp分块尺寸越大，意味着单个Warp内的计算量越大，但是需要并行计算的Warp数量越少，这是Warp维度计算量和并行度的权衡。

Warp数量过少的话，会导致Warp Occupancy过低，即调度器内的Eligible Warps数量过少，可能会严重影响调度器内部的指令发射周期，最终影响Kernel的性能。另一方面结合硬件规格的限制，Warp数量一般取4、8、16，再结合Block分块尺寸，即可确定最优的Warp分块尺寸，一般为8、16、32、64、128之间的组合。

4 访存优化

4.1 宽指令访存

A矩阵和B矩阵的数据尽可能按16Bytes（int4或float4）从Global Memory加载到Shared Memory，减少加载指令的数目，提高指令的计算访存比，同时减少指令发射的延迟。

4.2 数据复用

如下图所示，Block在计算C矩阵的绿色分块时，需要A矩阵的整个浅蓝色分块和B矩阵的整个浅黄色分块，对于Block内的不同Warp，当不同Warp计算相同的行或者列时，其所需要加载A矩阵的数据或者B矩阵的数据是相同的，也就是Block内部可能会存在重复加载A矩阵或B矩阵的数据。

解决方法是将A矩阵或B矩阵的数据先加载到Shared Memory，实现Block内共享，Block内的Warp再分别从Shared Memory加载数据到寄存器计算。由于Shared Memory的访问延迟远小于Global Memory，因此可以显著缓解Kernel的带宽瓶颈。

4.3 异步拷贝

cp.async是一个非阻塞的异步拷贝PTX指令，可以将数据从Global Memory拷贝到Shared Memory，不仅具有一定的缓存控制策略，还可以自主控制多组发射指令。在一定程度上扩展了数据迁移的优化范围和pipeline并行的优化策略。

cp.async.ca.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], cp-size{, src-size}{, cache-policy} ;
cp.async.cg.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], 16{, src-size}{, cache-policy} ;
cp.async.ca.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], cp-size{, ignore-src}{, cache-policy} ;
cp.async.cg.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], 16{, ignore-src}{, cache-policy} ;

.level::cache_hint =     { .L2::cache_hint }
.level::prefetch_size =  { .L2::64B, .L2::128B, .L2::256B }
cp-size =                { 4, 8, 16 }

4.4 消除Bank Conflict

一般来说，Nvidia GPU的Shared Memory分为32个Bank，每个Bank默认为4Bytes，总共128Bytes。通常情况下，单个Warp内的线程访问Shared Memory的不同Bank时，只需要一次显存请求。如果单个Warp内的线程访问Shared Memory同一Bank内的不同字段时，就会发生Bank Conflict，显存请求就会变成串行执行，显著增加访存延迟。

所以消除Shared Memory的Bank Conflict，几乎是所有CUDA Kernel优化的基本操作，通常采用Padding或者Permuted方式。

4.4.1 Padding

如下图所示，假设蓝色区域是单个Warp要访问的Shared Memory区域，直接访问的话，会发生明显的Bank Conflict。

如果在申请Shared Memory时，每一行再额外申请4个Bank的黄色区域，此时，单个Warp要访问的Shared Memory区域恰好处于不同的Bank，避免了Bank Conflict的发生。 

Padding方式可以有效解决Shared Memory的Bank Conflict问题，但是会申请额外的Shared Memory，增加Block内部Shared Memory的使用量。

4.4.2 Permuted

不申请额外的Shared Memory能否解决Bank Conflict问题，答案是肯定的。

如下图所示，假设A矩阵先从Global Memory加载到Shared Memory，同一颜色区域代表单个Warp的一次访存事务，对应位置存储的话，再使用ldmatrix PTX指令从Shared Memory加载到寄存器，会发生明显的Bank Conflict。如果加载到Shared Memory时进行Permuted操作，再使用ldmatrix PTX指令从Shared Memory加载到寄存器，会避免Bank Conflict的发生。

4.5 提高L2 Cache命中率

前面提到Block分块问题，假如Block分块已完成，那么在实际计算时这些Block分块的计算顺序该如何设计。最直接的方式是按行计算，这种计算方式会带来一个问题，显著降低L2 Cache命中率。

一般来说，对于同一行的C矩阵Block分块，其所需要的A矩阵的分块数据是相同的，同理对于同一列的C矩阵Block分块，其所需要的B矩阵的分块数据是相同的。如果按行计算，对于A矩阵来说，数据相同，L2 Cache命中率很高，但是对于B矩阵来说，数据都不相同，L2 Cache命中率很低，综合来看，考虑到L2 Cache的容量，L2 Cache命中率不会太高。

因此，可以采取如下图所示swizzle计算方式，即“牛耕式”计算，兼顾A矩阵和B矩阵的L2 Cache命中率，提高整体的L2 Cache命中率。同时可以结合Block分块大小和L2 Cache容量，调整“牛耕”步长，获取最高的L2 Cache命中率。

4.6 提高寄存器重复利用率

前面提到Warp分块问题，假如Warp分块已完成，那么在实际计算时Warp内部Tile（Tensor Core计算尺寸）的计算顺序该如何设计。此处，Cutlass源码里提供了一种思路。

对于计算能力在8.0及以上的设备，采取“Right Left Right Left”的计算方式，提高A矩阵寄存器的重复利用率。

对于计算能力在8.0以下的设备，采取“Down Up Down Up”的计算方式，提高B矩阵寄存器的重复利用率。 

5 Pipeline优化

一般来说，CUDA HGEMM的计算流程是先将A矩阵和B矩阵的分块数据按K维度逐批从Global Memory拷贝到Shared Memory，再将Shared Memory中的数据逐批拷贝到寄存器中进行矩阵乘计算，直至所有数据都完成矩阵乘计算，示意图如下。

5.1 Double Buffer

Double Buffer是一种数据预取的方式，也是GEMM优化中常用的方法之一，主要通过申请两个Shared Memory的Buffer，交替加载数据。其中一个Buffer在加载数据时，另一个Buffer已经完成数据加载，可以进行矩阵乘计算，这样就可以实现数据加载和矩阵乘计算的并行，隐藏数据加载的延迟。如下图所示为Double Buffer的工作流程，浅蓝色为数据加载Buffer，深蓝色为矩阵乘计算Buffer，两条流水线可并行执行。

5.2 Stage

在Double Buffer的基础上，思考这样一个问题，假如在Step 1时，数据加载和矩阵乘计算并行执行，但是矩阵乘计算先执行完，数据加载还没有完成，这时会出现Tensor Core计算单元在等待数据的情况，也就是SM会出现空闲状态，SM整体利用率会不高。怎么解决这个问题，也就是如何让Tensor Core计算单元不处于等待数据的情况。

既然是带宽瓶颈问题，那么可以再引入更多的Shared Memory Buffer，预取时可以完成多个Buffer的数据加载，这样就可以得避免Tensor Core计算单元处于等待数据的情况。如下图所示为Stage是3时的工作流程，每个Step的数据加载与两个Step的矩阵乘计算并行，进一步隐藏数据加载的延迟。同样地，在Shared Memory充足的情况下，Stage可以继续增大，直到可以充分隐藏数据加载的延迟。因此，Stage的选择取决于设备带宽和Tensor Core算力。

6 其他

6.1 优化方法

本文主要介绍了CUDA HGEMM的通用优化方法和一些特殊优化方法，后续可能会就某一个优化点详细聊聊优化经验。针对不同的GPU和CUDA版本，达到最优性能的优化策略是不一样的。

6.2 源码

本文使用的所有优化方法都开源在cuda_hgemm，包含WMMA API和MMA PTX的实现代码，Block分块尺寸（256*128）和Warp分块尺寸（64*64）是固定的，后续可能会结合源码进行分析。