CUDA GEMM优化:从内存访问到Tensor Core的深度性能调优

1. 从“能用”到“榨干”:为什么GEMM优化是CUDA程序员的必修课

如果你在GPU上跑过深度学习训练或者科学计算,那你一定对GEMM(通用矩阵乘法)这个词不陌生。它几乎是所有计算密集型应用的基石,从卷积神经网络的前向传播和反向传播,到物理模拟、金融建模,底层都在疯狂地调用GEMM。很多人觉得,用上CuBLAS或者类似的高性能库,性能问题就解决了。但现实是,当你面对一个定制化的、非标准的矩阵运算,或者库函数因为数据布局、问题规模等原因达不到预期性能时,你才会发现,亲手优化一个GEMM内核,是理解GPU编程精髓、真正掌控硬件性能的必经之路。

这不仅仅是把两个循环搬到GPU上那么简单。一个“能用”的CUDA GEMM内核和一个“榨干”硬件性能的内核,其执行效率可能相差几十甚至上百倍。优化的过程,是一个与GPU硬件架构(流多处理器SM、寄存器、共享内存、全局内存带宽、内存合并访问)、CUDA编程模型(线程层次结构、内存模型)以及具体问题特性(矩阵规模、数据精度、是否转置)深度对话的过程。最近社区里频繁出现的“RuntimeError: CUDA error: no kernel image is available for execution”这类错误,其根源往往也在于对目标GPU架构的计算能力(Compute Capability)理解不足,编译出的内核代码无法在目标设备上执行,这本身就是优化前需要扫清的基础障碍。

本文将从一个最基础的、每个CUDA初学者都能写出来的GEMM内核开始,逐步拆解优化策略。我们会像剥洋葱一样,一层层深入,探讨如何通过优化内存访问、利用内存层次结构、增加计算强度、隐藏访存延迟等手段,将一个“玩具级”内核,优化为一个具备实用性能的高效内核。无论你是正在为科研计算寻求加速,还是希望深入理解高性能计算库的底层原理,这篇手把手的优化指南都将为你提供清晰的路径和可复现的代码。

2. 起点:一个朴素到令人“心痛”的Baseline实现

在开始任何优化之前,我们必须先建立一个性能基线。这个基线内核实现了最基本的矩阵乘法算法,它“正确”,但几乎集齐了所有导致性能低下的“反面教材”。理解它为什么慢,是后续所有优化的出发点。

假设我们要计算C = A * B,其中AM x K矩阵,BK x N矩阵,CM x N矩阵。在CPU上,我们通常会写一个三重循环。在GPU上,最直观的想法是让每个线程负责计算输出矩阵C中的一个元素。

2.1 Baseline内核代码与性能分析

