1. 先把 tiling 这层窗户纸捅破如果你已经在 AI Infra 这条路上混了一阵子大概率见过这样一张图global memory 是仓库shared memory 是柜台register file 是收银员手边的小抽屉。数据一层层往内搬越靠近计算单元访问越快容量越小。tiling 就是决定“从哪个仓库往哪个抽屉搬多少货”的学问而寄存器 tiling 是其中最贴近执行单元、也最容易被忽视的一层。为什么单独拎出来问“在不同架构上长什么样”因为同一个概念在不同硬件上的呈现差异非常大甚至可以说你在 NVIDIA GPU 上理解的 register tiling直接套到 AMD 或移动端 GPU 上可能会理解错。原因是每家架构的寄存器堆register file组织方式、线程模型、矩阵指令抽象都不同于是寄存器里那块扁平的“小积木”到底长几乘几每一块像素级的 fragment 归属哪个 lane完全不是一个套路。我最初是从 CUTLASS 和 TVM 的代码里接触到这个概念的。一开始以为 register tiling 就是把一个 16x16 的 C tile 平均分给线程每人手里攥着几把 float 而已后来真正逐行追过 SASS、看过 AMD 的 wave32 layout、又跑到手机 GPU 上调了一些 compute shader才发现“在不同架构上长什么样”这个问题其实是在问硬件设计者对“计算主体”和“数据复用”的底层假设。这篇文章就当是把这个大问题拆成几个小问题系统盘一遍。先说结论所谓寄存器 tiling核心目标是让每个线程用尽量少的 load 指令把从 shared memory 或 global memory 取得的数据尽可能多地在寄存器里复用。寄存器访问延迟通常只有几个周期而 shared memory 需要几十个周期global 直接上百如果能把一块数据留在寄存器里被多次计算整体吞吐会非常可观。但每个架构能留在寄存器里的数据形状、每个线程能分到的数量、以及矩阵指令希望数据以什么格式出现差别极大。1.1 从“搬数”说起为什么先有 global、shared才轮到寄存器我们经常听说的 tiling 一般默认指的是 shared memory tiling。比如 GEMM 里把 4096x4096 的大矩阵切成 128x128 的 block再让一个 thread block 内的线程协作把这 128x128 的块加载到 shared memory然后在共享内存里反复读取、计算。这确实是最重要的宏观分块因为它决定了线程块之间的数据独立性也决定了访存复用率。但 shared memory tiling 做完之后每个线程依然需要从 shared memory 里取数。如果每个线程每算一个输出元素都去 shared memory 取一次 A 和 B那 shared memory 的带宽马上就会成为瓶颈而且会有大量浪费。所以真实实现的下一步是把每个线程负责的“微块”micro-tile继续切分让微块数据落在寄存器里内层循环完全基于寄存器读写。举例来说假设线程块负责 128x128block 内有 256 个线程那么每个线程负责 8x8 的 C 块。如果没有寄存器 tiling直接写成外层遍历 8x8、内层遍历 K那每个输出点都会从 shared memory 读 K 次 A 和 K 次 B而 8x8 的寄存器 tiling 会把属于这个线程的 8 行 A、8 列 B 按需放入寄存器K 循环内部只需从 shared memory 读新的 K 切片旧的 K 切片完全留在寄存器里反复参与 64 次乘加。所以寄存器 tiling 不是和 shared tiling 并列的另一个方法而是 shared tiling 底下必须再补的一层。两者一起构成了常见的三层分块global 分块决定 block 层面调度shared 分块决定 block 内协作register 分块决定单线程计算密度。1.2 寄存器 tiling 的定位真正决定指令级并行上限的一层寄存器 tiling 最本质的价值是它直接决定了一个线程一个循环里能做多少次独立的乘加运算也就是指令级并行ILP的上限。编译器只要能把这些独立运算调度开ALU 就会一直满载。而这个上限其实在写代码时就定死了如果你的寄存器 tile 设计得让同一时刻可并行的 mma矩阵乘加指令太少后面再怎么调优也救不回来。除此之外寄存器 tiling 还对 shared memory 的 bank conflict 有间接影响。你从 shared memory 读数据时如果想避免 bank conflict要设计成连续线程访问连续地址而由于每个线程读的数据最终要落到自己的寄存器 fragment 中线程与地址之间的映射关系就是由寄存器 tiling 的 layout 决定的。很多做 GEMM 优化的人都有过这种经历shared 阶段明明没冲突进了寄存器阶段换了一种 fragment 分配回头一看 shared load 又冲突了。所以这一层看起来是寄存器的事实际上和整个数据通路都耦合。看不同架构的 register tiling本质上是在看架构设计师怎么权衡线程数量、寄存器容量、矩阵指令宽度这三者的关系。下面进入正题。2. NVIDIA GPUwarp 级 micro-tile 的标准答案NVIDIA GPU 是大多数 AI Infra 工程师接触到的第一块并行计算硬件也是寄存器 tiling 感知最强的地方。以 Ampere 架构的 GA100 为例每个 SM 有 64K 个 32 位寄存器换算成字节是 256KB 寄存器堆跟 L1/Shared memory 容量差不多大。每个 thread 最多可以用 255 个寄存器每个 warp 是 32 个线程。对 GEMM 来说一个 warp 负责一个 warp tile比如 64x64然后 warp tile 内部再切成线程级 micro-tile通常形状是 8x8也有 16x8、32x8 之类的变体。这里有个关键点NVIDIA 从 Turing 开始引入了 HMMA 指令Ampere 又加了三代 tf32、bf16、fp16 的 mma 指令它的硬件矩阵单元天然希望数据以一种固定 fragment 布局出现在寄存器中。你自己把寄存器当数组来回读写当然也可以但除非用上 mma 类指令否则没法吃到 tensor core 的吞吐。而 mma 指令的 fragment 布局就是 NVIDIA 对寄存器 tiling 的“官方答案”。2.1 硬件基础寄存器堆、warp 与 mma/wgmma在写任何 CUDA 或 PTX 之前先把硬件事实摆清楚。SM 里的寄存器堆按 lane 组织每个 lane 有自己独立的一组寄存器文件。对 warp 内的第 L 个 lane 来说它只能访问自己名下的寄存器不能直接读其他 lane 的寄存器。这一点很关键你说的“寄存器 tiling 数据在寄存器之间流转”实际上只发生在同一个 lane 内部lane 之间交换数据必须通过 shared memory 或 shuffle。Tensor Core 的 mma 指令相当于把寄存器 tiling 逻辑固化成了硬件微码。以 Ampere 常见的 m16n8k8fp16为例一条指令由 warp 内 32 个 lane 协作完成 16x8x8 的矩阵乘加。指令规定A 矩阵是 16x8按 k 维 8 拆开后每个 lane 持有 A 的 4 个元素表现为寄存器称为 a-fragmentB 矩阵是 8x8每个 lane 持有 B 的 2 个元素b-fragmentC/D 是 16x8每个 lane 持有 4 个元素c/d-fragment。四个片段合在一起才是这个 warp 在当前 k 切片的寄存器 tile 全貌。Hopper 上进一步出现了 wgmma把 mma 的协作范围从 warp 扩大到 warpgroup四个 warp。wgmma 的一条指令可以同时指定 A 矩阵在 shared memory 中、B 矩阵在寄存器中的复杂布局也可以反过来。这让寄存器 tiling 的形态变得更加多样本质上是在用硬件调度器教你A 放 shared 或寄存器、B 放寄存器或 shared、累加器放寄存器可以按需组合。2.2 reg tile 长什么样以 CUTLASS 中常见的 128x128 block、256 线程8 warps为例每个 warp 分到 64x64 的 warp tile一个 warp 内再把 64x64 的 C 矩阵切成 8x8 的 micro-tile于是共有 8x864 份每份对应一个 lane恰好覆盖 32 个线程加循环展开的两份。每个线程持有这个 8x8 micro-tile 里的 8 个 float也就是 8 个 32 位寄存器。从内存布局上看A 矩阵的寄存器 fragment 通常是按 k 维度连续的。假设每轮从 shared memory 加载 16x16 的 A 块warp 中每个 lane 会持有若干 a-fragment这些 fragment 在寄存器中的编号是连续的方便编译器做双缓冲和乱序调度。B 矩阵类似。C 矩阵的 fragment 严格按 (row, col) 分布一个线程的 8 个值往往不是原始矩阵里连续的 8 个点而是按固定 stride 取到的。这就是官网文档里那张密密麻麻的“lane 编号占位图”的由来。2.3 代码实现示意下面用简化的伪代码展示 Ampere 上 m16n8k8 的寄存器 tiling 流程。假设我们已经在 shared memory 里有 A_tile 和 B_tile每个线程需要把随后 8 个 k 步的数据准备到寄存器里。// warp 内每 lane 持有 a_frag[4], b_frag[2], c_frag[4] float a_frag[4]; float b_frag[2]; float c_frag[4] {0, 0, 0, 0}; // k 维度迭代这里每步处理 k 8 for (int k0 0; k0 K_TILE; k0 8) { // 从 shared memory 加载 A 的 16x8 分片到寄存器 // 实际地址由 lane_id 和 mma 官方 layout 决定 load_a_fragment(a_frag, A_smem, k0); load_b_fragment(b_frag, B_smem, k0); asm volatile( mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n : f(c_frag[0]), f(c_frag[1]), f(c_frag[2]), f(c_frag[3]) : r(a_frag[0]), r(a_frag[1]), r(a_frag[2]), r(a_frag[3]), r(b_frag[0]), r(b_frag[1])); }这里 a_frag 的 4 个寄存器、b_frag 的 2 个寄存器就是当前 k 切片对应的寄存器 tile。它们不像数组那样在内存里连续而是被 mma 指令的逻辑硬编码成一种分布式布局。如果你用传统 SIMT 风格手写乘加循环当然也可以在寄存器里做 8x8 的 FMA但那样就吃不到 tensor core 了。2.4 ptxas 角度寄存器分配与 occupancy 权衡NVIDIA 这一侧寄存器 tiling 的形状最后会反映在编译结果里。用ptxas -v可以看到每个 kernel 用了多少寄存器有没有 spill。一个 128x128 block、8x8 micro-tile 的 HGEMM kernel通常会用到每个线程 168~232 个寄存器这正好卡在占用率比较舒服的区间一个 SM 上能同时驻留的 warp 数 65536 / (寄存器数 * 32)。如果每线程 224 个寄存器那么 65536 / (224*32) 9.14向下取整 9 个 warp因为 warp 调度器不支持 9 个 warp 均匀分配到 4 个 scheduler实际会退到 8 个 warpSM 利用率会明显下降。这里有个 tips不要一味追求大 tile。寄存器 tile 从 8x8 放大到 16x8C fragment 的寄存器数翻倍占用率立刻跳水。对于纯计算密集型的 GEMM你当然希望 tile 大一点好复用数据但对于带宽敏感或者需要高并发隐藏延迟的 kernel适当的 register tiling 加高占用率效果反而更好。具体得靠 profiling没有统一答案。3. AMD CDNAwave 级 tile 与 VGPR/SGPR 的分工AMD 的 GPU 和 NVIDIA 走了不完全相同的路。CDNA 架构MI100、MI200、MI300 系列使用 wavefront 作为调度单位一个 wavefront 是 64 个 lane也可以配置成 wave32RDNA 后原生支持 wave32CDNA 也提供选项。寄存器堆分为两类Vector GPRVGPR和 Scalar GPRSGPR。VGPR 是每个 lane 都有自己的副本存逐线程数据SGPR 是所有 lane 共享一份存常量、地址和统一数据。这个双寄存器堆设计直接影响了寄存器 tiling 的形态。在 AMD 的 MFMA 指令matrix core 指令类似 NVIDIA 的 mma里A、B、C 的 fragment 分布在 VGPR 中而地址、步长、即时数等“元数据”放在 SGPR 中。SGPR 的有效使用能减少 VGPR 的浪费因为你不用在每个 lane 里重复保存相同的地址。这是 CDNA 架构在编译器约束上比 NVIDIA 更精细的地方。3.1 CDNA 有别于 NVIDIA 的那几件事第一个区别是 wave64 的 lane 数。同样一个 16x16 的寄存器 tilewave64 里会被切成更多份每个 lane 持有的寄存器更少wave32 下每个 lane 持有更多元素指令数更少但寄存器压力更大。AMD 的编译器经常需要在 wave32 和 wave64 之间做选择CDNA 的矩阵核心在 wave64 下支持更多指令形状而 wave32 在纯 scalar 逻辑或低占用场景下更高效。第二个区别是指令的 fragment 尺寸选择。CDNA 的 MFMA 支持 m16n16k16、m32n32k8、m64n64k4 之类的组合m64n64k4 对每个 lane 的寄存器要求更大适合大 K 复用m16n16k16 则更均衡。这种按需选择的灵活性让寄存器 tiling 变成了一个真正的参数空间而不是固定模板。实际调优时我们往往先用 m16n16k16 跑通功能再切换到 m32n32k8 去看 latency 和 occupancy 的平衡。第三个区别在 shared memory 的命名上AMD 叫 LDSLocal Data Share它的 bank 机制和 NVIDIA 类似但访问竞争规则细节不同。由于寄存器 tiling 决定了 LDS 读取的数据如何分配到各 lane同一个 16x16 块用不同的 fragment 布局读入 LDS冲突率可能差几倍。AMD 的 rocBLAS 和 Composable Kernel 内部都维护了很多针对不同 CDNA 芯片的 layout 变体背后的原因就在这。3.2 一个 wave32 的 reg tile 示例假设我们用 wave32 做一个基于 MFMA 的 HGEMM选择 m16n16k16 的指令形状。每个 wave 负责 16x16 的 C tileA 是 16x16 的 tileB 是 16x16 的 tile。按 MFMA 的文档A 的 16x16 会被拆成若干 32 位 VGPR每个 lane 拿 8 个16x16 / 32 8B 同样每个 lane 拿 8 个C 的 16x16 累加器每个 lane 拿 4 个 VGPR这是由 16x16x4 字节 / 32 lane 32 字节 8 个 32 位寄存器不对准确说 C 是 16x16 的 f32即 1024 字节分给 32 个 lane每个 lane 32 字节即 8 个 VGPR。让我重新算一下。A 矩阵是 16x16 的 fp16共 512 字节每个 lane 拿到 16 字节即 4 个 16 位元素吗MFMA 的实际文档里m16n16k16 的 A fragment 每个 lane 持有 4 个 fp16B 是 4 个 fp16C 是 8 个 f32。这个具体数字因芯片世代而异读者不一定要背真正重要的是理解布局不是连续排列。写伪代码的话和 NVIDIA 的例子结构类似只是指令助记符从mma.sync.aligned变成了v_mfma_f32_16x16x16_f16。为了完整可以罗列一下伪代码// 假设 wave32v_mfma_f32_16x16x16_f16 // 每个 lane 的 VGPR // a_vgpr[4]A fragment // b_vgpr[4]B fragment // c_vgpr[8]C/D 累加 for (int k0 0; k0 K_TILE; k0 16) { lds_load_a_to_vgpr(a_vgpr, ...); lds_load_b_to_vgpr(b_vgpr, ...); asm volatile( v_mfma_f32_16x16x16_f16 %0, %1, %2, %0 : v(c_vgpr) : v(a_vgpr), v(b_vgpr)); }注意这里 C 累加器必须用 “” 约束因为 MFMA 是累加操作。在 CDNA2 之后MFMA 还支持双发行和 LDS 直读某些组合下 B 矩阵可以直接在指令里引用 LDS 地址减少一次 VGPR 中转。这种“寄存器 tiling 与 LDS 读取融合”的模式是 NVIDIA 在 Blackwell 之前都不太强调的设计。3.3 换到 RDNA/移动端后的变化RDNA 系列虽然也使用波前但主要面向图形和消费级计算矩阵核心能力弱很多。RDNA3 加入了 WMMAWave Matrix Multiply-Accumulate指令支持 16x16x16 的 fp16 计算但寄存器文件仍沿用 VGPR/SGPR 双 bank 机制。在 RDNA 上做寄存器 tiling 优化时要考虑的不仅是 matrix core 指令还包括 wave32 下 VGPR 数量对 shader occupancy 的影响。移动端 RDNA比如 RDNA3 的 iGPU往往没有独立显存带宽和 L2 都很有限寄存器 tiling 更需要和 LDS 配合尽量让每个 wave 在寄存器里完成多轮复用减少对全局内存和二级缓存的依赖。这部分实战中比较痛苦的地方在于移动端驱动对寄存器分配的干预程度较高你写的num_SGPR/num_VGPR约束不一定被严格尊重需要实测。4. 移动端 GPU寄存器被当作“向量寄存器组”用的一层手机 GPU 是 AI Infra 工程师经常忽略但实际非常重要的场景。端侧 AI 推理跑在 Adreno、Mali、PowerVR 这类 GPU 上它们没有像 NVIDIA CDNA 那样被明确文档化的mma指令体系Adreno 有一些uGQR矩阵指令Arm 也加入了smmla等但寄存器 tiling 同样成立只是表现形式更接近“把寄存器当成 SIMD 向量寄存器组来用”。这类 GPU 的底层执行模型是“一个着色器程序在多个 SIMD 通道上并行”。每个通道有自己的寄存器但编译器通常把它们组织成向量寄存器。你在 Adreno 的 compute shader 里看到一个 128 位寄存器r0.xyzw可以理解成同时容纳 4 个 fp32或者 8 个 fp16。寄存器 tiling 在这里就表现为每个线程或者每个像素/每个 work-item在自己的向量寄存器里维护一个小矩阵块比如 4x4 或 8x4然后一次 ALU 操作能同时对 4 个通道做 FMA。这个形态看起来没有 NVIDIA 那么复杂但约束反而更极端。移动端 GPU 的寄存器文件往往比桌面级小很多Adreno 某些代际上每个核心的寄存器总量一体超限就 spill 到 local memory而移动 GPU 的 local memory 带宽和延迟都远不如桌面 GPU 的 shared memoryspill 代价极其昂贵可能直接让 kernel 从 2ms 变 20ms。所以移动端的 register tiling 优化更像是在“如何用最少的寄存器表达最大的数据复用率”。4.1 Adreno 的 scalar/vector 双寄存器堆Adreno 的架构从 6xx 到 8xx 一直在演进但大体保留了 scalar 和 vector 两套寄存器。scalar 寄存器存每个线程都相同的值比如循环计数、统一缓冲区的参数、计算地址的公共部分vector 寄存器存随线程变化的数据。这个设计原理上和 AMD 的 SGPR/VGPR 是同构的。在 Adreno 上写 compute shader编译器自动把 uniform 数据和 per-thread 数据分开分配。如果你能手动保证某些矩阵 tile 的行列步长是一致的编译器可以把地址计算尽量下沉到 scalar 寄存器减少 vector 寄存器的浪费。实测中很多 kernel 的性能瓶颈并不在 ALU而在寄存器压力过大导致 occupancy 太低或者地址计算占用了大量 vector 寄存器。Qualcomm 后来引入了 uGQR 这类硬件矩阵指令目的就是让寄存器 tiling 的 fragment 布局有硬件标准避免编译器瞎猜。但这类指令主要服务于特定精度和特定大小比如 8x8x8实际开发中如果用不到仍要靠手写向量寄存器块来优化。4.2 Mali Valhall 的 reg tile 实践Arm Mali Valhall 架构G77 之后的架构代号和上一代 Bifrost 有本质区别完全转向了基于 warp 的标量/向量统一执行模型每个 warp 有 16 个 lane某些实现里是 8 个支持 fp16 的 ALU 对半分。Mali 的寄存器堆是按执行引擎Execution EngineEE组织的每个 EE 有固定数量的寄存器如果超出就 spill。在 Mali 上做寄存器 tiling 时我发现最合适的是 8x8 或 16x8 的 fp16 tile。原因很简单Mali 的 SIMD 宽度是 128 位一次性可以处理 8 个 fp16。假如每个 lane 持有 8x8 tile 里的一行那么 8 行数据正好填满 8 个 128 位向量寄存器加上 C 的 8 行又是 8 个寄存器A/B 各若干总计大约 24~32 个向量寄存器。这种情况下 occupancy 还能保持不错的水平。如果把 tile 加到 16x16寄存器会飙到 60 个以上占用率立刻崩掉。Mali 的编译器和驱动对寄存器分配也有很强的主导权它可能会在你意想不到的地方插入补丁指令patching。一个常见坑是你精心安排的寄存器 tile 会被编译器重排导致数据布局变化性能反而下降。我的经验是在这种平台上不要过度追求寄存器内部布局除非你愿意阅读反汇编和驱动调度日志。4.3 PowerVR 和其他移动架构的备注PowerVR 架构的 GPU在部分 SoC 中出现使用的是统一存储和基于 tile 的渲染方式计算着色器也沿用类似的寄存器组织。它的 TBDR 设计虽然对渲染很友好但在通用计算上寄存器 tiling 的发挥空间比较受限因为线程调度被硬件抽象得很完整你很难通过手工布局寄存器来压榨额外性能。一般的建议是用厂商提供的 compute library比如 IMG 的 OpenCL SDK的推荐模式不要自己手动搞复杂 tile。顺便说一句鸿蒙或国产 SoC 里的 GPU 很多基于 Mali 或 PowerVR 的 IP也有部分自研架构但寄存器相关的开发接口基本都是 OpenCL / Vulkan compute。指令集形态虽有差异寄存器 tiling 的思考方式完全一致先理解每个 lane 有几个可用的 32 位寄存器再决定 micro-tile 形状。5. CPU 侧向量化视角下的寄存器 tiling我们通常说寄存器 tiling 都是针对 GPU但其实 CPU 上的 SIMD 向量化也可以从这个角度理解。AVX-512 有 32 个 512 位寄存器ARM SVE 有可变的向量长度寄存器NEON 有 32 个 128 位寄存器。这些寄存器和 GPU 的 lane 寄存器不同它更像“一个寄存器就是一排小抽屉”SIMD 指令一次操作一排。所以在 CPU 上做寄存器 tiling本质是“把一个小矩阵块的每一行或每一列映射到一组向量寄存器里”。比如对 8x8 的 fp32 矩阵乘你用 AVX-512 的话可以一个向量寄存器装一行 8 个元素C tile 就是 8 个向量寄存器A 的一行广播到 8 个向量寄存器和 B 的 8 列逐列乘加。这就是典型的寄存器 tiling不经过内存完全在寄存器里完成 8 个向量寄存器和 8 个向量寄存器之间的外积累加。这个套路和 GPU 的区别在于GPU 的寄存器 tiling 往往要依赖硬件矩阵指令而 CPU 的寄存器 tiling 可以直接写成普通的向量 FMA。编译器GCC、Clang、MSVC通常会自动向量化一部分但高级优化还是得手写 intrinsics。AI Infra 里的推理引擎优化对 CPU 后端也常做类似工作比如 oneDNN 的 jit kernel 就会把 4x16 或 8x16 的 micro-tile 硬编码到寄存器和指令调度里。5.1 SIMD 寄存器也是 tiling用矩阵乘法来举例最直观。假设你要算一个 4x4 的 C A x B把 A 的每一行看成一组 SIMD 数据。最简单的 SIMD 方式C 的每一行用一个向量寄存器当累加器内层遍历 K 时把 A 元素的标量广播到所有 lane然后乘 B 的一行向量加到 C 行向量上。这样每个 C 行向量就相当于一个寄存器 tile 的行。如果进一步优化可以把 A 的连续 4 行同时保留在 4 个向量寄存器里B 的 4 列也在 4 个向量寄存器里然后做 4x4 个向量 FMA。这听起来和 GPU 的寄存器 tile 一模一样。区别只在于CPU 的“lane”和“寄存器”是同一个概念的不同表现一个 512 位寄存器相当于 16 个 fp32 lane和 GPU 中一个 lane 拥有若干寄存器恰好是反过来的。5.2 ARM SVE 与 x86 AVX-512 的例子ARM SVE 的一个特点是向量长度可以支持 128 到 2048 位软件不用硬编码长度这使得寄存器 tiling 的代码需要写成与 VL 无关的形式。比如你要实现一个基于外积的 GEMM micro-kernel写代码时不能假设一个向量寄存器能装 16 个 fp32只能用svcntw()得到当前机器的 lane 数再据此切分 tile。好处是同一份代码可以在不同 SVE 机型上获得不同性能坏处是编译器优化压力很大。x86 的 AVX-512 固定是 16 个 fp32 per register写代码时往往直接硬编码。我常用的一个 fp32 GEMM 寄存器 tile 是 8x16C 的 8 行每行占一个 zmm 寄存器16 个 fp32这样 C tile 一共 8 个 zmm。A 的一行 16 个元素可以打包进一个 zmmB 的 8 列也各占一个 zmm。内层循环时A 行广播到 8 个寄存器分别乘 B 的各列累加进对应的 8 个 C 行寄存器。一轮 k 步消耗 8 4 8 之类的寄存器不做特殊复杂处理也能保持 28~32 个寄存器以内不会造成溢出。5.3 与 GPU 的一致性和差异CPU 和 GPU 寄存器 tiling 的共性在于都是“数据复用优先减少访存次数”。差异在于数据分布模型GPU 是“多个线程各拿一小块”CPU 是“单线程用 SIMD 通道构成一个小块”。这两种模型各有优劣。GPU 的 tiling 适合大规模并行和硬件矩阵指令CPU 的 tiling 适合延迟敏感、并行度有限的场景。理解这个一致性后你会发现一个非常有用的迁移思路在 GPU 上研发的 tiling 策略可以先在 CPU SIMD 上做快速原型验证因为 CPU 上调试寄存器布局、观察 spill 都比 GPU 方便得多。等逻辑正确了再移植到 GPU 上做 mma / MFMA 适配。6. 一个完整例子HGEMM 里手工设计 reg tile这一节我们用实际可操作的步骤把之前的抽象概念落到一个 HGEMM kernel 上。假设目标平台是 NVIDIA A100Ampere输入的矩阵是 fp16输出累加成 fp32K 维度较大。我们不依赖 cuBLAS自己写一个 register tiling 的示例目的在于看清每一步的形状选择。6.1 从 GEMM shape 往下拆一个典型的 A100 HGEMM 分块流程如下全局分块每个 thread block 计算 128x128 的 C tile。block 内线程数256 线程组织成 8 个 warp每个 warp 负责 64x64 的 C tile。warp 内每个 warp 的 32 个 lane 共同负责一个 64x64 的 warp tile。寄存器层通常每个 thread 负责 8x8 的 micro-tile也就是 C 的 64 个元素。这里的 8x8 不是拍脑袋定的它和 mma 指令的 m16n8k8 形状有对应关系。一个 lane 同时会参与若干个 mma 片段经过两个 mma 步骤才能覆盖 8x8 的 C tile。如果改成 16x8 的 micro-tile寄存器占用会翻倍到 16 个 f32但同一个 lane 需要参与的 mma 次数变多指令级并行更高适合 K 较短、block 数多的场景。实际项目里A100 上用 8x8 或 16x8 都有人做具体以 profile 为准。寄存器预算方面8x8 的 C tile 需要 8 个 fp32 寄存器循环里同时要保留 A 的 16x8 fragment 和 B 的 8x8 fragment以 fp16 形式各自折合若干 32 位寄存器再加上地址、指针、循环变量整体大概 100~128 个寄存器。这个量级能保证 SM 上的 warp 数在 10 个左右ALU 有足够的并发来掩盖 shared memory 延迟。6.2 选 micro-tile 形状与 fragment 分配选 micro-tile 形状时要考虑两个指标数据复用率和寄存器占用。假设每个线程负责 8x8 的 C tile那么在 k 维内层循环里这个线程需要同时持有 A 的 8 个元素和 B 的 8 个元素每个 k 步加上 C 的 8 个累加器一共 24 个寄存器左右。这是一份低压力配置适合启动的很早、或者对 occupancy 要求很高的 kernel。如果选 16x8 的 C tile就需要 A 的 16 个元素、B 的 8 个元素、C 的 16 个累加器接近 40 个 32 位寄存器属于中等压力配置。再往上选 32x8C 本身就 32 个寄存器加上 A/B 很容易突破 80 个除非对 L2 命中率有极大自信否则不建议在 A100 上这么干。Fragment 分配则取决于你用什么指令。如果手写 PTX 的 mma.m16n8k8A fragment 的 4 个寄存器编号在 lane 间是错开的编译器会帮你插入合适的从 shared memory 到寄存器的加载序列。如果你用 CUTLASS 的MmaTensorOpfragment 分配由模板参数里的MmaPolicy决定内部的TiledMma会把每个线程负责的元素下标映射得清清楚楚。我个人的建议是一开始不要自己手工映射 lane 和寄存器先用 CUTLASS 的默认配置跑通再逐步魔改。6.3 主循环体k 循环中的 reg tiling假设 shared memory 里已经放好了 64x64 的 A tile 和 B tile寄存器 tiling 的主循环体大致如下// 每个线程注册 8 个 C 累加器 float c[8] {0}; // 以 16 为步长在 k 维上迭代对应一个 shared tile 切片 for (int k 0; k 64; k 16) { // 从 shared memory 取 A 的 16x8 到寄存器 a_frag float a_frag[8]; // 简化的描述实际可能拆成 4 个寄存器对 load_a_into_reg(a_frag, smem_a, k); // 从 shared memory 取 B 的 16x8 到寄存器 b_frag float b_frag[8]; load_b_into_reg(b_frag, smem_b, k); // 4 次 mma.m16n8k8 完成 8x8 的累加 // 每次 mma 指令处理 16x8x8这里需要按布局重组 mma_16x8x8(c[0], a_frag, b_frag); mma_16x8x8(c[1], a_frag, b_frag); }这段代码忽略了很多索引细节但结构对了A/B 各自从 shared 载入寄存器然后连续做 4 次 mma。这里的“连续”是重点它让寄存器中的 a_frag、b_frag 被复用多次而不是每次 mma 都重新读 shared。这也是寄存器 tiling 相对朴素实现的最高价值点。实际调优中还会在load_a_into_reg和 mma 之间做软件流水线也就是提前预取下一轮 k 的 A/B 到另一组寄存器让当前 mma 和下一次 shared load 并行。这样寄存器占用会上升但延迟能被隐藏掉吞吐往往更好。CUTLASS 的PipelineTmaAsync或传统cp.async双缓冲方案本质上就是在为这个目标服务。6.4 完整可参考的架构选型参数表这里整理一份我在不同架构上做 fp16 GEMM 时经常参考的初始参数表。注意只是起点不代表最优真实场景请按 profiling 调整。架构典型线程模型推荐 micro-tile寄存器类型核心指令/机制NVIDIA Turing/Amperewarp 32 lanes8x8 / 16x8普通 32 位寄存器mma.sync / wgmmaNVIDIA Hopperwarpgroup 128 lanes16x16 / 32x16普通 32 位寄存器wgmma TMEMAMD CDNA2/3wave32 / wave6416x16 / 32x32VGPR SGPRv_mfma / v_wmmaQualcomm Adreno 7xx每 SIMD 通道 thread8x8 / 8x16 fp16scalar vector regs手写向量 FMA / uGQRArm Mali Valhallwarp 16 lanes8x8 / 16x8vector regs手写向量 FMA / smmlax86 AVX-512单线程多 lane8x16 fp32zmm 向量寄存器FMA / 广播ARM SVE单线程多 laneVL 自适应 tileSVE 向量寄存器FMLA / 外积这张表不是让你死记硬背而是提醒一个问题不同架构的 micro-tile 形状差异极大这背后是寄存器堆容量、lane 数和指令宽度共同作用的结果。在 NVIDIA 上顺手的 16x8放到 Mali 可能因为 lane 数不同而完全发挥不出来反过来也一样。7. 常见问题与排查技巧实录做寄存器 tiling 优化时最头疼的不是理论不懂而是“编译器不按我想的来”。下面几个问题是我自己在实践中反复遇到的基本每天都在调优群里被问一遍。7.1 寄存器溢出最容易翻车的地方寄存器溢出spill是寄存器 tiling 的头号敌人。如果你的 kernel 编译后出现local memory访问大概率是寄存器压力超过了硬件上限。NVIDIA 上可以用ptxas -v看 spill 数量AMD 上可以用ROCm的编译日志或者rocprof看 VGPR 分配移动端则要看驱动或 profiling 工具的 register usage 指标。我的排查顺序是先看寄存器占用再看 occupancy再看实际运行时间。如果寄存器占用很高但 occupancy 没跌可以先跑一下如果 spill 已经出现必须降 tile 或者减少双缓冲级数。很多时候减少一个循环展开因子就能避免 spill性能不降反升。记住一个原则宁可让 ALU 动用不完美的指令序列也不要让数据跑到 local memory 去因为 spill 的代价常常是数倍的 memory 访问延迟。7.2 occupancy 掉下来tile 越大未必越好常有人觉得寄存器 tile 越大数据复用率越高性能一定越好。但寄存器总量是物理限制。NVIDIA A100 的 SM 只有 65536 个 32 位寄存器如果每个线程用 256 个一个 SM 只能放下 8 个 warp而一个 SM 的四个调度器理想情况下需要至少 16 个 warp 才能保持满负载。一旦 warp 数不足访存延迟就藏不住kernel 照样跑不快。AMD CDNA 同样有 VGPR 上限的约束。对于需要用 wave64 的 kernel如果每个 lane 分到的 VGPR 超过 128wave 数就会下降。这时候可以尝试切 wave32虽然一个 wave 的线程数少了一半但同一时刻可调度的 wave 数量可能翻倍总体吞吐反而提升。移动端更是如此Adreno 和 Mali 的寄存器堆更小过大的 tile 会让 occupancy 暴跌导致 latency hiding 完全失效。7.3 shared memory bank conflict 仍然存在的寄存器 tiling 并不能解决所有 shared memory 冲突。很多人以为数据进了寄存器就万事大吉实际上你从 shared memory 到寄存器的那次 load 本身就要面对 bank conflict。而这次 load 的地址模式直接由 register tiling 的 layout 决定。比如 NVIDIA 上常见的“跨步读取”导致同一 bank 被多个 lane 争抢就是 fragment 分配不合适的结果。排查方法是看 profiling 里的 shared load bank conflict 指标。如果冲突率很高可以尝试改变 micro-tile 的方向比如把 C tile 从 16x8 改成 8x16很多情况下 A/B 的加载模式会跟着改变冲突就自然消掉了。AMD 的 LDS 也有类似问题不过它的 bank 数和 conflict 检测方式和 NVIDIA 不完全一致需要单独看文档。7.4 工具链怎么看 reg tiling 效果实话说最直接的手段是看 SASS/汇编。NVIDIA 上用cuobjdump -sass看每条指令用哪些寄存器AMD 上用LLVM MC反汇编可以看 VGPR/SGPR 使用移动端则比较麻烦但多数厂商 profiling 工具会输出 register usage per thread。另一个实用技巧是故意把 tile 改成极端值比如 1x1 和 64x64分别测时间从两端逼出寄存器 tiling 的贡献就能判断当前 kernel 到底卡在寄存器还是卡在带宽。Nsight Compute 的Speed Of Light页面和Warp State Statistics面板对 NVIDIA 非常有用能看到 stall 原因是什么AMD 的rocprof加--stats也可以给出 instruction mix 和 VGPR 分配。移动端通常用厂商的 GPU profiler比如 Qualcomm Snapdragon Profiler 或 Arm Streamline观察 shader core 的 ALU 利用率和 register bank 压力。7.5 问题速查表下面是我自己整理的快速定位表照着排查往往能省很多时间。现象可能原因下一步动作kernel 运行时间远超预期register spill降 micro-tile / 减少双缓冲级数profiling 显示 memory stall 很高occupancy 不足降低寄存器占用提高 warp/wave 数shared load bank conflict 高fragment 布局和地址映射不匹配换 tile 方向检查 lane 到地址的映射编译器没生成 mma/MFMA数据类型或 fragment 布局不符合指令要求检查指令约束必要时写 PTX/intrinsic改了 tile 大小性能没什么变化瓶颈不在寄存器层改测 global/shared tiling或 L2 命中率这张表没有涵盖所有情况但覆盖了 80% 的初调问题。如果你遇到的现象不在这张表里建议回去看汇编一条STL或V_ADD出现的位置往往能告诉你编译器到底在寄存器层面做了什么奇怪的事。最后说一点个人体会。寄存器 tiling 这东西看十篇文档不如动手改一次 kernel。我见过很多人卡在“寄存器 tiling 和 shared tiling 的区别”这种概念问题上其实只要先写一个只做 shared tiling 的版本再写一个加了 register tiling 的版本两边跑一下性能立刻就能感受到这一层的价值。而且不同架构的差异并不可怕它们背后的原则其实是同一个把最热的数据放在最靠近 ALU 的存储里让它在被丢弃前尽可能地多算几次。搞明白这一点换架构只是换 API 和 layout 表而已。
寄存器Tiling跨架构解析：从NVIDIA到移动端GPU的寄存器分块设计
1. 先把 tiling 这层窗户纸捅破如果你已经在 AI Infra 这条路上混了一阵子大概率见过这样一张图global memory 是仓库shared memory 是柜台register file 是收银员手边的小抽屉。数据一层层往内搬越靠近计算单元访问越快容量越小。tiling 就是决定“从哪个仓库往哪个抽屉搬多少货”的学问而寄存器 tiling 是其中最贴近执行单元、也最容易被忽视的一层。为什么单独拎出来问“在不同架构上长什么样”因为同一个概念在不同硬件上的呈现差异非常大甚至可以说你在 NVIDIA GPU 上理解的 register tiling直接套到 AMD 或移动端 GPU 上可能会理解错。原因是每家架构的寄存器堆register file组织方式、线程模型、矩阵指令抽象都不同于是寄存器里那块扁平的“小积木”到底长几乘几每一块像素级的 fragment 归属哪个 lane完全不是一个套路。我最初是从 CUTLASS 和 TVM 的代码里接触到这个概念的。一开始以为 register tiling 就是把一个 16x16 的 C tile 平均分给线程每人手里攥着几把 float 而已后来真正逐行追过 SASS、看过 AMD 的 wave32 layout、又跑到手机 GPU 上调了一些 compute shader才发现“在不同架构上长什么样”这个问题其实是在问硬件设计者对“计算主体”和“数据复用”的底层假设。这篇文章就当是把这个大问题拆成几个小问题系统盘一遍。先说结论所谓寄存器 tiling核心目标是让每个线程用尽量少的 load 指令把从 shared memory 或 global memory 取得的数据尽可能多地在寄存器里复用。寄存器访问延迟通常只有几个周期而 shared memory 需要几十个周期global 直接上百如果能把一块数据留在寄存器里被多次计算整体吞吐会非常可观。但每个架构能留在寄存器里的数据形状、每个线程能分到的数量、以及矩阵指令希望数据以什么格式出现差别极大。1.1 从“搬数”说起为什么先有 global、shared才轮到寄存器我们经常听说的 tiling 一般默认指的是 shared memory tiling。比如 GEMM 里把 4096x4096 的大矩阵切成 128x128 的 block再让一个 thread block 内的线程协作把这 128x128 的块加载到 shared memory然后在共享内存里反复读取、计算。这确实是最重要的宏观分块因为它决定了线程块之间的数据独立性也决定了访存复用率。但 shared memory tiling 做完之后每个线程依然需要从 shared memory 里取数。如果每个线程每算一个输出元素都去 shared memory 取一次 A 和 B那 shared memory 的带宽马上就会成为瓶颈而且会有大量浪费。所以真实实现的下一步是把每个线程负责的“微块”micro-tile继续切分让微块数据落在寄存器里内层循环完全基于寄存器读写。举例来说假设线程块负责 128x128block 内有 256 个线程那么每个线程负责 8x8 的 C 块。如果没有寄存器 tiling直接写成外层遍历 8x8、内层遍历 K那每个输出点都会从 shared memory 读 K 次 A 和 K 次 B而 8x8 的寄存器 tiling 会把属于这个线程的 8 行 A、8 列 B 按需放入寄存器K 循环内部只需从 shared memory 读新的 K 切片旧的 K 切片完全留在寄存器里反复参与 64 次乘加。所以寄存器 tiling 不是和 shared tiling 并列的另一个方法而是 shared tiling 底下必须再补的一层。两者一起构成了常见的三层分块global 分块决定 block 层面调度shared 分块决定 block 内协作register 分块决定单线程计算密度。1.2 寄存器 tiling 的定位真正决定指令级并行上限的一层寄存器 tiling 最本质的价值是它直接决定了一个线程一个循环里能做多少次独立的乘加运算也就是指令级并行ILP的上限。编译器只要能把这些独立运算调度开ALU 就会一直满载。而这个上限其实在写代码时就定死了如果你的寄存器 tile 设计得让同一时刻可并行的 mma矩阵乘加指令太少后面再怎么调优也救不回来。除此之外寄存器 tiling 还对 shared memory 的 bank conflict 有间接影响。你从 shared memory 读数据时如果想避免 bank conflict要设计成连续线程访问连续地址而由于每个线程读的数据最终要落到自己的寄存器 fragment 中线程与地址之间的映射关系就是由寄存器 tiling 的 layout 决定的。很多做 GEMM 优化的人都有过这种经历shared 阶段明明没冲突进了寄存器阶段换了一种 fragment 分配回头一看 shared load 又冲突了。所以这一层看起来是寄存器的事实际上和整个数据通路都耦合。看不同架构的 register tiling本质上是在看架构设计师怎么权衡线程数量、寄存器容量、矩阵指令宽度这三者的关系。下面进入正题。2. NVIDIA GPUwarp 级 micro-tile 的标准答案NVIDIA GPU 是大多数 AI Infra 工程师接触到的第一块并行计算硬件也是寄存器 tiling 感知最强的地方。以 Ampere 架构的 GA100 为例每个 SM 有 64K 个 32 位寄存器换算成字节是 256KB 寄存器堆跟 L1/Shared memory 容量差不多大。每个 thread 最多可以用 255 个寄存器每个 warp 是 32 个线程。对 GEMM 来说一个 warp 负责一个 warp tile比如 64x64然后 warp tile 内部再切成线程级 micro-tile通常形状是 8x8也有 16x8、32x8 之类的变体。这里有个关键点NVIDIA 从 Turing 开始引入了 HMMA 指令Ampere 又加了三代 tf32、bf16、fp16 的 mma 指令它的硬件矩阵单元天然希望数据以一种固定 fragment 布局出现在寄存器中。你自己把寄存器当数组来回读写当然也可以但除非用上 mma 类指令否则没法吃到 tensor core 的吞吐。而 mma 指令的 fragment 布局就是 NVIDIA 对寄存器 tiling 的“官方答案”。2.1 硬件基础寄存器堆、warp 与 mma/wgmma在写任何 CUDA 或 PTX 之前先把硬件事实摆清楚。SM 里的寄存器堆按 lane 组织每个 lane 有自己独立的一组寄存器文件。对 warp 内的第 L 个 lane 来说它只能访问自己名下的寄存器不能直接读其他 lane 的寄存器。这一点很关键你说的“寄存器 tiling 数据在寄存器之间流转”实际上只发生在同一个 lane 内部lane 之间交换数据必须通过 shared memory 或 shuffle。Tensor Core 的 mma 指令相当于把寄存器 tiling 逻辑固化成了硬件微码。以 Ampere 常见的 m16n8k8fp16为例一条指令由 warp 内 32 个 lane 协作完成 16x8x8 的矩阵乘加。指令规定A 矩阵是 16x8按 k 维 8 拆开后每个 lane 持有 A 的 4 个元素表现为寄存器称为 a-fragmentB 矩阵是 8x8每个 lane 持有 B 的 2 个元素b-fragmentC/D 是 16x8每个 lane 持有 4 个元素c/d-fragment。四个片段合在一起才是这个 warp 在当前 k 切片的寄存器 tile 全貌。Hopper 上进一步出现了 wgmma把 mma 的协作范围从 warp 扩大到 warpgroup四个 warp。wgmma 的一条指令可以同时指定 A 矩阵在 shared memory 中、B 矩阵在寄存器中的复杂布局也可以反过来。这让寄存器 tiling 的形态变得更加多样本质上是在用硬件调度器教你A 放 shared 或寄存器、B 放寄存器或 shared、累加器放寄存器可以按需组合。2.2 reg tile 长什么样以 CUTLASS 中常见的 128x128 block、256 线程8 warps为例每个 warp 分到 64x64 的 warp tile一个 warp 内再把 64x64 的 C 矩阵切成 8x8 的 micro-tile于是共有 8x864 份每份对应一个 lane恰好覆盖 32 个线程加循环展开的两份。每个线程持有这个 8x8 micro-tile 里的 8 个 float也就是 8 个 32 位寄存器。从内存布局上看A 矩阵的寄存器 fragment 通常是按 k 维度连续的。假设每轮从 shared memory 加载 16x16 的 A 块warp 中每个 lane 会持有若干 a-fragment这些 fragment 在寄存器中的编号是连续的方便编译器做双缓冲和乱序调度。B 矩阵类似。C 矩阵的 fragment 严格按 (row, col) 分布一个线程的 8 个值往往不是原始矩阵里连续的 8 个点而是按固定 stride 取到的。这就是官网文档里那张密密麻麻的“lane 编号占位图”的由来。2.3 代码实现示意下面用简化的伪代码展示 Ampere 上 m16n8k8 的寄存器 tiling 流程。假设我们已经在 shared memory 里有 A_tile 和 B_tile每个线程需要把随后 8 个 k 步的数据准备到寄存器里。// warp 内每 lane 持有 a_frag[4], b_frag[2], c_frag[4] float a_frag[4]; float b_frag[2]; float c_frag[4] {0, 0, 0, 0}; // k 维度迭代这里每步处理 k 8 for (int k0 0; k0 K_TILE; k0 8) { // 从 shared memory 加载 A 的 16x8 分片到寄存器 // 实际地址由 lane_id 和 mma 官方 layout 决定 load_a_fragment(a_frag, A_smem, k0); load_b_fragment(b_frag, B_smem, k0); asm volatile( mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n : f(c_frag[0]), f(c_frag[1]), f(c_frag[2]), f(c_frag[3]) : r(a_frag[0]), r(a_frag[1]), r(a_frag[2]), r(a_frag[3]), r(b_frag[0]), r(b_frag[1])); }这里 a_frag 的 4 个寄存器、b_frag 的 2 个寄存器就是当前 k 切片对应的寄存器 tile。它们不像数组那样在内存里连续而是被 mma 指令的逻辑硬编码成一种分布式布局。如果你用传统 SIMT 风格手写乘加循环当然也可以在寄存器里做 8x8 的 FMA但那样就吃不到 tensor core 了。2.4 ptxas 角度寄存器分配与 occupancy 权衡NVIDIA 这一侧寄存器 tiling 的形状最后会反映在编译结果里。用ptxas -v可以看到每个 kernel 用了多少寄存器有没有 spill。一个 128x128 block、8x8 micro-tile 的 HGEMM kernel通常会用到每个线程 168~232 个寄存器这正好卡在占用率比较舒服的区间一个 SM 上能同时驻留的 warp 数 65536 / (寄存器数 * 32)。如果每线程 224 个寄存器那么 65536 / (224*32) 9.14向下取整 9 个 warp因为 warp 调度器不支持 9 个 warp 均匀分配到 4 个 scheduler实际会退到 8 个 warpSM 利用率会明显下降。这里有个 tips不要一味追求大 tile。寄存器 tile 从 8x8 放大到 16x8C fragment 的寄存器数翻倍占用率立刻跳水。对于纯计算密集型的 GEMM你当然希望 tile 大一点好复用数据但对于带宽敏感或者需要高并发隐藏延迟的 kernel适当的 register tiling 加高占用率效果反而更好。具体得靠 profiling没有统一答案。3. AMD CDNAwave 级 tile 与 VGPR/SGPR 的分工AMD 的 GPU 和 NVIDIA 走了不完全相同的路。CDNA 架构MI100、MI200、MI300 系列使用 wavefront 作为调度单位一个 wavefront 是 64 个 lane也可以配置成 wave32RDNA 后原生支持 wave32CDNA 也提供选项。寄存器堆分为两类Vector GPRVGPR和 Scalar GPRSGPR。VGPR 是每个 lane 都有自己的副本存逐线程数据SGPR 是所有 lane 共享一份存常量、地址和统一数据。这个双寄存器堆设计直接影响了寄存器 tiling 的形态。在 AMD 的 MFMA 指令matrix core 指令类似 NVIDIA 的 mma里A、B、C 的 fragment 分布在 VGPR 中而地址、步长、即时数等“元数据”放在 SGPR 中。SGPR 的有效使用能减少 VGPR 的浪费因为你不用在每个 lane 里重复保存相同的地址。这是 CDNA 架构在编译器约束上比 NVIDIA 更精细的地方。3.1 CDNA 有别于 NVIDIA 的那几件事第一个区别是 wave64 的 lane 数。同样一个 16x16 的寄存器 tilewave64 里会被切成更多份每个 lane 持有的寄存器更少wave32 下每个 lane 持有更多元素指令数更少但寄存器压力更大。AMD 的编译器经常需要在 wave32 和 wave64 之间做选择CDNA 的矩阵核心在 wave64 下支持更多指令形状而 wave32 在纯 scalar 逻辑或低占用场景下更高效。第二个区别是指令的 fragment 尺寸选择。CDNA 的 MFMA 支持 m16n16k16、m32n32k8、m64n64k4 之类的组合m64n64k4 对每个 lane 的寄存器要求更大适合大 K 复用m16n16k16 则更均衡。这种按需选择的灵活性让寄存器 tiling 变成了一个真正的参数空间而不是固定模板。实际调优时我们往往先用 m16n16k16 跑通功能再切换到 m32n32k8 去看 latency 和 occupancy 的平衡。第三个区别在 shared memory 的命名上AMD 叫 LDSLocal Data Share它的 bank 机制和 NVIDIA 类似但访问竞争规则细节不同。由于寄存器 tiling 决定了 LDS 读取的数据如何分配到各 lane同一个 16x16 块用不同的 fragment 布局读入 LDS冲突率可能差几倍。AMD 的 rocBLAS 和 Composable Kernel 内部都维护了很多针对不同 CDNA 芯片的 layout 变体背后的原因就在这。3.2 一个 wave32 的 reg tile 示例假设我们用 wave32 做一个基于 MFMA 的 HGEMM选择 m16n16k16 的指令形状。每个 wave 负责 16x16 的 C tileA 是 16x16 的 tileB 是 16x16 的 tile。按 MFMA 的文档A 的 16x16 会被拆成若干 32 位 VGPR每个 lane 拿 8 个16x16 / 32 8B 同样每个 lane 拿 8 个C 的 16x16 累加器每个 lane 拿 4 个 VGPR这是由 16x16x4 字节 / 32 lane 32 字节 8 个 32 位寄存器不对准确说 C 是 16x16 的 f32即 1024 字节分给 32 个 lane每个 lane 32 字节即 8 个 VGPR。让我重新算一下。A 矩阵是 16x16 的 fp16共 512 字节每个 lane 拿到 16 字节即 4 个 16 位元素吗MFMA 的实际文档里m16n16k16 的 A fragment 每个 lane 持有 4 个 fp16B 是 4 个 fp16C 是 8 个 f32。这个具体数字因芯片世代而异读者不一定要背真正重要的是理解布局不是连续排列。写伪代码的话和 NVIDIA 的例子结构类似只是指令助记符从mma.sync.aligned变成了v_mfma_f32_16x16x16_f16。为了完整可以罗列一下伪代码// 假设 wave32v_mfma_f32_16x16x16_f16 // 每个 lane 的 VGPR // a_vgpr[4]A fragment // b_vgpr[4]B fragment // c_vgpr[8]C/D 累加 for (int k0 0; k0 K_TILE; k0 16) { lds_load_a_to_vgpr(a_vgpr, ...); lds_load_b_to_vgpr(b_vgpr, ...); asm volatile( v_mfma_f32_16x16x16_f16 %0, %1, %2, %0 : v(c_vgpr) : v(a_vgpr), v(b_vgpr)); }注意这里 C 累加器必须用 “” 约束因为 MFMA 是累加操作。在 CDNA2 之后MFMA 还支持双发行和 LDS 直读某些组合下 B 矩阵可以直接在指令里引用 LDS 地址减少一次 VGPR 中转。这种“寄存器 tiling 与 LDS 读取融合”的模式是 NVIDIA 在 Blackwell 之前都不太强调的设计。3.3 换到 RDNA/移动端后的变化RDNA 系列虽然也使用波前但主要面向图形和消费级计算矩阵核心能力弱很多。RDNA3 加入了 WMMAWave Matrix Multiply-Accumulate指令支持 16x16x16 的 fp16 计算但寄存器文件仍沿用 VGPR/SGPR 双 bank 机制。在 RDNA 上做寄存器 tiling 优化时要考虑的不仅是 matrix core 指令还包括 wave32 下 VGPR 数量对 shader occupancy 的影响。移动端 RDNA比如 RDNA3 的 iGPU往往没有独立显存带宽和 L2 都很有限寄存器 tiling 更需要和 LDS 配合尽量让每个 wave 在寄存器里完成多轮复用减少对全局内存和二级缓存的依赖。这部分实战中比较痛苦的地方在于移动端驱动对寄存器分配的干预程度较高你写的num_SGPR/num_VGPR约束不一定被严格尊重需要实测。4. 移动端 GPU寄存器被当作“向量寄存器组”用的一层手机 GPU 是 AI Infra 工程师经常忽略但实际非常重要的场景。端侧 AI 推理跑在 Adreno、Mali、PowerVR 这类 GPU 上它们没有像 NVIDIA CDNA 那样被明确文档化的mma指令体系Adreno 有一些uGQR矩阵指令Arm 也加入了smmla等但寄存器 tiling 同样成立只是表现形式更接近“把寄存器当成 SIMD 向量寄存器组来用”。这类 GPU 的底层执行模型是“一个着色器程序在多个 SIMD 通道上并行”。每个通道有自己的寄存器但编译器通常把它们组织成向量寄存器。你在 Adreno 的 compute shader 里看到一个 128 位寄存器r0.xyzw可以理解成同时容纳 4 个 fp32或者 8 个 fp16。寄存器 tiling 在这里就表现为每个线程或者每个像素/每个 work-item在自己的向量寄存器里维护一个小矩阵块比如 4x4 或 8x4然后一次 ALU 操作能同时对 4 个通道做 FMA。这个形态看起来没有 NVIDIA 那么复杂但约束反而更极端。移动端 GPU 的寄存器文件往往比桌面级小很多Adreno 某些代际上每个核心的寄存器总量一体超限就 spill 到 local memory而移动 GPU 的 local memory 带宽和延迟都远不如桌面 GPU 的 shared memoryspill 代价极其昂贵可能直接让 kernel 从 2ms 变 20ms。所以移动端的 register tiling 优化更像是在“如何用最少的寄存器表达最大的数据复用率”。4.1 Adreno 的 scalar/vector 双寄存器堆Adreno 的架构从 6xx 到 8xx 一直在演进但大体保留了 scalar 和 vector 两套寄存器。scalar 寄存器存每个线程都相同的值比如循环计数、统一缓冲区的参数、计算地址的公共部分vector 寄存器存随线程变化的数据。这个设计原理上和 AMD 的 SGPR/VGPR 是同构的。在 Adreno 上写 compute shader编译器自动把 uniform 数据和 per-thread 数据分开分配。如果你能手动保证某些矩阵 tile 的行列步长是一致的编译器可以把地址计算尽量下沉到 scalar 寄存器减少 vector 寄存器的浪费。实测中很多 kernel 的性能瓶颈并不在 ALU而在寄存器压力过大导致 occupancy 太低或者地址计算占用了大量 vector 寄存器。Qualcomm 后来引入了 uGQR 这类硬件矩阵指令目的就是让寄存器 tiling 的 fragment 布局有硬件标准避免编译器瞎猜。但这类指令主要服务于特定精度和特定大小比如 8x8x8实际开发中如果用不到仍要靠手写向量寄存器块来优化。4.2 Mali Valhall 的 reg tile 实践Arm Mali Valhall 架构G77 之后的架构代号和上一代 Bifrost 有本质区别完全转向了基于 warp 的标量/向量统一执行模型每个 warp 有 16 个 lane某些实现里是 8 个支持 fp16 的 ALU 对半分。Mali 的寄存器堆是按执行引擎Execution EngineEE组织的每个 EE 有固定数量的寄存器如果超出就 spill。在 Mali 上做寄存器 tiling 时我发现最合适的是 8x8 或 16x8 的 fp16 tile。原因很简单Mali 的 SIMD 宽度是 128 位一次性可以处理 8 个 fp16。假如每个 lane 持有 8x8 tile 里的一行那么 8 行数据正好填满 8 个 128 位向量寄存器加上 C 的 8 行又是 8 个寄存器A/B 各若干总计大约 24~32 个向量寄存器。这种情况下 occupancy 还能保持不错的水平。如果把 tile 加到 16x16寄存器会飙到 60 个以上占用率立刻崩掉。Mali 的编译器和驱动对寄存器分配也有很强的主导权它可能会在你意想不到的地方插入补丁指令patching。一个常见坑是你精心安排的寄存器 tile 会被编译器重排导致数据布局变化性能反而下降。我的经验是在这种平台上不要过度追求寄存器内部布局除非你愿意阅读反汇编和驱动调度日志。4.3 PowerVR 和其他移动架构的备注PowerVR 架构的 GPU在部分 SoC 中出现使用的是统一存储和基于 tile 的渲染方式计算着色器也沿用类似的寄存器组织。它的 TBDR 设计虽然对渲染很友好但在通用计算上寄存器 tiling 的发挥空间比较受限因为线程调度被硬件抽象得很完整你很难通过手工布局寄存器来压榨额外性能。一般的建议是用厂商提供的 compute library比如 IMG 的 OpenCL SDK的推荐模式不要自己手动搞复杂 tile。顺便说一句鸿蒙或国产 SoC 里的 GPU 很多基于 Mali 或 PowerVR 的 IP也有部分自研架构但寄存器相关的开发接口基本都是 OpenCL / Vulkan compute。指令集形态虽有差异寄存器 tiling 的思考方式完全一致先理解每个 lane 有几个可用的 32 位寄存器再决定 micro-tile 形状。5. CPU 侧向量化视角下的寄存器 tiling我们通常说寄存器 tiling 都是针对 GPU但其实 CPU 上的 SIMD 向量化也可以从这个角度理解。AVX-512 有 32 个 512 位寄存器ARM SVE 有可变的向量长度寄存器NEON 有 32 个 128 位寄存器。这些寄存器和 GPU 的 lane 寄存器不同它更像“一个寄存器就是一排小抽屉”SIMD 指令一次操作一排。所以在 CPU 上做寄存器 tiling本质是“把一个小矩阵块的每一行或每一列映射到一组向量寄存器里”。比如对 8x8 的 fp32 矩阵乘你用 AVX-512 的话可以一个向量寄存器装一行 8 个元素C tile 就是 8 个向量寄存器A 的一行广播到 8 个向量寄存器和 B 的 8 列逐列乘加。这就是典型的寄存器 tiling不经过内存完全在寄存器里完成 8 个向量寄存器和 8 个向量寄存器之间的外积累加。这个套路和 GPU 的区别在于GPU 的寄存器 tiling 往往要依赖硬件矩阵指令而 CPU 的寄存器 tiling 可以直接写成普通的向量 FMA。编译器GCC、Clang、MSVC通常会自动向量化一部分但高级优化还是得手写 intrinsics。AI Infra 里的推理引擎优化对 CPU 后端也常做类似工作比如 oneDNN 的 jit kernel 就会把 4x16 或 8x16 的 micro-tile 硬编码到寄存器和指令调度里。5.1 SIMD 寄存器也是 tiling用矩阵乘法来举例最直观。假设你要算一个 4x4 的 C A x B把 A 的每一行看成一组 SIMD 数据。最简单的 SIMD 方式C 的每一行用一个向量寄存器当累加器内层遍历 K 时把 A 元素的标量广播到所有 lane然后乘 B 的一行向量加到 C 行向量上。这样每个 C 行向量就相当于一个寄存器 tile 的行。如果进一步优化可以把 A 的连续 4 行同时保留在 4 个向量寄存器里B 的 4 列也在 4 个向量寄存器里然后做 4x4 个向量 FMA。这听起来和 GPU 的寄存器 tile 一模一样。区别只在于CPU 的“lane”和“寄存器”是同一个概念的不同表现一个 512 位寄存器相当于 16 个 fp32 lane和 GPU 中一个 lane 拥有若干寄存器恰好是反过来的。5.2 ARM SVE 与 x86 AVX-512 的例子ARM SVE 的一个特点是向量长度可以支持 128 到 2048 位软件不用硬编码长度这使得寄存器 tiling 的代码需要写成与 VL 无关的形式。比如你要实现一个基于外积的 GEMM micro-kernel写代码时不能假设一个向量寄存器能装 16 个 fp32只能用svcntw()得到当前机器的 lane 数再据此切分 tile。好处是同一份代码可以在不同 SVE 机型上获得不同性能坏处是编译器优化压力很大。x86 的 AVX-512 固定是 16 个 fp32 per register写代码时往往直接硬编码。我常用的一个 fp32 GEMM 寄存器 tile 是 8x16C 的 8 行每行占一个 zmm 寄存器16 个 fp32这样 C tile 一共 8 个 zmm。A 的一行 16 个元素可以打包进一个 zmmB 的 8 列也各占一个 zmm。内层循环时A 行广播到 8 个寄存器分别乘 B 的各列累加进对应的 8 个 C 行寄存器。一轮 k 步消耗 8 4 8 之类的寄存器不做特殊复杂处理也能保持 28~32 个寄存器以内不会造成溢出。5.3 与 GPU 的一致性和差异CPU 和 GPU 寄存器 tiling 的共性在于都是“数据复用优先减少访存次数”。差异在于数据分布模型GPU 是“多个线程各拿一小块”CPU 是“单线程用 SIMD 通道构成一个小块”。这两种模型各有优劣。GPU 的 tiling 适合大规模并行和硬件矩阵指令CPU 的 tiling 适合延迟敏感、并行度有限的场景。理解这个一致性后你会发现一个非常有用的迁移思路在 GPU 上研发的 tiling 策略可以先在 CPU SIMD 上做快速原型验证因为 CPU 上调试寄存器布局、观察 spill 都比 GPU 方便得多。等逻辑正确了再移植到 GPU 上做 mma / MFMA 适配。6. 一个完整例子HGEMM 里手工设计 reg tile这一节我们用实际可操作的步骤把之前的抽象概念落到一个 HGEMM kernel 上。假设目标平台是 NVIDIA A100Ampere输入的矩阵是 fp16输出累加成 fp32K 维度较大。我们不依赖 cuBLAS自己写一个 register tiling 的示例目的在于看清每一步的形状选择。6.1 从 GEMM shape 往下拆一个典型的 A100 HGEMM 分块流程如下全局分块每个 thread block 计算 128x128 的 C tile。block 内线程数256 线程组织成 8 个 warp每个 warp 负责 64x64 的 C tile。warp 内每个 warp 的 32 个 lane 共同负责一个 64x64 的 warp tile。寄存器层通常每个 thread 负责 8x8 的 micro-tile也就是 C 的 64 个元素。这里的 8x8 不是拍脑袋定的它和 mma 指令的 m16n8k8 形状有对应关系。一个 lane 同时会参与若干个 mma 片段经过两个 mma 步骤才能覆盖 8x8 的 C tile。如果改成 16x8 的 micro-tile寄存器占用会翻倍到 16 个 f32但同一个 lane 需要参与的 mma 次数变多指令级并行更高适合 K 较短、block 数多的场景。实际项目里A100 上用 8x8 或 16x8 都有人做具体以 profile 为准。寄存器预算方面8x8 的 C tile 需要 8 个 fp32 寄存器循环里同时要保留 A 的 16x8 fragment 和 B 的 8x8 fragment以 fp16 形式各自折合若干 32 位寄存器再加上地址、指针、循环变量整体大概 100~128 个寄存器。这个量级能保证 SM 上的 warp 数在 10 个左右ALU 有足够的并发来掩盖 shared memory 延迟。6.2 选 micro-tile 形状与 fragment 分配选 micro-tile 形状时要考虑两个指标数据复用率和寄存器占用。假设每个线程负责 8x8 的 C tile那么在 k 维内层循环里这个线程需要同时持有 A 的 8 个元素和 B 的 8 个元素每个 k 步加上 C 的 8 个累加器一共 24 个寄存器左右。这是一份低压力配置适合启动的很早、或者对 occupancy 要求很高的 kernel。如果选 16x8 的 C tile就需要 A 的 16 个元素、B 的 8 个元素、C 的 16 个累加器接近 40 个 32 位寄存器属于中等压力配置。再往上选 32x8C 本身就 32 个寄存器加上 A/B 很容易突破 80 个除非对 L2 命中率有极大自信否则不建议在 A100 上这么干。Fragment 分配则取决于你用什么指令。如果手写 PTX 的 mma.m16n8k8A fragment 的 4 个寄存器编号在 lane 间是错开的编译器会帮你插入合适的从 shared memory 到寄存器的加载序列。如果你用 CUTLASS 的MmaTensorOpfragment 分配由模板参数里的MmaPolicy决定内部的TiledMma会把每个线程负责的元素下标映射得清清楚楚。我个人的建议是一开始不要自己手工映射 lane 和寄存器先用 CUTLASS 的默认配置跑通再逐步魔改。6.3 主循环体k 循环中的 reg tiling假设 shared memory 里已经放好了 64x64 的 A tile 和 B tile寄存器 tiling 的主循环体大致如下// 每个线程注册 8 个 C 累加器 float c[8] {0}; // 以 16 为步长在 k 维上迭代对应一个 shared tile 切片 for (int k 0; k 64; k 16) { // 从 shared memory 取 A 的 16x8 到寄存器 a_frag float a_frag[8]; // 简化的描述实际可能拆成 4 个寄存器对 load_a_into_reg(a_frag, smem_a, k); // 从 shared memory 取 B 的 16x8 到寄存器 b_frag float b_frag[8]; load_b_into_reg(b_frag, smem_b, k); // 4 次 mma.m16n8k8 完成 8x8 的累加 // 每次 mma 指令处理 16x8x8这里需要按布局重组 mma_16x8x8(c[0], a_frag, b_frag); mma_16x8x8(c[1], a_frag, b_frag); }这段代码忽略了很多索引细节但结构对了A/B 各自从 shared 载入寄存器然后连续做 4 次 mma。这里的“连续”是重点它让寄存器中的 a_frag、b_frag 被复用多次而不是每次 mma 都重新读 shared。这也是寄存器 tiling 相对朴素实现的最高价值点。实际调优中还会在load_a_into_reg和 mma 之间做软件流水线也就是提前预取下一轮 k 的 A/B 到另一组寄存器让当前 mma 和下一次 shared load 并行。这样寄存器占用会上升但延迟能被隐藏掉吞吐往往更好。CUTLASS 的PipelineTmaAsync或传统cp.async双缓冲方案本质上就是在为这个目标服务。6.4 完整可参考的架构选型参数表这里整理一份我在不同架构上做 fp16 GEMM 时经常参考的初始参数表。注意只是起点不代表最优真实场景请按 profiling 调整。架构典型线程模型推荐 micro-tile寄存器类型核心指令/机制NVIDIA Turing/Amperewarp 32 lanes8x8 / 16x8普通 32 位寄存器mma.sync / wgmmaNVIDIA Hopperwarpgroup 128 lanes16x16 / 32x16普通 32 位寄存器wgmma TMEMAMD CDNA2/3wave32 / wave6416x16 / 32x32VGPR SGPRv_mfma / v_wmmaQualcomm Adreno 7xx每 SIMD 通道 thread8x8 / 8x16 fp16scalar vector regs手写向量 FMA / uGQRArm Mali Valhallwarp 16 lanes8x8 / 16x8vector regs手写向量 FMA / smmlax86 AVX-512单线程多 lane8x16 fp32zmm 向量寄存器FMA / 广播ARM SVE单线程多 laneVL 自适应 tileSVE 向量寄存器FMLA / 外积这张表不是让你死记硬背而是提醒一个问题不同架构的 micro-tile 形状差异极大这背后是寄存器堆容量、lane 数和指令宽度共同作用的结果。在 NVIDIA 上顺手的 16x8放到 Mali 可能因为 lane 数不同而完全发挥不出来反过来也一样。7. 常见问题与排查技巧实录做寄存器 tiling 优化时最头疼的不是理论不懂而是“编译器不按我想的来”。下面几个问题是我自己在实践中反复遇到的基本每天都在调优群里被问一遍。7.1 寄存器溢出最容易翻车的地方寄存器溢出spill是寄存器 tiling 的头号敌人。如果你的 kernel 编译后出现local memory访问大概率是寄存器压力超过了硬件上限。NVIDIA 上可以用ptxas -v看 spill 数量AMD 上可以用ROCm的编译日志或者rocprof看 VGPR 分配移动端则要看驱动或 profiling 工具的 register usage 指标。我的排查顺序是先看寄存器占用再看 occupancy再看实际运行时间。如果寄存器占用很高但 occupancy 没跌可以先跑一下如果 spill 已经出现必须降 tile 或者减少双缓冲级数。很多时候减少一个循环展开因子就能避免 spill性能不降反升。记住一个原则宁可让 ALU 动用不完美的指令序列也不要让数据跑到 local memory 去因为 spill 的代价常常是数倍的 memory 访问延迟。7.2 occupancy 掉下来tile 越大未必越好常有人觉得寄存器 tile 越大数据复用率越高性能一定越好。但寄存器总量是物理限制。NVIDIA A100 的 SM 只有 65536 个 32 位寄存器如果每个线程用 256 个一个 SM 只能放下 8 个 warp而一个 SM 的四个调度器理想情况下需要至少 16 个 warp 才能保持满负载。一旦 warp 数不足访存延迟就藏不住kernel 照样跑不快。AMD CDNA 同样有 VGPR 上限的约束。对于需要用 wave64 的 kernel如果每个 lane 分到的 VGPR 超过 128wave 数就会下降。这时候可以尝试切 wave32虽然一个 wave 的线程数少了一半但同一时刻可调度的 wave 数量可能翻倍总体吞吐反而提升。移动端更是如此Adreno 和 Mali 的寄存器堆更小过大的 tile 会让 occupancy 暴跌导致 latency hiding 完全失效。7.3 shared memory bank conflict 仍然存在的寄存器 tiling 并不能解决所有 shared memory 冲突。很多人以为数据进了寄存器就万事大吉实际上你从 shared memory 到寄存器的那次 load 本身就要面对 bank conflict。而这次 load 的地址模式直接由 register tiling 的 layout 决定。比如 NVIDIA 上常见的“跨步读取”导致同一 bank 被多个 lane 争抢就是 fragment 分配不合适的结果。排查方法是看 profiling 里的 shared load bank conflict 指标。如果冲突率很高可以尝试改变 micro-tile 的方向比如把 C tile 从 16x8 改成 8x16很多情况下 A/B 的加载模式会跟着改变冲突就自然消掉了。AMD 的 LDS 也有类似问题不过它的 bank 数和 conflict 检测方式和 NVIDIA 不完全一致需要单独看文档。7.4 工具链怎么看 reg tiling 效果实话说最直接的手段是看 SASS/汇编。NVIDIA 上用cuobjdump -sass看每条指令用哪些寄存器AMD 上用LLVM MC反汇编可以看 VGPR/SGPR 使用移动端则比较麻烦但多数厂商 profiling 工具会输出 register usage per thread。另一个实用技巧是故意把 tile 改成极端值比如 1x1 和 64x64分别测时间从两端逼出寄存器 tiling 的贡献就能判断当前 kernel 到底卡在寄存器还是卡在带宽。Nsight Compute 的Speed Of Light页面和Warp State Statistics面板对 NVIDIA 非常有用能看到 stall 原因是什么AMD 的rocprof加--stats也可以给出 instruction mix 和 VGPR 分配。移动端通常用厂商的 GPU profiler比如 Qualcomm Snapdragon Profiler 或 Arm Streamline观察 shader core 的 ALU 利用率和 register bank 压力。7.5 问题速查表下面是我自己整理的快速定位表照着排查往往能省很多时间。现象可能原因下一步动作kernel 运行时间远超预期register spill降 micro-tile / 减少双缓冲级数profiling 显示 memory stall 很高occupancy 不足降低寄存器占用提高 warp/wave 数shared load bank conflict 高fragment 布局和地址映射不匹配换 tile 方向检查 lane 到地址的映射编译器没生成 mma/MFMA数据类型或 fragment 布局不符合指令要求检查指令约束必要时写 PTX/intrinsic改了 tile 大小性能没什么变化瓶颈不在寄存器层改测 global/shared tiling或 L2 命中率这张表没有涵盖所有情况但覆盖了 80% 的初调问题。如果你遇到的现象不在这张表里建议回去看汇编一条STL或V_ADD出现的位置往往能告诉你编译器到底在寄存器层面做了什么奇怪的事。最后说一点个人体会。寄存器 tiling 这东西看十篇文档不如动手改一次 kernel。我见过很多人卡在“寄存器 tiling 和 shared tiling 的区别”这种概念问题上其实只要先写一个只做 shared tiling 的版本再写一个加了 register tiling 的版本两边跑一下性能立刻就能感受到这一层的价值。而且不同架构的差异并不可怕它们背后的原则其实是同一个把最热的数据放在最靠近 ALU 的存储里让它在被丢弃前尽可能地多算几次。搞明白这一点换架构只是换 API 和 layout 表而已。
 方法详解：从对象数据转换到 OOXML 报错排查)
