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其中A是M x K矩阵B是K x N矩阵C是M 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);。即一个线程块包含16x16256个线程整个网格由足够覆盖输出矩阵C所有元素的线程块组成。为什么它这么慢核心问题在于内存访问模式。全局内存的合并访问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.x32。但我们的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是偏移。对于同一个Warpk*N相同col是连续变化的假设线程在col方向连续。这理论上可以形成合并访问。但是这要求N不是某些特定的值比如非常大的质数并且内存地址对齐良好。在简单情况下这可能是Baseline中唯一稍微“高效”一点的访问但依然不理想。极高的全局内存访问与计算比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%的时间都在等待数据从显存中读取计算单元几乎在“空转”。没有利用任何高速缓存Baseline内核反复从全局内存读取A和B的数据。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:blockRowBM, k:kBK]加载到共享内存As中 // 每个线程加载一个或多个元素。这里假设线程数 BM*BK每个线程加载一个。 if (threadRow BM (k threadCol) K) { // threadCol在这里充当BK维度索引 As[threadRow][threadCol] A[(blockRow threadRow) * K (k threadCol)]; } // 协作加载将全局内存中B[k:kBK, blockCol:blockColBN]加载到共享内存Bs中 // 注意B的索引计算行索引是kthreadRow列索引是blockColthreadCol 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 性能提升原理与参数选择性能提升的关键数据复用对于C子块中的每个元素原来需要从全局内存访问A的BM个元素和B的BN个元素各K次。现在A的BM x BK块和B的BK x BN块被加载到共享内存后在计算C子块时A的每一行被复用了BN次B的每一列被复用了BM次。这极大地降低了对全局内存的访问需求。访问模式优化在从全局内存向共享内存加载数据时我们可以通过精心设计线程的加载任务确保对全局内存A和B的访问是合并的。例如让一个Warp内的线程连续读取A矩阵的一小段行连续数据或B矩阵的一小段列连续数据可能涉及转置存储以优化访问。计算强度提升数据从全局内存加载到共享内存后后续的BK次乘加运算都发生在高速的共享内存上。计算强度提升为大约(2*BK次FLOP) / (从全局内存加载2*BM*BK2*BK*BN字节)。当BM和BN较大时分母中的加载次数被平摊计算强度显著增加。参数选择经验BM,BN,BK的选择受限于共享内存容量。每个线程块需要的共享内存大小为BM*BK BK*BN以元素计。例如若BMBN128,BK8使用float则共享内存需求为(128*8 8*128) * 4字节 8192字节 8KB。这通常可以接受。选择BK时需要考虑共享内存bank冲突。共享内存被组织成多个bank例如32个。如果同一个Warp内的多个线程访问同一个bank的不同地址就会发生bank conflict导致串行化访问降低性能。因此在存储As和Bs时有时会故意增加一个填充padding维度来错开访问避免bank conflict。BM和BN的大小也影响寄存器使用和占用率Occupancy即每个SM上活跃的线程块/线程数。更大的块需要更多寄存器来存储中间累加值Csub可能降低占用率但提升了数据复用。需要在两者间取得平衡。经过这一轮优化性能通常能有数量级的提升。但距离硬件峰值还有很大差距。4. 优化第二战寄存器优化、双缓冲与指令级并行在利用了共享内存之后下一个瓶颈往往出现在寄存器使用、内存访问延迟隐藏和指令流水线上。4.1 寄存器分块让每个线程计算一个小矩阵在基础的分块内核中每个线程只负责输出矩阵C中的一个标量元素。这意味着每个线程只使用很少的寄存器主要就是一个累加器Csub。现代GPU拥有大量的寄存器每个线程多达255个我们可以利用这些寄存器让每个线程计算一个小的TM x TN的输出块。这被称为寄存器分块Register Tiling。这样做的好处进一步增加数据复用线程从共享内存中加载一小块ATM x BK和一小块BBK x TN到寄存器中然后用它们计算TM x TN个结果。这比每次只计算一个点复用数据的程度更高。减少共享内存的访问频率原来每计算一个输出点需要从共享内存读取BK次A和B的元素。现在为了计算TM x TN个点只需要从共享内存加载TM行A和TN列B的数据各一次在BK循环内然后在寄存器中进行所有组合的乘加运算。这显著降低了共享内存的带宽压力。隐藏访存延迟当线程有多个独立的乘加运算TM*TN个可以执行时GPU的指令调度器可以更好地在等待一次共享内存加载数据的同时执行其他不依赖该数据的计算指令从而隐藏延迟。实现上线程块的大小会调整为(BN/TN, BM/TM)每个线程在寄存器中声明一个float Creg[TM][TN]的数组。在BK循环内部线程先协作将共享内存中需要的A和B的数据块加载到寄存器变量中例如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 指令级优化循环展开与向量化内存访问编译器优化可以帮助我们但显式地给出提示通常效果更好。循环展开Loop Unrolling对于内部的BK循环或寄存器分块内的循环如果BK是编译时常量我们可以使用#pragma unroll指令或手动展开。这减少了循环开销分支预测、递增、比较增加了指令级并行ILP的机会让编译器能更好地调度指令。#pragma unroll for (int ki 0; ki BK; ki) { Csub As[threadRow][ki] * Bs[ki][threadCol]; }向量化内存访问GPU支持一次加载或存储多个数据如float2, float4。使用向量化加载可以减少指令数量提高内存吞吐。例如在从全局内存加载到共享内存时如果地址对齐且连续可以让每个线程一次加载一个float44个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内核结构大致如下声明Warp级别的累加器矩阵片段fragment。在外循环中Warp协作从全局内存通过共享内存加载矩阵A和B的片段到寄存器。调用wmma::mma_sync执行矩阵乘积累加。循环结束后将结果片段存储回全局内存。注意事项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];。这样同一列中相邻行的元素在内存地址上相差(BN1)*sizeof(float)字节而不是BN*sizeof(float)。通过精心选择填充大小可以使线程束的访问分散到不同的bank。改变数据布局例如将Bs存储为[BN][BK]列主序然后让线程按列读取。但这需要与计算时的访问模式相匹配可能增加索引计算的复杂性。调整分块参数选择BN和BK为奇数或不是bank数量32的约数可以减少规律性冲突的概率。5.3 占用率Occupancy与资源平衡占用率是指每个SM上活跃的线程束数量与最大可能支持的线程束数量之比。更高的占用率有助于隐藏延迟如全局内存访问延迟但并非总是越高越好。寄存器限制每个线程使用的寄存器数量是影响占用率的主要因素。更复杂的算法如更大的寄存器分块TM x TN需要更多寄存器存储中间变量可能导致占用率下降。共享内存限制每个线程块使用的共享内存总量也影响SM上能同时驻留的线程块数量。平衡策略目标是最大化“吞吐量”而非单纯最大化占用率。有时降低一点占用率以换取更大的寄存器分块从而增加计算强度和指令级并行反而能获得更高的整体性能。这需要通过性能分析工具如Nsight Compute进行实测和权衡。6. 实战中的调试、性能剖析与边界处理优化过程中正确性和性能验证至关重要。6.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; }边界条件确保内核正确处理非整除的矩阵维度。在加载数据到共享内存和写回结果时必须进行严格的边界检查防止越界访问。未初始化的共享内存或越界访问可能导致不可预知的结果或程序崩溃。6.2 性能剖析工具的使用nvprof / Nsight Systems用于分析应用程序的整体时间线了解内核执行时间、内存拷贝时间、API调用开销等。可以快速定位是哪个内核或操作是性能瓶颈。Nsight Compute这是深入分析CUDA内核性能的利器。它可以提供占用率实际 vs 理论最大值。内存吞吐量全局内存、共享内存、L1/L2缓存的读写吞吐量与硬件峰值的对比。指令统计各种类型指令的发射数量、效率。延迟分析查看流水线停滞的原因如内存依赖、执行依赖、同步等待。Shared Memory Bank Conflicts直接报告发生的Bank冲突次数。DRAM效率评估全局内存访问的合并程度。 通过Nsight Compute你可以精确地知道你的内核在哪个环节没有达到硬件极限从而进行针对性优化。6.3 常见性能问题与排查链路当性能不如预期时可以按以下链路排查检查计算强度用Nsight Compute查看Achieved FLOP/s和DRAM Throughput。如果计算强度很低FLOP/s远低于峰值而DRAM吞吐量接近峰值说明瓶颈在内存访问。优化方向是增加数据复用更大的分块、寄存器分块。检查占用率如果占用率很低例如低于30%可能是寄存器或共享内存使用过多限制了活跃线程块数量。尝试减少寄存器分块大小TM, TN或共享内存分块大小BM, BN, BK或者调整编译选项如-maxrregcount。检查内存访问模式全局内存查看DRAM效率。效率低通常意味着非合并访问。检查从全局内存加载数据到共享内存的代码确保一个Warp内的线程访问连续地址。共享内存查看Bank Conflict指标。如果冲突严重尝试添加填充或调整数据布局。检查指令效率查看指令发射效率。效率低可能意味着存在大量的分支分歧Divergent Branch或内存依赖停滞。确保内核中的控制流尽可能简单避免线程束内的分支。验证Tensor Core使用如果目标硬件支持Tensor Core但性能未达到预期使用Nsight Compute检查是否成功发射了Tensor Core指令如HMMA,IMMA。确保数据精度、矩阵形状和对齐方式符合要求。6.4 一个完整的优化步骤示例以FP32 GEMM为例假设我们从最基础的Naive内核开始目标是优化在Ampere架构GPU上的性能。步骤1实现并验证基础分块版本。使用共享内存实现BMBN128, BK8的分块。确保功能正确获得一个稳定的性能基线例如200 GFLOPS。步骤2引入寄存器分块。让每个线程计算一个2x2的小块TMTN2。调整线程块大小。性能预期提升例如到500 GFLOPS。步骤3优化共享内存访问。分析Bank Conflict使用Nsight Compute。如果冲突高为共享内存数组As和Bs添加填充例如As[BM][BK1]。优化全局内存加载确保加载A和B到共享内存时线程束访问是合并的。可能需要让线程加载多个元素或调整加载的维度映射。步骤4应用循环展开和向量化。对内部的BK循环使用#pragma unroll。如果条件允许地址对齐连续访问尝试使用float4进行全局内存加载/存储。步骤5尝试双缓冲。实现共享内存双缓冲观察是否能进一步隐藏延迟。步骤6调整参数。尝试不同的BM, BN, BK, TM, TN组合。这是一个参数搜索过程可以使用脚本自动化测试。注意平衡寄存器使用、共享内存使用和占用率。步骤7进阶转向WMMA/Tensor Core。如果硬件支持且追求极致性能重写内核使用WMMA API。这需要完全不同的数据加载、存储和计算流程。性能可能跃升至数TFLOPS甚至更高。步骤8处理非标准情况。优化代码以处理非方阵、非2的幂次方尺寸、以及矩阵转置A^T * B,A * B^T等情况。这通常涉及更复杂的索引计算和边界处理。整个优化过程是迭代和实验性的需要结合性能分析工具的数据进行理性决策而不是盲目尝试。每一次改动后都必须验证计算的正确性。最终一个高度优化的GEMM内核性能可以达到硬件峰值性能的70%甚至更高这是一个非常了不起的成就。这个过程深刻体现了对GPU计算层次结构的理解从全局内存到共享内存再到寄存器从线程到线程束再到线程块通过精细的数据搬运和计算调度最终让强大的计算单元持续饱和地工作。