__global__ void gemm_naive(float* A, float* B, float* C, int M, int N, int K) { // 计算当前线程负责的C矩阵的行列索引 int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; // 边界检查:确保线程不会计算矩阵范围之外的元素 if (row < M && col < N) { float sum = 0.0f; // 内积循环:计算A的一行和B的一列的点积 for (int k = 0; k < K; ++k) { // A[row, k] 和 B[k, col] sum += A[row * K + k] * B[k * N + col]; } // 将结果写回全局内存的C矩阵 C[row * N + col] = sum; } }

这个内核的线程网格配置通常是这样:dim3 blockDim(16, 16); dim3 gridDim((N + blockDim.x - 1) / blockDim.x, (M + blockDim.y - 1) / blockDim.y);。即,一个线程块包含16x16=256个线程,整个网格由足够覆盖输出矩阵C所有元素的线程块组成。

为什么它这么慢?核心问题在于内存访问模式。

  1. 全局内存的合并访问(Coalesced Access)灾难:GPU的全局内存(显存)访问带宽很高,但延迟巨大。为了高效利用带宽,GPU设计了一次性从连续地址读取一大块数据(例如128字节)的机制,称为“合并访问”。在我们的Baseline内核中:

    • 对于矩阵A的读取:A[row * K + k]。在同一线程束(Warp,通常是32个线程)中,相邻的threadIdx.x(即col不同)的线程,它们需要读取的A元素都在同一行(row相同),但列索引k是连续的。这看起来是连续的,对吗?但注意,A是按行存储的。当k变化时,地址row*K + k确实是连续变化的。然而,这里有一个更致命的问题。实际上,一个Warp内的线程是沿着threadIdx.x方向连续的。在我们的配置中,一个Warp会覆盖32x1的线程(假设blockDim.x=32)。但我们的block是16x16,一个Warp会先取满threadIdx.x方向的16个线程,再取下一行的16个线程。这导致同一个Warp内线程的row值可能不同,它们访问的A的行首地址row*K相差了整整K个元素。这导致Warp内线程访问的全局内存地址完全不连续,无法合并,相当于发起了32次低效的小内存事务。这是性能的第一大杀手。
    • 对于矩阵B的读取:B[k * N + col]。这更糟糕。k是循环变量,对于内层循环的每次迭代,一个Warp内所有线程的k值相同,但col不同。这意味着它们访问的是B矩阵的同一行(第k行)的不同列。由于矩阵是按行存储的,同一行不同列的元素在内存中是连续的。这看起来是连续的!但慢着,这里有一个“跨步”问题。B[k * N + col]中,k * N是行首地址,col是偏移。对于同一个Warp,k*N相同,col是连续变化的(假设线程在col方向连续)。这理论上可以形成合并访问。但是,这要求N不是某些特定的值(比如非常大的质数),并且内存地址对齐良好。在简单情况下,这可能是Baseline中唯一稍微“高效”一点的访问,但依然不理想。
  2. 极高的全局内存访问与计算比(Arithmetic Intensity):在这个内核中,每个输出元素C[i][j]的计算需要进行K次乘加运算(2*K次浮点操作),但同时需要从全局内存读取2*K个浮点数(A和B各K个)。计算强度(每次内存访问对应的浮点操作数)大约是(2*K次FLOP) / (2*K*4字节) = 0.25 FLOP/Byte。这个值极低。现代GPU(如NVIDIA A100)的峰值计算能力(FP32)超过19 TFLOPS,而显存带宽约2TB/s。要喂饱计算单元,计算强度需要达到19e12 FLOP/s / 2e12 Byte/s ≈ 9.5 FLOP/Byte。我们的Baseline强度差了近40倍!这意味着内核99%的时间都在等待数据从显存中读取,计算单元几乎在“空转”。

  3. 没有利用任何高速缓存:Baseline内核反复从全局内存读取AB的数据。GPU有L1/L2缓存,但由于糟糕的访问模式(非合并、跨大步长),缓存命中率会非常低,无法有效缓解带宽压力。

实测中,对于一个1024x1024的方阵乘法,这个Baseline内核在RTX 4090上的性能可能只有几十GFLOPS,不到硬件峰值性能的1%。它为我们后续的优化提供了巨大的提升空间。

3. 优化第一战:利用共享内存实现数据复用

优化GEMM最经典、最有效的一步就是引入共享内存(Shared Memory)。共享内存是位于每个流多处理器(SM)上的片上高速内存,其带宽比全局内存高一个数量级,延迟低得多。我们的核心思想是:将计算一个输出块所需的数据块从全局内存加载到共享内存中,然后在共享内存中进行高速的数据复用。

3.1 分块(Tiling)策略与内核设计

我们不再让一个线程计算一个输出点,而是让一个线程块(Thread Block)协作计算输出矩阵C的一个子块(Tile)。假设我们决定每个线程块计算BM x BN大小的C子块。为了计算这个子块,我们需要从A中读取BM x BK的子块,从B中读取BK x BN的子块。这里BK是内积维度K上的分块大小。

由于K可能很大,我们无法一次性将整个A的行和B的列都塞进共享内存(容量有限,通常几十KB)。因此,我们需要沿K维度进行循环分块。在每一次外循环中,我们将A的一个BM x BK块和B的一个BK x BN块加载到共享内存中,然后线程块内的所有线程协作,利用这两块共享内存中的数据,更新它们各自负责的C子块的部分和。循环遍历完K维度所有块后,每个线程将其累加的部分和写回全局内存的C矩阵。

#define BM 128 // C子块的行维度 #define BN 128 // C子块的列维度 #define BK 8 // 内积维度的分块大小(共享内存中A、B子块的内部维度) __global__ void gemm_tiled(float* A, float* B, float* C, int M, int N, int K) { // 声明共享内存,用于存储A和B的数据块 __shared__ float As[BM][BK]; __shared__ float Bs[BK][BN]; // 线程块负责的C子块在整体矩阵中的起始位置 int blockRow = blockIdx.y * BM; int blockCol = blockIdx.x * BN; // 每个线程在C子块内的相对位置,以及它负责计算的元素 int threadRow = threadIdx.y; int threadCol = threadIdx.x; // 寄存器中累加C的子块,初始化为0 float Csub = 0.0f; // 沿K维度循环分块 for (int k = 0; k < K; k += BK) { // 协作加载:将全局内存中A[blockRow:blockRow+BM, k:k+BK]加载到共享内存As中 // 每个线程加载一个或多个元素。这里假设线程数 >= BM*BK,每个线程加载一个。 if (threadRow < BM && (k + threadCol) < K) { // threadCol在这里充当BK维度索引 As[threadRow][threadCol] = A[(blockRow + threadRow) * K + (k + threadCol)]; } // 协作加载:将全局内存中B[k:k+BK, blockCol:blockCol+BN]加载到共享内存Bs中 // 注意B的索引计算:行索引是k+threadRow,列索引是blockCol+threadCol if ((k + threadRow) < K && threadCol < BN) { Bs[threadRow][threadCol] = B[(k + threadRow) * N + (blockCol + threadCol)]; } // 等待块内所有线程完成共享内存的加载 __syncthreads(); // 利用共享内存中的As和Bs块,计算部分和 for (int ki = 0; ki < BK; ++ki) { Csub += As[threadRow][ki] * Bs[ki][threadCol]; } // 等待块内所有线程完成本次计算,确保共享内存中的数据不再被使用,才能加载下一块 __syncthreads(); } // 将最终结果写回全局内存C if ((blockRow + threadRow) < M && (blockCol + threadCol) < N) { C[(blockRow + threadRow) * N + (blockCol + threadCol)] = Csub; } }

注意:上面的代码是一个高度简化的示意图,它假设线程块的大小(blockDim.x, blockDim.y)至少为(BN, BM),并且BK被巧妙地映射到了线程索引上。在实际的高性能实现中,加载逻辑和线程映射要复杂得多,通常会让每个线程加载多个元素以减少线程同步开销,并精心设计索引映射以达成全局内存的合并访问。

3.2 性能提升原理与参数选择

性能提升的关键

  1. 数据复用:对于C子块中的每个元素,原来需要从全局内存访问ABM个元素和BBN个元素各K次。现在,ABM x BK块和BBK x BN块被加载到共享内存后,在计算C子块时,A的每一行被复用了BN次,B的每一列被复用了BM次。这极大地降低了对全局内存的访问需求。
  2. 访问模式优化:在从全局内存向共享内存加载数据时,我们可以通过精心设计线程的加载任务,确保对全局内存AB的访问是合并的。例如,让一个Warp内的线程连续读取A矩阵的一小段行连续数据,或B矩阵的一小段列连续数据(可能涉及转置存储以优化访问)。
  3. 计算强度提升:数据从全局内存加载到共享内存后,后续的BK次乘加运算都发生在高速的共享内存上。计算强度提升为大约(2*BK次FLOP) / (从全局内存加载2*BM*BK+2*BK*BN字节)。当BMBN较大时,分母中的加载次数被平摊,计算强度显著增加。

参数选择经验

  • BM,BN,BK的选择受限于共享内存容量。每个线程块需要的共享内存大小为BM*BK + BK*BN(以元素计)。例如,若BM=BN=128,BK=8,使用float,则共享内存需求为(128*8 + 8*128) * 4字节 = 8192字节 = 8KB。这通常可以接受。
  • 选择BK时,需要考虑共享内存bank冲突。共享内存被组织成多个bank(例如32个)。如果同一个Warp内的多个线程访问同一个bank的不同地址,就会发生bank conflict,导致串行化访问,降低性能。因此,在存储AsBs时,有时会故意增加一个填充(padding)维度来错开访问,避免bank conflict。
  • BMBN的大小也影响寄存器使用和占用率(Occupancy,即每个SM上活跃的线程块/线程数)。更大的块需要更多寄存器来存储中间累加值Csub,可能降低占用率,但提升了数据复用。需要在两者间取得平衡。

经过这一轮优化,性能通常能有数量级的提升。但距离硬件峰值还有很大差距。

4. 优化第二战:寄存器优化、双缓冲与指令级并行

在利用了共享内存之后,下一个瓶颈往往出现在寄存器使用、内存访问延迟隐藏和指令流水线上。

4.1 寄存器分块:让每个线程计算一个小矩阵

在基础的分块内核中,每个线程只负责输出矩阵C中的一个标量元素。这意味着每个线程只使用很少的寄存器(主要就是一个累加器Csub)。现代GPU拥有大量的寄存器(每个线程多达255个),我们可以利用这些寄存器,让每个线程计算一个小的TM x TN的输出块。这被称为寄存器分块(Register Tiling)。

这样做的好处

  1. 进一步增加数据复用:线程从共享内存中加载一小块ATM x BK)和一小块BBK x TN)到寄存器中,然后用它们计算TM x TN个结果。这比每次只计算一个点复用数据的程度更高。
  2. 减少共享内存的访问频率:原来每计算一个输出点,需要从共享内存读取BKAB的元素。现在,为了计算TM x TN个点,只需要从共享内存加载TMATNB的数据各一次(在BK循环内),然后在寄存器中进行所有组合的乘加运算。这显著降低了共享内存的带宽压力。
  3. 隐藏访存延迟:当线程有多个独立的乘加运算(TM*TN个)可以执行时,GPU的指令调度器可以更好地在等待一次共享内存加载数据的同时,执行其他不依赖该数据的计算指令,从而隐藏延迟。

