ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

CUDA矩阵乘法优化实战:共享内存tiling与向量化技巧

CUDA矩阵乘法优化实战:共享内存tiling与向量化技巧 矩阵乘法这个题目我在前几篇已经铺垫了不少基础。如果你是从第一篇一路看过来的应该对线程格、线程块、共享内存这些概念不陌生了也能把简单的 kernel 跑起来、用cudaEvent量出耗时。但这一步离“写得好”还差得很远。矩阵乘法是我认为练习 kernel 优化最适合的靶子它计算密集、访存规律、性能天花板明确而且几乎能把 kernel 优化里的经典问题——内存合并、共享内存分块、bank conflict、指令级并行——全部踩一遍。这篇文章是系列第五篇的“中篇”我打算从一个最朴素的乘法 kernel 开始一步步改成 shared memory tiling 版本再把 padding 和向量化加载加进去。整个过程不搞玄学每个改动都能用工具观测到效果。适合已经能写简单 CUDA 程序、但对性能分析没什么感觉的人也适合想把矩阵乘法彻底吃透的算法工程师和竞赛党。1. 先给快速矩阵乘法做个基线1.1 从朴素内核出发优化第一步不是动手改而是先有一个能跑、能测、能复现的基线版本。我这里说的“朴素内核”就是每个线程负责计算输出矩阵 C 里的一个元素内层循环一路累加 K 次__global__ void matmul_naive(const float* A, const float* B, float* C, int M, int N, int K) { 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; for (int k 0; k K; k) { sum A[row * K k] * B[k * N col]; } C[row * N col] sum; } }启动配置直接用二维网格dim3 block(16, 16); dim3 grid((N 15) / 16, (M 15) / 16); matmul_naivegrid, block(d_A, d_B, d_C, M, N, K);这个版本好不好从“能跑”的角度说完全没问题但性能只能用一个字形容——惨。以 1024x1024 的方阵为例我的 RTX 3060 上跑出来大概只有 60~80 GFLOPS而这张卡的 FP32 理论峰值在 12 TFLOPS 以上利用率连 1% 都不到。对比一下cuBLAS 手写一个高性能 kernel 可以跑到 8~10 TFLOPS 左右差距是两三个数量级。问题出在哪不是 GPU 算得慢而是数据喂不过去。1.2 用 Roofline 模型看懂瓶颈这里必须引入一个非常实用的分析工具Roofline 模型。它的核心思想很简单一个 kernel 的性能上限取决于它到底是“算得快”还是“数据搬得快”。GPU 的 FP32 峰值算力很高但全局内存带宽有限如果每个浮点运算都要从全局内存搬好几个字节那么性能就会被带宽卡死。以朴素矩阵乘法为例每个输出元素需要读取 A 的一行K 个 float和 B 的一列K 个 float也就是 2K 次浮点读取外加一次写入。计算量是 2K 次浮点操作。算下来每个浮点操作大约需要 4~6 字节的全局内存搬运算术强度只有 0.2~0.3 FLOP/Byte。而 RTX 3060 的机器平衡点大约是 35 FLOP/Byte 左右。算术强度远低于平衡点所以这个 kernel 从理论上就不可能跑得快瓶颈就是全局内存带宽。你可能想问L2 cache 难道不能帮忙吗矩阵规模小的时候可以但一旦数据超过 L2 容量每个线程重复读同一行、同一列的行为就会把带宽打满。所以矩阵乘法优化的第一步就是提高数据复用率而共享内存就是干这个用的。1.3 性能记录到底该怎么记在开始改代码之前建议先把计时和误差校验的框架搭好不然后面每一步优化都像在黑箱里猜。我一般用cudaEvent计时比clock()和std::chrono都准因为它直接测量 GPU 端的时间戳cudaEvent_t start, stop; cudaEventCreate(start); cudaEventCreate(stop); cudaEventRecord(start); kernelgrid, block(...); cudaEventRecord(stop); cudaEventSynchronize(stop); float ms 0.0f; cudaEventElapsedTime(ms, start, stop);GFLOPS 的计算公式是2 * M * N * K / (time_in_seconds) / 1e9。注意那个因子 2因为一次乘法和一次加法各算一次浮点操作。矩阵乘法是 FMA 密集操作计算量天然要乘 2。还要做一个 CPU 参考实现每次 kernel 跑完都和 CPU 结果对比一下最大绝对误差控制在1e-3量级大矩阵可以放宽到1e-2。我踩过太多次“性能很好看、结果全错”的坑没有校验就继续优化最后排查问题的时候加倍痛苦。2. 共享内存 tiling摆脱全局内存带宽束缚2.1 为什么要用共享内存共享内存是 GPU 上每个 SM 内部的一块高速存储它和 L1 cache 在硬件上是同一块物理存储但最大的区别在于共享内存是程序员显式控制的。你可以把它理解成一个由软件管理的小型“缓存”数据需要你手动搬进去、手动同步、手动释放。虽然麻烦但换来的是极高的带宽和极低的延迟。在矩阵乘法里共享内存的核心价值是数据复用。看一个分块思路把输出矩阵 C 分成若干个小 tile每个线程块负责计算一个 tile。计算这个 tile 时需要 A 的子块和 B 的子块。如果每个子块被读入共享内存一次然后在块内被反复使用那么全局内存的访问量就从 O(N^3) 降到了 O(N^3 / tile_size) 左右。tile 越大复用次数越多对带宽的压力就越小。我记得第一次把朴素 kernel 改成 shared memory 版本时性能提升了 10 倍以上那是在 GPU 上学到的最直观的一课很多时候瓶颈不是算力而是数据搬运。共享内存相当于在全局内存和寄存器之间加了一层“中间仓库”把重复读取转成了片上访问。2.2 tile 大小和 block 维度如何定tile 大小的选择是一个经典的空间和时间权衡。tile 越大数据复用率越高但每个线程块占用的共享内存也越多会限制 SM 上能同时驻留的线程块数量影响占用率occupancy。我一直喜欢用 32x32 的 tile 做教学不是因为它一定是最优解而是因为 32 这个数字正好对齐一个 warp 的大小分析内存访问模式时非常方便。一个 32x32 的 tile 需要 1024 个 floatA 和 B 各一份就是 8KB。RTX 30 系显卡每个 SM 有 128KB 共享内存32x32 的 tile 完全放得下还不会对 occupancy 造成太大压力。对应的线程块维度我通常设置成dim3 block(32, 32)刚好 1024 个线程达到每个线程块的最大线程数上限。虽然块内线程数多但每个线程的寄存器消耗不高实测 occupancy 还是能保持在一个不错的水平。如果你用的是老一些的卡或者想留更多寄存器给后面优化也可以先用 16x16 的 tile 做练习思路完全一样。2.3 边界处理从固定尺寸到通用矩阵做性能优化的时候我建议先把“矩阵尺寸能被 tile 整除”这个条件当成默认假设。这不是偷懒而是为了把核心逻辑讲清楚。实际生产环境当然会遇到任意尺寸的矩阵处理方式是在加载阶段加边界判断超出范围就填 0int row by * TILE_SIZE ty; int col bx * TILE_SIZE tx; if (row M ktile * TILE_SIZE tx K) { As[ty][tx] A[row * K ktile * TILE_SIZE tx]; } else { As[ty][tx] 0.0f; }加载 B 的时候同理。这种填充哨兵值的做法能让计算阶段不用再做任何边界判断内层循环可以完全展开性能损失几乎可以忽略。另外说一下动态共享内存。当你把 tile 大小作为运行时参数传递时可以用extern __shared__ float smem[];配合启动配置的第三个参数指定共享内存大小size_t smem_size 2 * TILE_SIZE * TILE_SIZE * sizeof(float); kernelgrid, block, smem_size(...);动态共享内存更灵活但代码可读性差一点索引计算也要手工偏移。入门阶段先用静态共享内存就好跑通了再去折腾动态版本。3. 动手写一个分块矩阵乘法内核3.1 代码结构与初始化直接上一个可以跑的 shared memory tiling 版本。假设矩阵尺寸都能被TILE_SIZE整除先保证主线逻辑干净#define TILE_SIZE 32 __global__ void matmul_smem(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[TILE_SIZE][TILE_SIZE]; __shared__ float Bs[TILE_SIZE][TILE_SIZE]; int bx blockIdx.x; int by blockIdx.y; int tx threadIdx.x; int ty threadIdx.y; int row by * TILE_SIZE ty; int col bx * TILE_SIZE tx; float acc 0.0f; for (int t 0; t K / TILE_SIZE; t) { // 协同加载 A 和 B 的 tile 到共享内存 As[ty][tx] A[row * K t * TILE_SIZE tx]; Bs[ty][tx] B[(t * TILE_SIZE ty) * N col]; __syncthreads(); // 计算当前 tile 的贡献 #pragma unroll for (int k 0; k TILE_SIZE; k) { acc As[ty][k] * Bs[k][tx]; } __syncthreads(); } C[row * N col] acc; }host 侧调用和朴素版本几乎一样只是把 block 改成dim3 block(TILE_SIZE, TILE_SIZE);grid 对应(N / TILE_SIZE, M / TILE_SIZE)。记得编译时加-O3不然性能差异会误导你。3.2 内核逐步拆解这个 kernel 的关键点有三个。第一个是“加载”。每一轮循环当前线程块负责读取 A 的一个 32x32 子块和 B 的一个 32x32 子块分别存进As和Bs。为什么这里可以直接映射因为线程块内的线程坐标(ty, tx)恰好对应子块内的坐标每个线程读一个元素协作完成整个 tile 的搬运非常规整。第二个是“同步”。__syncthreads()用在两个地方一是加载完成之后确保所有线程都把数据写进了共享内存才能开始读二是计算完成之后确保没有线程在下一轮循环加载时覆盖掉还在被其他线程读取的数据。这里如果漏掉第二个__syncthreads()你会得到一个时对时错的诡异结果而且很难定位。第三个是“计算”。内层循环从As[ty][k]和Bs[k][tx]读取数据累加到acc。每个线程块计算一个 32x32 的输出 tile而As和Bs上的每个元素会被本块内 32 个不同线程复用这就是共享内存带来的数据复用效果。3.3 跑分与第一次对比我在 RTX 3060 上用 1024x1024 的方阵跑了一下结果大概是这样的版本耗时(ms)GFLOPS相对朴素加速比朴素版本~28~771xshared memory tiling~2.3~930约 12x从 77 GFLOPS 到 930 GFLOPS提升非常可观。但你要知道离 RTX 3060 的 12 TFLOPS 还差得远。为什么还有很多提升空间因为当前版本还有两个关键问题一是加载阶段可能存在 bank conflict二是内层计算循环的指令级并行度不足。接下来就一个一个解决。4. 合并访问与 bank conflict 的实战影响4.1 全局内存加载真的合并了吗先看全局内存的加载方式。GPU 执行内存访问时是以 warp 为单位的一个 warp 里的 32 个线程会同时发起一次内存请求如果它们访问的地址在内存中是连续的硬件就能合并成少数几个事务一次搞定这叫“合并访问”。如果地址是跳跃的、互相错开的硬件就不得不拆成多个事务带宽利用率直线下降。在我们这个 kernel 里加载As[ty][tx] A[row * K t * TILE_SIZE tx]一个 warp 内线程的tx是从 0 到 31 连续变化的对应的地址也是连续的所以是完美的合并访问。加载Bs[ty][tx] B[(t * TILE_SIZE ty) * N col]同样warp 内col连续地址连续合并访问也没有问题。所以这个版本在全局内存访问上没什么大毛病。真正需要仔细分析的是共享内存内部的访问冲突也就是 bank conflict。4.2 bank conflict 到底怎么触发共享内存为了高带宽在硬件上被分成了 32 个 bank每个 bank 可以同时服务一个地址的访问。理想情况下一个 warp 的 32 个线程访问 32 个不同 bank一次完成如果多个线程访问同一个 bank就会发生冲突硬件会把这次访问拆成多次串行处理性能成倍下降。在当前的 tiling 版本里计算阶段的内层循环访问As[ty][k]和Bs[k][tx]。由于tx在 warp 内连续Bs[k][tx]的地址也连续落在不同 bank 上没有冲突。As[ty][k]由于同一 warp 内ty基本不变k固定所有线程访问的是同一地址会触发广播机制也没有冲突。那 bank conflict 什么时候出现一个典型场景是如果你在共享内存里以“转置”的方式存放 B也就是Bs[tx][k]这样索引让 warp 内不同线程访问同一列的不同行地址间隔就是一行的大小例如 32 个 float。32 个线程访问到的地址全部落在同一个 bank 上32-way conflict性能直接掉一个数量级。这是很多人在做“共享内存转置”优化时踩过的大坑。为了直观演示你可以试一下把计算循环里的索引改成acc As[ty][k] * Bs[tx][k];然后转置存储 B 的 tile。编译不会报错计算也能算对但性能会严重跳水。这就是 bank conflict 的威力。4.3 padding 修复冲突的两行改动修复 bank conflict 最经典的手段是 padding也就是给共享内存数组的每一行多加一个 float让行与行之间的起始地址错开__shared__ float As[TILE_SIZE][TILE_SIZE 1]; __shared__ float Bs[TILE_SIZE][TILE_SIZE 1];这个改动背后的原理很简单原来每一行是 32 个 float正好占满 32 个 bank。第 0 行的第 0 个元素和第 1 行的第 0 个元素落在同一个 bank 上。加上 1 个 float 的偏移后第 1 行的起始地址就错开了一个 bank不同行的同列元素不再碰撞了。如果你用转置访问模式去测试加 padding 前后性能差异会非常明显。即使在标准 tiling 版本里padding 对某些特殊访问模式也有保护作用而且只多占用了 32x2 个 float 的空间性价比极高。我个人的习惯是只要写共享内存数组一律加 padding。这不是洁癖而是成本几乎为零的防御性编程。5. 向量化加载与循环展开再进一步5.1 先决条件16字节对齐在引入float4之前先强调一个硬性条件地址对齐。float4一次读取 16 字节要求指针地址必须是 16 字节对齐的。用cudaMalloc分配的设备内存起点天然满足对齐要求但真正的坑在于矩阵行起始地址如果K不是 4 的倍数A[row * K ...]这一行的起始位置就可能不对齐。解决办法有两种。第一种干脆让K在分配时向上取整到 4 的倍数在 kernel 里按 padding 后的行宽访问。第二种使用cudaMallocPitch分配带 pitch 的二维数组保证每一行起始地址对齐。我在手写高性能 kernel 时通常选第一种因为逻辑简单而且对计算阶段没有任何额外开销。另一个好习惯是用reinterpret_castconst float4*之前的指针最好断言一下地址确实对齐了否则在部分平台上会出现莫名其妙的访问错误或性能下降。5.2 用 float4 减少加载指令接下来这个版本让每个线程一次计算 4 个输出元素同时用float4加速全局内存到共享内存的搬运。核心思路是block 还是负责一个 32x32 的输出 tile但线程块维度改成(8, 32)共 256 个线程。每个线程的tx取值范围是 0 到 7对应 4 个连续的列索引4 * tx 0到4 * tx 3。加载阶段可以这样写__shared__ float As[TILE_SIZE][TILE_SIZE 1]; __shared__ float Bs[TILE_SIZE][TILE_SIZE 1]; int tx8 threadIdx.x; // 0..7 int ty threadIdx.y; // 0..31 int col4 tx8 * 4; // 每个线程用 float4 一次加载 A 的 4 个连续元素 const float4* Arow reinterpret_castconst float4*( A (by * TILE_SIZE ty) * K bx * TILE_SIZE col4); float4 a4 *Arow; As[ty][col4 0] a4.x; As[ty][col4 1] a4.y; As[ty][col4 2] a4.z; As[ty][col4 3] a4.w;B 的加载类似。计算阶段每个线程维护 4 个累加器float acc0 0.0f, acc1 0.0f, acc2 0.0f, acc3 0.0f; for (int t 0; t K / TILE_SIZE; t) { // 加载 As 和 Bs ... __syncthreads(); #pragma unroll for (int k 0; k TILE_SIZE; k) { float a As[ty][k]; acc0 a * Bs[k][4 * tx8 0]; acc1 a * Bs[k][4 * tx8 1]; acc2 a * Bs[k][4 * tx8 2]; acc3 a * Bs[k][4 * tx8 3]; } __syncthreads(); }注意这里读取Bs[k][4 * tx8 0..3]warp 内tx8取 0..7所以这 4 次读取分别覆盖了 32 个不同列地址是连续的没有 bank conflict。而 4 个累加器同时更新硬件可以把 4 次 FMA 指令并行发射指令级并行度比原来单个累加器高了很多。5.3 多累加器与指令级并行为什么单个累加器不够好因为内层循环每次迭代都要把新读到的a乘上b再加回同一个acc这形成了一个串行的依赖链。GPU 虽然能同时执行大量线程来隐藏延迟但每个线程内部的指令依赖还是会把 FMA 的吞吐往下拉。当你有 4 个独立的累加器时4 个 FMA 之间没有数据依赖编译器可以把它们调度成交错执行大大提升每个线程的指令吞吐。这个优化在 Ampere 之前的架构上效果特别明显因为老的架构没有足够多的执行单元来完全隐藏一个线程内部的长依赖链。实测下来这个“一个线程算 4 个输出”的版本在我机器上已经能跑到 2.5~3 TFLOPS是把朴素版本又提升了 30 多倍。当然每个线程计算 4 个输出只是开始。更激进的版本会让每个线程计算 8x8 甚至 16x16 个输出此时寄存器使用量会非常大需要配合寄存器数组和更大的 tile 设计。这部分内容比较多适合放到下一篇再细讲。6. 性能问题排查与实用技巧6.1 计算结果不对先查什么做矩阵乘法优化最痛苦的不是性能低而是结果错了还找不到原因。我总结了一个固定的排查顺序按这个顺序查基本都能定位。第一步查边界。矩阵尺寸不能被 tile 整除时越界读写是最常见的错误根源。核心里所有对A、B、C的访问都要检查row、col、k的索引范围。第二步查同步。__syncthreads()是不是配对出现的是不是有人在条件分支里用了它赋值循环里是否有线程提前退出导致同步错位第三步查共享内存初始化。如果加载阶段有边界填充确保填充的是 0 而不是随机值。忘记填充会在输出 tile 边缘产生垃圾数据。第四步查浮点误差。矩阵乘法本来就是大量浮点累加计算顺序不同会导致结果略有差异。如果最大误差在1e-2量级基本是正常浮点问题如果差到离谱那就是逻辑 bug。6.2 性能上不去时如何定位不要靠“感觉”判断瓶颈要上工具。NVIDIA 官方的 Nsight Compute 是排查 kernel 性能的首选工具。常用的定位方式是用两个关键指标dram__throughput.avg.pct_of_peak_sustained_elapsed全局内存带宽利用率。sm__throughput.avg.pct_of_peak_sustained_elapsedSM 整体吞吐率包含计算、访存、调度等综合流水线。如果前者接近 90% 以上而后者不高说明 kernel 被内存带宽限制重点优化数据复用和访存模式。如果后者很高而前者很低说明 kernel 是计算限制或者指令限制考虑减少冗余计算、提高 ILP、降低寄存器溢出。还有一个值得看的指标是smsp__average_warp_latency_per_inst_issued它反映 warp 的平均执行延迟。如果这个值很高说明线程大量时间在等待访存结果共享内存 tiling 就是为了解决这类等待。ncu --set full ./matmul完整 profile 的信息量很大第一次跑会被海量指标淹没。我的建议是只看上面两三个指标跑一版优化记录一版对比趋势比对着几百个指标发呆有用得多。6.3 我自己常用的检查清单最后分享一个日常开发用的速查表每当性能没有达到预期我会按顺序过一遍检查项定位思路是否发生了大量 global memory 重复读取用 Roofline 粗算算术强度或用dram__throughput确认带宽瓶颈共享内存加载是否合并检查 warp 内线程访问的地址是否连续是否存在 bank conflict检查共享内存索引模式尝试加 padding循环依赖是否太重增加独立累加器看 ILP 是否提升线程块大小和占用率是否合理用cudaOccupancyMaxPotentialBlockSize粗查理论占用率是否用了-O3调试模式编译的性能数据没有参考价值这是最容易犯的低级错误内存分配是否用了cudaMallocManaged且在 kernel 内被反复访问Managed memory 有时会触发 page fault 和迁移性能测出来会偏慢测试时尽量用cudaMalloc排查性能问题最忌讳一次改多个变量。我每次只改一处重新计时、重新 profile记录到一张表里。优化的过程本质上是在建立“改动 - 性能指标”的因果关系而不是碰运气。在矩阵乘法这条优化路径上我最大的体会是不要指望一次就能写出最优 kernel慢慢来每一步都量化哪怕只提升 5% 也是有用的情报。下一篇我会继续往下走讲寄存器 tiling、更深的循环展开以及如何逼近甚至调用 Tensor Core 的硬件能力。先把这个 shared memory 版本吃透你会发现自己对 CUDA 内存模型的理解上了一个大台阶。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进