这是“AI Infra 每日一问”系列的第 6 篇。今天想聊一个特别底层、但只要碰过 CUDA kernel 优化就绕不开的问题寄存器 tiling 在不同架构上到底长什么样这里的“寄存器”不是嵌入式里那种 PHY 寄存器地址而是 CPU、GPU、AI 加速器上离计算单元最近的高速暂存空间“tiling”也不是图片处理里的瓦片分割而是把循环切块后让数据尽量留在寄存器里重复使用。如果你写过 GEMM、看过 Attention 的 kernel或者翻过 Cutlass/Triton 生成的代码一定见过一堆 16×8、64×64、128×128 之类的数字。这些数字往小了看是“一块数据多大”往大了看是某一种硬件对数据复用方式给出的硬约束。这篇不是教科书我会从“为什么需要它”开始分别拆一拆 NVIDIA、AMD、CPU 以及脉动阵列上的寄存器 tiling每个场景尽量落到“形状怎么来的、代码长什么样、踩过什么坑”上。适合谁读正在学算子优化的人、想读懂 Cutlass/Triton 代码的人以及被 occupancy 和 register usage 绕晕的调参选手应该都能从里面找到自己踩过的坑。1. 从访存说起寄存器 tiling 到底在优化什么1.1 计算与访存的“剪刀差”逼出了 tiling先看最朴素的 GEMM 写法三层循环里每条语句都从全局内存拿 A[i][k] 和 B[k][j]对每个输出元素做一次乘加。这个时候数据是“一次性”的——读进来算完就丢下一个输出继续重新读。把搬运和计算分开看问题就来了。假设每个 FP32 计算需要读 4 字节数据一台 FP32 算力 19.5 TFLOPS 的 A100需要大约 78 TB/s 的片外带宽才能喂饱。而实际 HBM 带宽只有 2 TB/s 左右差了一个数量级以上。这还只是每次只读一个操作数的情况如果考虑读两个操作数缺口会更大。这就是大家常说的“内存墙”。tiling 的思路就是打破“每条数据只用一次”的局面把循环顺序换一下让一个数据从 HBM 搬到寄存器之后被多条指令反复使用。所以 tiling 的本质不是切块而是通过切块把“访存次数 / 计算次数”这个比例压下来。寄存器是硬件里最快的存储但容量极小。一个线程手头上能用的寄存器数量通常只有几十到几百个。如何在这么小的空间里让每个数据被尽量多次地复用其实就是寄存器 tiling 的核心命题。一个很直观的生活类比把 HBM 想成大仓库shared memory 是操作台寄存器就是你手里能捧住的食材。每次从仓库往操作台搬一大块肉再从操作台往手里抓一小把手里的量决定你一次能同时处理几道菜。1.2 寄存器 tiling 和 shared memory tiling 的分工很多时候大家把“分块”统称为 tiling但寄存器 tiling 是更靠里的一层。shared memory tiling 解决的是线程间数据共享一个线程块把一块数据从全局内存搬进 shared memory块内所有线程都能复用。寄存器 tiling 解决的是线程内数据复用一个线程的寄存器里同时保有多少个输出累加器、多少个 A/B fragment以及在 K 方向上展开多长。两者通常是配对的。一个典型的 GEMM kernel会把整块 tile 从 HBM 搬进 shared memory然后再从 shared memory 搬进每个线程自己的寄存器。shared memory 决定“一个 block 一次吃多大”寄存器 tiling 决定“一个线程一次能攒多少活”。如果把整块 tile 比作一桌菜shared memory 就是那个可以反复取菜的转盘寄存器就是你碗里同时堆着的几道菜。转盘上的菜再多碗里放不下一次也只能吃那么几口。所以寄存器 tiling 的尺寸并不只由“数据有多大”决定而是由“每个线程手里能捧多少”决定。这也引出了后面最关键的问题不同架构能“捧”的量差别很大。2. NVIDIA、AMD、CPU 与加速器tile 的“形状”由谁决定2.1 NVIDIA GPU从外积寄存器 tile 到 Tensor Core 指令形状在 Volta 之前的特斯拉架构上没有 Tensor Core寄存器 tiling 靠的就是 FMA 指令做外积。一个线程通常在寄存器里维护一个 4×4 甚至 8×4 的 C 累加块循环内部每次取一个 a 片段和一个 b 片段用外层积的方式同时更新所有累加器。这样做的好处是同一个 a[i] 可以被同时用于计算多个不同的输出b 也一样数据复用率一下子提上来了。Volta 之后有了 MMA 指令事情变得更有意思。以 Ampere 上最常见的mma.m16n8k16为例一个 warp 的 32 个线程共同完成一个 16×8 的输出块。注意这里的“16×8”不是拍脑袋定的它跟线程数、寄存器宽度、硬件里的 FMA 阵列布局紧密相关。每个线程的 C fragment 是 4 个 fp32 累加器A 和 B 的 fragment 按固定模式散布在 32 个线程里。在实际写 kernel 的时候你不需要手动去摆每个线程的 A/B fragment编译器会帮你生成对应的寄存器排布。但这不意味着你可以完全不管形状。如果你把每线程的 C tile 设成 4×4大概符合 m16n8k16 的粒度如果设成 8×8编译器很可能要做更多寄存器重排或者退化成多次 mma 的拼接。Hopper 上的 wgmma 更进一步由 4 个 warp 组成的 warpgroup 一起发射输入形状可以更大寄存器数据的排布由硬件描述符直接指定灵活性更小但效率更高。这里有个很关键的点寄存器 tiling 的形状本质上是“线程模型 指令集”共同决定的。你不能脱离硬件约束去谈“最优 tile 大小”。同一个 128×128 的 block tile在 A100 上和 H100 上最佳线程配置可能完全不同。2.2 AMD CDNAwavefront 64 带来的形状变化AMD 的调度单位是 wavefront一个 wavefront 是 64 个线程比 NVIDIA 的 warp 大一倍。这个看似简单的差别直接影响了寄存器 tiling 的形状。CDNA 架构上的 MFMA 指令常见形状是 16×16×16 和 32×8×16正好由 64 个线程一起协作。以 16×16×16 为例输出是 16×16 的矩阵64 个线程每线程拿到 4 个 fp32 累加器这和 NVIDIA m16n8k16 的每线程 4 个累加器很接近但因为线程总数不同fragment 在 lane 上的分布方式完全不一样。如果你只写过 CUDA第一次看 AMD 的汇编时特别容易懵明明都是做矩阵乘为什么每线程持有的 A 片段、B 片段跟 N 卡对不上这就是架构差异最真实的一面。做算子库的人比如 ROCm 的 rocBLAS 或者 Triton 的 AMD 后端会专门维护两套 fragment 排布逻辑原因就在这里。好的一面是像 Triton 这类编译器已经把这些差异封装掉了你写一遍计算逻辑它自动帮你映射到不同硬件。但如果你在做底层性能调优理解“wavefront 64 导致 tile 更宽”这一点能帮你少走很多弯路。2.3 CPU向量寄存器与 AMX 的整块 tile 寄存器CPU 这边的情况更“复古”一些但也更直观。AVX-512 一个向量寄存器能装 512 bit也就是 16 个 fp32。CPU 上的寄存器 tiling通常表现为把内层循环 unroll 8 或 16 次同时维护 8~16 个 zmm 向量寄存器每个向量装的是同一行不同列的结果。没有 Tensor Core全靠 FMA 指令一个一个算靠的就是向量寄存器之间的数据复用。Intel AMX 发布之后x86 CPU 也有了真正的“tile 寄存器”。8 个 1KB 的 tile 寄存器一条指令可以做 16×16 甚至更大的矩阵乘加。这个思路已经非常接近 GPU 的 wgmma 了寄存器本身就是按矩阵形状组织的编译器负责把 tile 映射到物理寄存器文件。不过 AMX 的使用限制更多比如数据摆放要和 tile 寄存器形状严格对齐同时 CPU 的访存延迟比 GPU 高很多寄存器 tiling 的收益往往被主存带宽卡住。下面这张表能帮你快速建立不同架构的尺寸感架构典型指令/机制线程/协作单位典型单次累加块每线程 C 累加器数NVIDIA Amperemma.m16n8k1632 线程 warp16×84 个 fp32NVIDIA Hopperwgmma128 线程 warpgroup64×N由描述符指定AMD CDNAMFMA 16×16×1664 线程 wavefront16×164 个 fp32x86 AVX-512FMA 展开单线程16 个输出取决于 unrollIntel AMXTMMULT单线程16×16由 tile 寄存器布局决定2.4 TPU/脉动阵列把 tiling 熔进硬件的极端形态说到极端就必须提脉动阵列比如 TPU 的 systolic array。在脉动阵列里没有我们习惯的“寄存器文件 指令”每个 PE 里的暂存器就是整条数据通路的一部分。矩阵乘法时A 的数据从左边流进阵列B 的数据从上面流进部分和沿对角线一路累积。在这个模型里寄存器 tiling 的形状直接由阵列大小决定。硬件是 128×128 的阵列单次就能累加 128×128 的一块编译器能做的只是决定数据流进阵列的次序和分块粒度。作为软件开发者你没法像在 GPU 上那样自由决定“每个线程手里拿 4×4 还是 8×8”。好处是数据几乎零搬运开销坏处是灵活性很低一旦模型尺寸和硬件阵列不匹配就需要 padding 或者拆分。对比到这里你应该能感受到一个光谱CPU 最自由寄存器 tiling 靠手动 unrollGPU 介于中间指令给了一个硬形状但线程内还有选择空间TPU 最受限tiling 已经变成了物理形状。理解自己在光谱的哪一端才能合理地设计 kernel。3. 手把手设计一个 GEMM 的寄存器 tiling参数怎么定3.1 从 naive 循环到块状外积先换循环次序用一个具体的例子来走一遍。假设我们要算 C[M][N] A[M][K] * B[K][N]naive 版本是三层循环for i in range(M): for j in range(N): sum 0 for k in range(K): sum A[i][k] * B[k][j] C[i][j] sum这里每个内层迭代只做一个乘加A[i][k] 读一次、B[k][j] 读一次数据完全不复用。要把数据复用提上来第一步是让每个线程同时算多个输出。改成这样for i_tile in range(M / MT): for j_tile in range(N / NT): for k in range(K): load a_frag[MT] A[i_tile][k] into registers load b_frag[NT] B[k][j_tile] into registers for i in range(MT): for j in range(NT): C[i][j] a_frag[i] * b_frag[j]这个过程就是把“每次只算一个输出”变成“每次同时算 MT×NT 个输出”。a_frag 和 b_frag 被复用 MT 次和 NT 次访存次数从 MT*NT 次降到了 MTNT 次。这就是寄存器 tiling 的外积形态也是所有优化版本的基础。落实到代码上一个线程内部的累加器往往长这样float c[4][4] {}; // 16 个寄存器4x4 的 C 累加块 for (int k 0; k BK; k) { float a_reg[4] load_a_from_shared(thread_row, k); float b_reg[4] load_b_from_shared(k, thread_col); #pragma unroll for (int i 0; i 4; i) { for (int j 0; j 4; j) { c[i][j] a_reg[i] * b_reg[j]; } } }这里c[4][4]就是寄存器 tiling 的具体体现。16 个 fp32 累加器每个外层 k 的迭代里用 4 个 a 和 4 个 b 同时更新全部 16 个输出。如果你只是用共享内存分块却没有在寄存器里同时维护多个累加器性能会停在很低的水平。因为 shared memory 虽然比 HBM 快很多但和寄存器还是有数量级的差距。只有把真正频繁更新的数据放到寄存器里才能把计算流水线压满。3.2 用“每线程寄存器预算”逆向推导 tile 大小现在关键问题来了寄存器 tiling 应该取多大这里不能凭感觉必须算账。以 NVIDIA GPU 为例每线程最多可用 255 个寄存器每 SM 有 64K 个 32 位寄存器。假设我们想用一个 block 的 256 个线程8 warps处理一个 128×128 的输出块每个线程平均要负责 128×128 / 256 64 个输出。如果排成 8×8需要 64 个 fp32 累加器也就是 64 个寄存器。看起来 255 个寄存器足够放下但账不能只算 C。A/B 的 fragment、shared memory 的地址指针、循环索引、常驻变量这些都要占寄存器。一个 8×8 的 C 已经占了 64 个寄存器再加上 A fragment比如 8 个半精度或 8 个 fp32、B fragment、指针和索引很容易到 100 个以上。如果编译器为了循环展开再预留一些中间变量突破 128 是分分钟的事。一旦超过 255硬件只能把多余的数据 spill 到 local memorylocal memory 实际是落在 HBM 上的性能瞬间崩盘。所以生产级 kernel 往往把 C 控制在 16~32 个寄存器对应每线程 4×4 或 8×4 的输出块。下面这张表是个参考每线程 C tileC 累加器占用单次 k 迭代访问 A/B 次数数据复用比风险4×416 个寄存器8 次4低8×432 个寄存器12 次约 5.3中8×864 个寄存器16 次8高容易溢出权衡的核心是C tile 越大复用比越高但寄存器压力也越大同时指令级并行度不一定能跟上。很多时候你把 C tile 从 8×8 降到 8×4性能反而更高因为 spill 消失了每线程 occupancy 也上来了。3.3 从伪代码到 Tensor Core代码不是手摆寄存器上面的手动外积版本在 Tensor Core 出来之后已经显得有点“古典”了。如果你在 Ampere 之后的 GPU 上写 GEMM真正高效的做法是直接发 MMA 指令让硬件来算 16×8 或更大的一块而不是自己在寄存器里循环 FMA。但明白手动外积的过程依然很重要。因为无论是mma还是wgmma它们本质上都在做同一件事把若干线程的寄存器片段组成一个大矩阵乘法的输入和输出。只是数据排布的职责从你手里移交给了编译器和硬件描述符。用 CUDA 的 WMMA API 时你会写类似wmma::fragmentmatrix_a, 16, 16, 16, half, row_major这样的类型这个 fragment 在语义上就是“分布在 warp 内所有线程寄存器里的一块 A 矩阵”。你不需要手动告诉每个线程做哪几个乘加硬件结构已经固定了。对于做算子库的人这时候真正需要关心的是两个问题一个是 tile 的几何形状能不能和指令的固定形状对齐另一个是 fragment 在寄存器里的排布会不会导致额外的 move 指令。例如用mma.m16n8k16时C tile 的粒度是每个线程 4 个 fp32A/B 的 fragment 需要在 warp 内做固定排列。如果你上层用了一个让编译器很难生成高效排布的 tile 形状汇编里就会出现一大堆mov指令白白浪费吞吐。所以我的建议是先用手动外积理解原理再用 Tensor Core 指令拿性能。不要跳过原理直接抄 WMMA 代码否则遇到 fragment 排布问题你根本不知道编译器在对你的 kernel 做什么操作。4. 寄存器 tiling 的实战排查与常见坑4.1 寄存器溢出你以为在读寄存器其实在读 local memory我在实际调 kernel 时遇到最多的坑就是寄存器溢出。典型症状代码感觉写得很规整但性能比预期低 20% 到 50%甚至更多。打开 Nsight Compute 一看local memory 的访问次数吓人寄存器使用量已经接近 255甚至编译日志里直接出现spill stores和spill loads。寄存器溢出的本质是每个线程想用的寄存器超过了硬件极限编译器只能把一部分变量挪到 local memory。而 local memory 在物理上就是 HBM读写它意味着每个线程都在偷偷访问全局内存。这时候你的寄存器 tiling 再漂亮也等于白搭因为数据反复在 HBM 和寄存器之间倒腾。排查方法NVIDIA 这边最简单的是编译时加--ptxas-options-v直接看寄存器用量和 spill 情况。如果是运行期用 Nsight Compute 看Register Usage和Local Memory两个指标。一旦确认 spill优先按这个顺序调把每线程 C tile 变小比如从 8×8 降到 8×4、减少内层循环的 unroll 倍数、用__launch_bounds__或maxrregcount128强制编译器收敛寄存器用量。AMD 那边对应看 vgpr 占用用 rocprof 也能看到类似的信息。有一种情况很容易被忽略不是 C tile 太大而是编译器为了生成向量化访存把 shared memory 里连续的一段数据一次性加载到多个寄存器里导致峰值寄存器用量暴增。这时候你可以手动把共享内存读取拆成几次独立的 load让编译器在每次 load 后尽快使用降低同时存活的寄存器数量。4.2 tile 形状与指令形状不匹配Tensor Core 根本没用到另一个反直觉的坑是你觉得自己用了 WMMA 或 MMA实际上代码编译出来跑的却是一堆 FMATensor Core 一个都没参与。原因通常是 tile 形状和指令要求对不上导致编译器只能退化为标量计算。例如mma.m16n8k16需要 C 累加块至少是 16×8 的粒度。如果你把一个 warp 的线程分工设计成每个线程独立算一个 4×4 的小块而没有在 warp 层面凑出 16×8 的协作关系编译器就很难把这几条独立的计算合并成一条 MMA只能老老实实用 FMA 去算。结果就是Tensor Core 单元空转SM 的普通 FMA 流水线忙死性能只比纯手动外积好一点点。遇到这种情况建议直接在汇编或者 SASS 层面确认 MMA 指令是否出现。Nsight Compute 的指令统计里可以看到mma类指令的数量。如果比预期少很多就要回头检查 tile 形状和线程协作方式是否和指令形状对齐。还有一个隐蔽问题Hopper 的 wgmma 对数据对齐要求更高A 和 B 的起始地址经常要求 16 字节对齐你要是只做 4 字节对齐编译不出错但性能会差一截。4.3 不要忘了 shared memory 这条腿很多初学者把注意力全放在寄存器 tiling 上结果忘了数据要到寄存器必须先经过一层 global、再经过 shared。如果 shared memory 的加载模式乱七八糟每个线程读的地址不连续或者 bank conflict 频繁即便寄存器 tiling 做得完美HBM 带宽照样跑不满。正确做法是把 tiling 当成一个三级流水来设计global→shared→register→计算。每一层都有它自己的“颗粒度”。shared memory 层解决线程间的数据复用寄存器层解决线程内的数据复用两者缺一不可。我个人的调试顺序是这样的先保证 global memory 合并访问用一个最简单的 shared memory tile 建立一个基线用 profiling 工具看当前是 memory-bound 还是 compute-bound如果已经是 memory-bound优化 shared memory 的加载和 bank conflict如果是 compute-bound再逐步增大寄存器 tiling观察寄存器使用量和 spill 变化每改一版只动一个变量对比 profiling 数据不要一次改一堆。这个顺序能帮你快速定位瓶颈到底在哪一层。很多人一开始就上大 C tile发现性能没提升其实是 shared memory 加载阶段早就把带宽占满了寄存器层再怎么优化都是白费力。4.4 常见错误速查表症状可能原因排查手段性能比预期低local memory 访问多寄存器溢出C tile 过大或循环展开过猛看编译日志查 spill缩减 C tilecompute-bound 但 Tensor Core 利用率低tile 形状与 mma/wgmma 形状不匹配退化成了 FMA看 SASS 里 mma 指令数量调整线程协作模式带宽上不去但寄存器 tiling 看着很合理shared memory 层 bank conflict 或非合并访问用 profiler 看 shared memory 吞吐修加载模式改了 tile 大小性能没变化可能根本不是 compute-bound瓶颈在数据搬运先判断 memory-bound 还是 compute-bound 再动手最后分享一个我自己的习惯拿到一个新架构第一件事不是去查理论峰值而是写一个最简单的 64×64 GEMM在汇编里数一数每个线程手上同时维护着多少个累加器。因为寄存器 tiling 的“长相”会直接告诉你这个架构的真实脾气它能塞下多大的 tile它的 fragment 怎么在 lane 间分布它的指令对形状有多挑剔。等踩过几次寄存器溢出的坑之后你对 tile 尺寸的直觉就会变得非常准再看到别人 kernel 里的128×128 block、8×8 per-thread tile脑子里马上能浮现出每个线程的寄存器里大概是什么状况。这种直觉就是调优经验最值钱的部分。