实现上,线程块的大小会调整为(BN/TN, BM/TM),每个线程在寄存器中声明一个float Creg[TM][TN]的数组。在BK循环内部,线程先协作将共享内存中需要的AB的数据块加载到寄存器变量中(例如float Areg[TM], Breg[TN]),然后通过嵌套循环更新Creg

4.2 共享内存双缓冲(Double Buffering)

在我们之前的内核中,存在一个明显的同步点:__syncthreads()。线程块必须先同步等待所有线程完成共享内存加载,然后才能进行计算;计算完后,又需要同步等待所有线程完成计算,才能加载下一块数据。这个同步点强制线程等待,浪费了计算资源。

双缓冲技术可以缓解这个问题。我们分配两套共享内存缓冲区:As[2][BM][BK]Bs[2][BK][BN]。概念上,我们使用一个“当前”缓冲区进行计算,同时使用另一个“预备”缓冲区异步加载下一块数据。通过一个巧妙的索引切换(例如buf = 1 - buf),在每次BK循环迭代中交换当前和预备缓冲区。

__shared__ float As[2][BM][BK]; __shared__ float Bs[2][BK][BN]; int write_idx = 0; // 用于写入(加载)的缓冲区索引 int read_idx = 0; // 用于读取(计算)的缓冲区索引 for (int k = 0; k < K; k += BK) { // 异步加载下一块数据到 write_idx 缓冲区 // ... 加载代码,使用 write_idx ... __syncthreads(); // 等待本次迭代的数据加载完成(对于read_idx缓冲区) // 使用 read_idx 缓冲区的数据进行计算 // ... 计算代码,使用 read_idx ... // 在计算进行的同时,理论上可以开始准备下一次迭代的加载(但需要下一次循环) // 交换缓冲区索引 int temp = read_idx; read_idx = write_idx; write_idx = temp; // 这个同步点是为了确保所有线程都完成了对当前read_idx缓冲区的计算,然后才能覆盖它(作为下一次的write_idx) __syncthreads(); }

