寄存器Tiling跨架构实战：从CUDA到AMD、ARM与NPU的差异与调优

每天在 AI Infra 团队里回答一个问题的习惯坚持到第 6 天了。今天这个问题是寄存器 tiling 在不同架构上到底长什么样问出这个问题的人多半已经啃过一些 GEMM 优化相关的文章知道 tiling 能提性能但一换硬件就懵——在 CUDA 上写好的 tile 逻辑拿到 AMD、ARM 甚至自研 NPU 上好像完全不是一回事。这篇文章想把这层窗户纸捅破它省的是什么钱、不同架构下的物理形态差异在哪、以及我给的一些可以直接上手的调参经验。看完你至少能理解一件事为什么同一个优化术语在不同硬件上代码长相可以天差地别。1. 先算一笔账寄存器 tiling 到底在省哪头的钱1.1 一次乘加背后的搬砖开销很多人第一次接触 tiling 时注意力都放在循环改写上觉得把循环切成块儿就是 tiling。这个理解没错但切在哪个存储层级效果完全不一样。寄存器 tiling 切的是最内层那一次循环也就是数据从寄存器到计算单元那一步。先看一个朴素 GEMM 内核里发生了什么。计算 C[m][n]sum_k(A[m][k]*B[k][n])最 naive 的写法是每个线程算一个输出元素内层循环里每个 k 都要读一次 A[m][k] 和 B[k][n]做一次 FMA。看起来一次乘加只搬了 2 个数挺便宜的。但真算起来2 个全局内存 load 对 1 次 FMA这个比例烂到没法看。以 A100 为例FP32 算力大概 19.5 TFLOPSHBM 带宽在 2TB/s 量级。要让算力跑满算术强度至少得 10 FLOP/byte 以上而 naive 内核每个输出元素需要搬 8 字节两个 4 字节浮点换来 2 次浮点运算0.25 FLOP/byte差了四十倍。也就是说这时候 GPU 的实际利用率不是被计算卡住而是被搬砖卡住99% 的时间都在等数据。tiling 解决的就是这个搬运问题。把输出切成一块一块让同一个 A 值、同一个 B 值在寄存器里多待一会儿多参与几次乘加搬运和计算的比例就健康了。1.2 三层仓库各管一段复用现代处理器上的存储层次可以理解成三个递进的仓库全局内存HBM/DDR是大仓库共享内存shared memory / LDS是楼下的中转站寄存器是手边的工作台。tiling 每一层都在做但管的事情不一样。块级 tilingCTA tile把整个大矩阵切成 128×128 这种块塞进共享内存。这一步管的是全局内存到共享内存之间的搬运量。线程级 tilingthread tile在一个线程内部把输出再切成 8×8、4×4 这样的小块数据装进寄存器。这一步管的是共享内存到寄存器之间的搬运量。指令级 tiling矩阵指令Tensor Core、MFMA、DPAS内部硬件自己决定数据怎么在寄存器里摆。这部分你只能顺着硬件的脾气来。寄存器 tiling 最直观的收益可以从数字上看。假设一个线程算 8×864 个输出每个 k 迭代需要往寄存器里装 8 个 A 值和 8 个 B 值共 16 个 load换来 64 次 FMA。每个 FMA 对应的 load 数从 naive 的 2 次降到了 0.25 次共享内存的压力直接降了一个量级。注意这还没算 bank conflict 和指令发射开销实际收益比纸面更大。提示tiling 的本质是用寄存器容量换带宽。寄存器是芯片上最贵的存储资源所以这个兑换比例是有限的——这就是为什么后面所有架构的 tiling 参数本质都是在算寄存器够不够放这笔账。2. NVIDIA 上常见的两副面孔FFMA 与 Tensor Core2.1 经典手写模板每个线程盯住一块 8×8在 NVIDIA 上如果不用 Tensor Core寄存器 tiling 的标准长相就是块级 tiling 把 A、B 的切片准备好放进共享内存然后每个线程从共享内存里取一个小块数据到寄存器算一块输出。最典型的是每线程 8×8 的 FFMA 内核简化代码长这样// 伪代码每线程 8×8 输出块As/Bs 是已经由 block 级 tiling 填好的共享内存切片 __global__ void gemm_kernel(const float* As, const float* Bs, float* Cs, int K, int N) { float c[8][8] {0.f}; // 64 个累加寄存器 float a[8], b[8]; // A/B 分片寄存器 for (int k 0; k K; k) { #pragma unroll for (int i 0; i 8; i) a[i] As[threadIdx.y * 8 i][k]; #pragma unroll for (int j 0; j 8; j) b[j] Bs[k][threadIdx.x * 8 j]; #pragma unroll for (int i 0; i 8; i) for (int j 0; j 8; j) c[i][j] fmaf(a[i], b[j], c[i][j]); } // 写回 C 矩阵 }注意这里的关键是什么是 a[i] 和 b[j] 被反复用了 8 次。一个 a[0] 一次 load跟 8 个不同的 b[j] 做了 8 次 FMA这就是寄存器复用。编译器看到你手动展开的循环后会把 a[0] 一直钉在寄存器里不会重复从共享内存 load。算一下寄存器开销c[][] 64 个寄存器a[] 8 个、b[] 8 个加上指针、索引、循环计数轻松到 90~100 个寄存器。NVIDIA 每个线程上限 255 个寄存器所以这个配置不会溢出但占用量不小。一个 SM 有 64K 个 32 位寄存器256 线程的 block 如果每人用 96 个就要 24K 个寄存器一个 SM 最多塞 2~3 个 block占用率大概 25% 到 40%。这个占用率听着不高但 8×8 的 tile 本身给了足够的指令级并行很多时候反而比高占用率的小 tile 跑得快。性能调优不是越高占用率越好这个观念从一开始就要立住。2.2 Tensor Core 把分块形状焊死进了指令到了 Volta 之后的 GPU情况变了。Tensor Core 的出现让寄存器 tiling从代码模式变成了硬件指令规格。以 Ampere 的mma.m16n8k16为例一个 warp32 线程协同算一个 16×8 的输出块硬件规定好了每个线程手里拿什么A 片段16×16 的 fp16 矩阵分到 32 个线程每个线程 8 个元素占 4 个 32 位寄存器B 片段16×8 的 fp16 矩阵每个线程 4 个元素占 2 个寄存器C 累加器16×8 的 fp32 输出每个线程 4 个 fp32占 4 个寄存器。也就是说一个线程总共也就 10 个左右的寄存器在管数据比 FFMA 路线的 90 多个少得多。这也是为什么 Tensor Core 内核反而能把占用率做高、把延迟藏得更好的原因之一。但代价是你的数据必须按硬件规定的 fragment 布局摆放通常不是朴素的行主序。实际工程里这一步靠ldmatrix指令从共享内存加载加载模式也是固定的叫 swizzled layout。所以在 NVIDIA 上回答寄存器 tiling 长什么样要分两种情况不用 Tensor Core 时长成你这个线程自己写的展开循环用 Tensor Core 时长成 PTX ISA 里写死的 fragment 分布表。后面这种你没法自由发挥只能顺着来。2.3 寄存器双缓冲让搬运和计算重叠起来寄存器 tiling 不止是放得下还有一个容易被忽略的点共享内存 load 和 FMA 之间的重叠。如果你每个 k 迭代都是先 load 再算load 的延迟大约几十个周期就会被原样吃掉。解决方式是寄存器双缓冲。准备两套 A/B 分片寄存器一套在算第 k 轮的时候另一套提前 load 第 k1 轮的数据。这样 load 延迟被藏在计算里吞吐能再上一个台阶。代价是分片寄存器翻倍8×8 的话就是 1616 个寄存器加上 64 个累加器一下又多了 32 个。这也是为什么很多工程内核宁可选 8×4 或 4×8 的小 tile 也要做双缓冲因为寄存器预算实在有限。CUTLASS 里的Pipeline机制做的就是这件事把 global-shared 和 shared-register 两级流水全部重叠起来。你要是手写内核至少要把 shared-register 这一级做出来否则 tiling 省下的带宽会被延迟浪费掉不少。2.4 CUTLASS 的三级 tile 命名与一套常见默认值如果去看 CUTLASS 的代码或者它生成的 profile会看到几个固定名词Threadblock tile比如 128×128一个 CTA 的输出范围对应共享内存切片大小。Warp tile比如 32×64一个 warp 管的输出范围。128×128 的 CTA 里8 个 warp 排成 4×2正好切出 4×2 个 32×64。Thread tile比如 8×8一个线程管的输出范围。32 个线程排成 4×8铺满 32×64 的 warp tile。很多人调 CUTLASS 不懂这些数字怎么来的。其实很简单Threadblock tile 受限于共享内存容量Warp tile 和 Thread tile 受限于寄存器文件。128×128 的 CTA、32×64 的 warp、8×8 的线程这套组合下来每线程 100 个左右的寄存器是 FFMA 内核在 A100 上一个非常稳的平衡点。你要改就按这个逻辑去改先定线程 tile再乘上线程数看 warp tile 是否匹配最后看寄存器占用和共享内存占用有没有超。3. AMD CDNA 和 Intel GPU看起来是镜像处处微妙不同3.1 wave64 让线程这个单位变了AMD 的 CDNA 架构从 GCN 时代就习惯了 64 个线程组成一个 wavefront也叫 wave而 NVIDIA 是 32 个线程一个 warp。这个差异直接改变了 tiling 的形状。比如 CDNA2 的矩阵指令v_mfma_f32_16x16x16_f16一组 64 个 lane 协同计算 16×16 的输出块。每个 lane 拿到的累加器数量是16×16 除以 64正好 4 个 fp32 寄存器。这个数字跟 NVIDIA m16n8k16 相同但注意参与的是一个 64 线程的组而不是 32 线程的 warp。所以你在设计 warp-level tile 时矩阵块的边长、线程排布全都不一样了。AMD 上你也可以开 wave32 模式但 wave64 默认下的 fragment 分布、bank conflict 行为、指令调度节奏都需要重新适配。3.2 LDS 容量和 VGPR 限制对 tile 选择的实际影响AMD 的寄存器文件在 TU 级别同样是 64K 个 32 位 VGPR 的规模单个 work-item 最多 256 个 VGPR跟 NVIDIA 的 255 基本一致。但共享内存AMD 叫 LDS传统上比 NVIDIA 小GCN/CDNA 一代常用的配置是 64KB 级别而 NVIDIA 这边 A100 单 SM 可以配到 164KBH100 更高。这意味着什么呢你在 AMD 上如果把 Threadblock tile 做得跟 NVIDIA 一样大共享内存可能放不下。所以实践中常见的选择是稍微缩小块级 tile把更多复用压力转移到寄存器层也就是加大每个线程的寄存器 tile。反正寄存器文件总量差不多但 LDS 少了就得靠寄存器去补。这正是同样叫寄存器 tiling不同架构下调参方向相反的典型例子。3.3 Intel 的 DPAS 和 XMXtiling 被做成脉动式点积Intel 的 Xe 架构包括 Ponte Vecchio 和后来的数据中心 GPU用的是 DPAS 指令全称 Dot Product Accumulate Systolic。从名字就能看出来这是把脉动阵列的点积累加直接做成了指令。它的矩阵乘形状、每个线程拿几个累加寄存器也是硬件定死的跟 NVIDIA 的 fragment 概念类似。区别在于 Intel 的 SIMD 宽度是可变的8/16/32这影响了 subgroup 的划分进而影响你在软件层怎么铺 tile。这几个架构总结起来核心差异可以用一张表看明白维度NVIDIA CUDAAMD CDNAIntel Xe基本线程组warp 32 线程wavefront 64可 wave32subgroup 8/16/32矩阵指令示例mma.m16n8k16v_mfma 16x16x16dpas 8x8x...每线程累加寄存器4 个 fp324 个 fp32由指令形状决定共享内存单 SM 最大 160KBLDS 常见 64KB 级别SLM 按 slice 分配tiling 调节空间较大受 LDS 限制偏重寄存器受指令形状限制提示别指望一份调好的 GEMM 内核换个厂商直接跑。寄存器 tiling 是硬件相关的换架构后第一件事是查 ISA 手册里的 fragment 布局和线程组大小然后重新算一遍寄存器预算。4. 离开 GPUCPU 向量寄存器和 NPU 脉动阵列里的 tiling4.1 ARM NEON/SVE没有共享内存编译器替你分了半边天很多人做 AI Infra 会把 ARM CPU 的算子优化忽略掉但端侧推理、边缘场景里CPU 的 GEMM 微内核反而是性能最敏感的部分。在 ARM 上寄存器 tiling 的物理形态又变了一个样没有显式共享内存数据全走 L1/L2 缓存能用的寄存器和显式控制的只有 NEON/SVE 向量寄存器。NEON 有 32 个 128 位向量寄存器一个寄存器正好装 4 个 FP32。经典的微内核结构是每个核心算一个 8×4 或 6×4 的输出块8 个向量寄存器放累加器每个向量对应一行输出的 4 个列再拿几个向量寄存器滚动存 A 和 B 的分片用vmla系列指令反复做乘加。这里没有 Threadblock tile、Warp tile 那套概念tiling 只剩最后一层——寄存器层。而块级的数据搬运完全靠缓存替换策略和预取指令pld/cache hint来管理。到了 SVE 架构事情更有意思。向量长度不再是固定的 128 位而是硬件相关的从 128 位到 2048 位都有Neoverse V1 是 512 位。也就是说你的 microkernel 到底一次算几行几列得运行时根据向量长度决定。很多 SVE 矩阵库的做法是写 sizeless 代码让编译器根据实际 VL 生成对应的 tile 形状。这就是我说的tiling 的另一种长相它不是编译期写死的常量而是运行期多态的形状。这种代码写起来复杂但换来的是同一套 binary 在不同 ARM 芯片上都接近最优。4.2 自研 NPU / 脉动阵列寄存器 tiling 变成空间布局再往上走到自研 NPU比如 TPU 风格或者 Eyeriss 论文里那种脉动阵列寄存器 tiling 的长相彻底变了没有线程、没有 warptiling 发生在数据流图的物理空间里。以 weight-stationary 数据流为例权重矩阵的元素在计算前就加载到了每个 PE 的本地寄存器里整个计算过程中不再移动激活值像流水一样从相邻 PE 传过来每过一个 PE 就跟本地权重做一次乘累加。这时候你问寄存器 tiling 长什么样答案就是每个 PE 的累加器就是它的寄存器 tile而tile 大小变成了一行输出条带在阵列里占据多长一段。分块策略变成了设计数据流的调度顺序——先推哪一行激活、权重在哪一级缓冲驻留、累加结果什么时候回写。在 SIMT 架构上你是靠循环展开寄存器分配在时间维度上做 tiling在脉动阵列上你是在空间维度上做 tiling。这两者对应的代码完全不同但目标一致让数据在离计算最近的地方多待一会儿减少远距离搬运。理解到这一层再回头看你手头的硬件就不会被各种架构的专有名词绕晕了。5. 调 register tiling 这些年我踩过最深的几个坑5.1 寄存器溢出是头号性能杀手且不报错寄存器 tiling 调大收益是复用率提高但超过硬件上限后编译器不会报错而是默默把多余的寄存器spill到本地内存local memory本质上是一块全局内存。问题是spill 访存路径的延迟和带宽比共享内存差得多一旦发生你的内核看似寄存器占用高实际跑起来全是内存访问性能断崖式下跌。最典型的现象就是把线程 tile 从 8×8 调到 12×12理论复用率提高了一倍实际跑出来反而慢了 30%。排查方法很简单。编译 CUDA 内核时加-Xptxas -v输出里能看到每个内核的寄存器数和 spill 数量spill 一旦出现就要警惕。AMD 侧用 ROCm 工具链看 VGPR 数和 local spill 同理。我的经验是宁可 tile 小一点、多开几个 block 保占用率也不要赌编译器能把溢出压回去。寄存器溢出的代价不是线性增长而是断崖。5.2 怎么在动手前先估一个合理的 tile 尺寸不用每次写完整内核才知道 tile 合不合适。拿纸笔就能估个八九不离十公式很简单累加器寄存器数 线程 tile 面积。8×8 的 FP32 累加器就是 64 个寄存器。分片寄存器数 每轮 k 迭代要装的 A 和 B 的个数。8×8 的话是 8816如果双缓冲再翻倍到 32。加上地址、指针、掩码等杂项算 8~12 个。总寄存器数小于硬件上限的 70% 是一个相对安全的区间留出的余量给编译器调度。比如 NVIDIA 单线程 255 个寄存器8×8 双缓冲大概 643212108 个很安全想上 16×16 的 FP32 tile光累加器就 256 个不用算就知道死路一条得换 Tensor Core 或降精度才有戏。占用的另一面是占用率。A100 SM 上 64K 寄存器能同时驻留的线程数 64K / 每线程寄存器数 × 线程数。每线程 108 个寄存器、一个 block 256 线程时约 2.3 个 block占用率在 30% 上下。这个数字对 FFMA 内核来说往往是够用的但你要是做小矩阵、延迟敏感的场景就值得牺牲 tile 尺寸换取更高的占用率。5.3 换架构后我的迁移顺序是固定的每次拿到新的芯片或者从 CUDA 切到 AMD/ARM/NPU我一般按这个顺序来做能少走很多弯路先查 ISA 手册把矩阵指令的 fragment 布局、线程组大小、寄存器上限这几个硬参数抄下来写进自己的调参表。用高层的 DSLTriton、TVM、XLA先跑一版 baseline把 tile 声明成变量用自动调优扫一圈拿到一个不差的参考点。Triton 里num_warps、num_stages这两个参数其实就是手动调 tiling 和流水线的入口值得先摸一遍。在 baseline 基础上手写一版最要紧的内核用 profiler 看两点内存吞吐是否接近峰值、计算单元利用率是否接近上限。哪个没到问题就出在哪一层——前者是块级 tiling 或数据布局的问题后者是寄存器级 tiling 或指令调度的问题。最后才做微调。微调的顺序也固定先调线程 tile 形状再调双缓冲级数最后调数据 swizzle 和 bank conflict 规避。这套流程帮我跳过了很多玄学调参的泥潭。寄存器 tiling 虽然名字里有寄存器三个字但它从来不是一个独立的优化而是跟数据布局、缓存容量、指令调度绑在一起的整体。想把它玩明白最好的办法就是在不同的架构上各写一遍 GEMM感受一下同一个原理在不同硬件约束下的不同表达。等你在两三种架构上都经历过一遍再回头看任何新的芯片扫一眼寄存器文件和向量宽度基本就能猜出它的 tiling 长什么样了。