通过重叠计算和通信(加载),双缓冲可以有效隐藏从全局内存加载数据到共享内存的延迟,提升SM的利用率。

4.3 指令级优化:循环展开与向量化内存访问

编译器优化可以帮助我们,但显式地给出提示通常效果更好。

  1. 循环展开(Loop Unrolling):对于内部的BK循环(或寄存器分块内的循环),如果BK是编译时常量,我们可以使用#pragma unroll指令或手动展开。这减少了循环开销(分支预测、递增、比较),增加了指令级并行(ILP)的机会,让编译器能更好地调度指令。
    #pragma unroll for (int ki = 0; ki < BK; ++ki) { Csub += As[threadRow][ki] * Bs[ki][threadCol]; }
  2. 向量化内存访问:GPU支持一次加载或存储多个数据(如float2, float4)。使用向量化加载可以减少指令数量,提高内存吞吐。例如,在从全局内存加载到共享内存时,如果地址对齐且连续,可以让每个线程一次加载一个float4(4个float),而不是4次单独的float加载。
    float4* A_vec = (float4*)A; float4 loaded = A_vec[global_index / 4]; // 然后将loaded的四个分量分别存入共享内存的相应位置
    这要求数据在全局内存中对齐到向量类型的边界,并且线程的访问模式要适配。

5. 优化第三战:适应特定硬件架构的微调

不同的GPU架构(如Ampere, Hopper)有不同的硬件特性,最优的GEMM实现需要针对目标架构进行微调。

5.1 Tensor Core的利用

从Volta架构开始,NVIDIA引入了Tensor Core,这是一种专门为混合精度矩阵乘积累加运算设计的硬件单元。它能在一个时钟周期内执行D = A * B + C,其中A,B,C,D可以是特定大小的矩阵(如4x4 for FP16/FP32混合精度)。使用Tensor Core可以获得比传统CUDA Core高一个数量级的吞吐量。

在CUDA中,可以通过Warp级矩阵操作(WMMA API)来使用Tensor Core。这需要将数据准备成特定的格式(如.row.col布局),并使用wmma::load_matrix_sync,wmma::mma_sync,wmma::store_matrix_sync等函数。编程模型从线程/线程块操作标量,提升到了Warp操作小型矩阵块。

一个使用WMMA的GEMM内核结构大致如下:

  1. 声明Warp级别的累加器矩阵片段(fragment)。
  2. 在外循环中,Warp协作从全局内存通过共享内存加载矩阵AB的片段到寄存器。
  3. 调用wmma::mma_sync执行矩阵乘积累加。
  4. 循环结束后,将结果片段存储回全局内存。

注意事项

  • Tensor Core对数据布局、精度、矩阵大小有严格限制。
  • 需要确保共享内存中的数据排列符合Tensor Core加载指令的要求,否则会导致性能下降甚至错误。
  • 这是目前实现极致性能GEMM的必由之路,CuBLAS等库在支持Tensor Core的硬件上默认会使用它。

5.2 异步拷贝与Shared Memory Bank冲突

异步拷贝(Async Copy):在Ampere及更新架构中,CUDA引入了cp.async指令,允许线程在等待数据从全局内存传输到共享内存的同时继续执行其他不依赖该数据的指令。这比传统的通过寄存器中转的加载方式更高效,能进一步实现计算与数据迁移的重叠。编程上可以通过__pipeline接口或PTX内联汇编来使用。

Shared Memory Bank冲突的避免:共享内存Bank冲突是性能隐形杀手。例如,如果多个线程访问同一个bank的不同32-bit字,这些访问会被串行化。在GEMM中,当多个线程按行读取共享内存中的矩阵Bs(假设Bs[BK][BN],按行存储)时,如果BN是bank数量的倍数(如32),且线程束中的线程索引threadIdx.x连续,那么它们访问的Bs[ki][threadIdx.x]很可能落在同一个bank,导致冲突。

解决方案

  • 添加填充(Padding):将共享内存数组声明为__shared__ float Bs[BK][BN + 1];。这样,同一列中相邻行的元素在内存地址上相差(BN+1)*sizeof(float)字节,而不是BN*sizeof(float)。通过精心选择填充大小,可以使线程束的访问分散到不同的bank。
  • 改变数据布局:例如,将Bs存储为[BN][BK](列主序),然后让线程按列读取。但这需要与计算时的访问模式相匹配,可能增加索引计算的复杂性。
  • 调整分块参数:选择BNBK为奇数,或不是bank数量(32)的约数,可以减少规律性冲突的概率。

5.3 占用率(Occupancy)与资源平衡

占用率是指每个SM上活跃的线程束数量与最大可能支持的线程束数量之比。更高的占用率有助于隐藏延迟(如全局内存访问延迟),但并非总是越高越好。

  • 寄存器限制:每个线程使用的寄存器数量是影响占用率的主要因素。更复杂的算法(如更大的寄存器分块TM x TN)需要更多寄存器存储中间变量,可能导致占用率下降。
  • 共享内存限制:每个线程块使用的共享内存总量也影响SM上能同时驻留的线程块数量。
  • 平衡策略:目标是最大化“吞吐量”,而非单纯最大化占用率。有时,降低一点占用率以换取更大的寄存器分块(从而增加计算强度和指令级并行)反而能获得更高的整体性能。这需要通过性能分析工具(如Nsight Compute)进行实测和权衡。

6. 实战中的调试、性能剖析与边界处理

优化过程中,正确性和性能验证至关重要。

6.1 内核正确性验证与数值精度

  1. 单元测试:使用小规模随机数据,与一个经过验证的参考实现(如CPU上的朴素算法或CuBLAS)进行逐元素对比。考虑到浮点运算的非结合性,GPU并行累加的顺序可能与CPU不同,导致细微差异。通常使用相对误差|gpu - cpu| / |cpu|进行检查,对于单精度浮点数,误差在1e-5量级通常可以接受。
    bool validate(float* gpu_C, float* cpu_C, int size) { float eps = 1e-5; for (int i = 0; i < size; ++i) { if (fabs(gpu_C[i] - cpu_C[i]) > eps * fabs(cpu_C[i])) { printf("Mismatch at %d: GPU %f, CPU %f\n", i, gpu_C[i], cpu_C[i]); return false; } } return true; }
  2. 边界条件:确保内核正确处理非整除的矩阵维度。在加载数据到共享内存和写回结果时,必须进行严格的边界检查,防止越界访问。未初始化的共享内存或越界访问可能导致不可预知的结果或程序崩溃。

6.2 性能剖析工具的使用

  1. nvprof / Nsight Systems:用于分析应用程序的整体时间线,了解内核执行时间、内存拷贝时间、API调用开销等。可以快速定位是哪个内核或操作是性能瓶颈。
  2. Nsight Compute:这是深入分析CUDA内核性能的利器。它可以提供:
    • 占用率:实际 vs 理论最大值。
    • 内存吞吐量:全局内存、共享内存、L1/L2缓存的读写吞吐量,与硬件峰值的对比。
    • 指令统计:各种类型指令的发射数量、效率。
    • 延迟分析:查看流水线停滞的原因(如内存依赖、执行依赖、同步等待)。
    • Shared Memory Bank Conflicts:直接报告发生的Bank冲突次数。
    • DRAM效率:评估全局内存访问的合并程度。 通过Nsight Compute,你可以精确地知道你的内核在哪个环节没有达到硬件极限,从而进行针对性优化。

6.3 常见性能问题与排查链路

当性能不如预期时,可以按以下链路排查:

  1. 检查计算强度:用Nsight Compute查看Achieved FLOP/s和DRAM Throughput。如果计算强度很低(FLOP/s远低于峰值,而DRAM吞吐量接近峰值),说明瓶颈在内存访问。优化方向是增加数据复用(更大的分块、寄存器分块)。
  2. 检查占用率:如果占用率很低(例如低于30%),可能是寄存器或共享内存使用过多限制了活跃线程块数量。尝试减少寄存器分块大小(TM, TN)或共享内存分块大小(BM, BN, BK),或者调整编译选项(如-maxrregcount)。
  3. 检查内存访问模式
    • 全局内存:查看DRAM效率。效率低通常意味着非合并访问。检查从全局内存加载数据到共享内存的代码,确保一个Warp内的线程访问连续地址。
    • 共享内存:查看Bank Conflict指标。如果冲突严重,尝试添加填充或调整数据布局。
  4. 检查指令效率:查看指令发射效率。效率低可能意味着存在大量的分支分歧(Divergent Branch)或内存依赖停滞。确保内核中的控制流尽可能简单,避免线程束内的分支。
  5. 验证Tensor Core使用:如果目标硬件支持Tensor Core,但性能未达到预期,使用Nsight Compute检查是否成功发射了Tensor Core指令(如HMMA,IMMA)。确保数据精度、矩阵形状和对齐方式符合要求。

6.4 一个完整的优化步骤示例(以FP32 GEMM为例)

假设我们从最基础的Naive内核开始,目标是优化在Ampere架构GPU上的性能。

  1. 步骤1:实现并验证基础分块版本。使用共享内存,实现BM=BN=128, BK=8的分块。确保功能正确,获得一个稳定的性能基线(例如200 GFLOPS)。
  2. 步骤2:引入寄存器分块。让每个线程计算一个2x2的小块(TM=TN=2)。调整线程块大小。性能预期提升(例如到500 GFLOPS)。
  3. 步骤3:优化共享内存访问
    • 分析Bank Conflict:使用Nsight Compute。如果冲突高,为共享内存数组AsBs添加填充,例如As[BM][BK+1]
    • 优化全局内存加载:确保加载AB到共享内存时,线程束访问是合并的。可能需要让线程加载多个元素,或调整加载的维度映射。
  4. 步骤4:应用循环展开和向量化
    • 对内部的BK循环使用#pragma unroll
    • 如果条件允许(地址对齐,连续访问),尝试使用float4进行全局内存加载/存储。
  5. 步骤5:尝试双缓冲。实现共享内存双缓冲,观察是否能进一步隐藏延迟。
  6. 步骤6:调整参数。尝试不同的BM, BN, BK, TM, TN组合。这是一个参数搜索过程,可以使用脚本自动化测试。注意平衡寄存器使用、共享内存使用和占用率。
  7. 步骤7(进阶):转向WMMA/Tensor Core。如果硬件支持且追求极致性能,重写内核使用WMMA API。这需要完全不同的数据加载、存储和计算流程。性能可能跃升至数TFLOPS甚至更高。
  8. 步骤8:处理非标准情况。优化代码以处理非方阵、非2的幂次方尺寸、以及矩阵转置(A^T * B,A * B^T)等情况。这通常涉及更复杂的索引计算和边界处理。

整个优化过程是迭代和实验性的,需要结合性能分析工具的数据进行理性决策,而不是盲目尝试。每一次改动后,都必须验证计算的正确性。最终,一个高度优化的GEMM内核性能可以达到硬件峰值性能的70%甚至更高,这是一个非常了不起的成就。这个过程深刻体现了对GPU计算层次结构的理解:从全局内存到共享内存再到寄存器,从线程到线程束再到线程块,通过精细的数据搬运和计算调度,最终让强大的计算单元持续饱和地工作。