最近在看 CMU 15-442 的 Hardware Acceleration。课件翻到 Register Tiling,一个 2×3 output tile 的例子。我盯着那页 slide 看了一会儿——Register Tiling 这个词不陌生,但要我说清楚它和 Cache Tiling 的区别,以及它们到底在优化什么,我发现自己的理解一直停在"把矩阵切成小块"这个表面。
之前写过 Arithmetic Intensity 和 Compute/Memory Bound,也讨论过 GPU 存储层级 和 FlashAttention 的分块策略。但那些都在 GPU 场景下。CPU 上的 GEMM 优化反而没有系统地想过。
这次想从最简单的三层循环出发,把这件事想清楚。
朴素的三层循环#
矩阵乘法 \(C_{ij} = \sum_k A_{ik} B_{kj}\),最直接的实现就是三层循环:
1for (int i = 0; i < M; i++)
2 for (int j = 0; j < N; j++)
3 for (int k = 0; k < K; k++)
4 C[i][j] += A[i][k] * B[k][j];每次内层循环迭代计算一个 \(C_{ij}\) 的一次累加。
想一下数据访问模式:\(A_{i0}\) 这个元素,在计算 \(C_{i0}, C_{i1}, \ldots, C_{i,N-1}\) 时都会被用到——被访问 \(N\) 次。类似地,\(B_{0j}\) 在计算 \(C_{0j}, C_{1j}, \ldots, C_{M-1,j}\) 时也都会用到,被访问 \(M\) 次。
这些重复访问是矩阵乘法的数学结构决定的。A 的每一行要和 B 的每一列配对,每个元素都会参与多次计算。
现代 CPU 有 Cache,重复访问同一个地址不一定每次都从 DRAM 读。但有两件事需要注意。第一,Cache Hit 不是免费的——L1 Cache 的延迟大约 4-5 个时钟周期,和 Register 的 0-1 个周期比还是有差距。第二,Cache 容量有限,当矩阵足够大,之前缓存的数据可能在被再次使用之前就已经被驱逐了。
于是问题就变成了:能不能在数据还在快速存储层级的时候,尽量多地利用它?
一次加载,多次计算#
Register Tiling 的想法:不是每次只算一个 \(C_{ij}\),而是同时维护一小块 output——比如 2 行 3 列,6 个输出元素——的累加器,全放在寄存器里。
每次 K 循环迭代的数据流:
- 从 A 加载 2 个元素:\(A_{ik}\) 和 \(A_{i+1,k}\)
- 从 B 加载 3 个元素:\(B_{kj}\),\(B_{k,j+1}\),\(B_{k,j+2}\)
- 执行 6 次 FMA
同一个 \(A_{ik}\) 参与了 3 次计算(分别和 3 个 B 元素相乘),同一个 \(B_{kj}\) 参与了 2 次计算(分别和 2 个 A 元素相乘)。加载 5 个值,做了 6 次乘加。
代码大概是这样:
1float c00 = 0, c01 = 0, c02 = 0;
2float c10 = 0, c11 = 0, c12 = 0;
3for (int k = 0; k < K; k++) {
4 float a0 = A[i][k], a1 = A[i+1][k];
5 float b0 = B[k][j], b1 = B[k][j+1], b2 = B[k][j+2];
6 c00 += a0*b0; c01 += a0*b1; c02 += a0*b2;
7 c10 += a1*b0; c11 += a1*b1; c12 += a1*b2;
8}对比一下:朴素版本每次 K 迭代加载 1 个 A 和 1 个 B,做 1 次乘加。这个版本加载 2+3=5 个值,做 6 次乘加。如果 tile 更大——比如 4×4——加载 4+4=8 个值,做 16 次乘加,复用率更高。
从另一个角度看,每次 K 迭代实际上是一个 outer product:一个 2 维列向量和一个 3 维行向量的外积,产生一个 2×3 的矩阵。整个 K 循环就是把 K 个外积累加起来。
Register Tiling 不减少计算量——FLOPs 完全没变——它通过重新组织计算,让加载到寄存器的数据在被覆盖之前参与更多的运算。
这和 Loop Unrolling 是一回事吗#
这个问题我想了一会儿。
表面上看,Register Tiling 的代码确实像是把 i 和 j 的循环展开了——2×3 的 tile “展开"成了 6 个独立的累加器。所以很容易产生一个直觉:Register Tiling 就是 Loop Unrolling。
但不完全对。
Loop Unrolling 有两个明确的收益:减少循环控制开销(比较、跳转指令),以及暴露更多指令级并行,让 CPU 的流水线和乱序执行更容易填满。之前那篇也讨论过这些。
但单纯展开 K 循环——把 4 次 K 迭代的代码平铺在一个循环体里——确实减少了循环开销,也暴露了更多 ILP。问题是,A 和 B 的复用次数并没有增加:每个 A 元素仍然只参与 1 次乘加,每个 B 元素也只参与 1 次。
Register Tiling 做的事情不同。它在 i 和 j 维度同时展开多个输出元素,让同一个 A 元素参与多次乘加(对应不同的 j),同一个 B 元素也参与多次乘加(对应不同的 i)。这才是 data reuse 的来源。
有一个需要小心的地方。朴素实现里 \(A_{ik}\) 在 j 循环的不同迭代中也会被使用——概念上确实是"复用”。但在朴素实现中,\(A_{ik}\) 在 \(j=0\) 时加载到寄存器,做完乘加后寄存器就被覆盖了;下一次 \(j=1\) 迭代需要同一个 \(A_{ik}\) 时,可能需要重新从 L1 甚至更慢的层级读取。现代编译器确实可能做 Load Hoisting,把这个加载提到 j 循环外面——但这取决于编译器能否识别出这个机会。Register Tiling 显式地把多个 j 的计算放在同一个循环体里,直接保证复用发生在寄存器级别,不依赖编译器优化。
Loop Unrolling 是一种实现手段,Data Reuse 才是 Register Tiling 的核心收益。 展开循环可以帮助编译器看到更多的复用机会,但光展开 K 循环不会自动产生跨 i/j 的数据复用。
当矩阵大到 Cache 装不下#
Register Tiling 解决的是"让数据在寄存器里多用几次"。但寄存器只是存储层级的最底层。
假设 A 和 B 都是 4096×4096 的 float 矩阵,每个大约 64 MB。典型的 CPU L1 Cache 只有 32-48 KB,L2 大约 256 KB-1 MB,L3 几 MB 到几十 MB。整个矩阵放不进任何一级 Cache。
即使 Register Tiling 已经做得很好,如果计算顺序不合理,仍然会反复从 DRAM 读取数据。比如按最朴素的循环顺序,计算 C 的第一行需要遍历整个 B 矩阵。等到计算第二行,又要重新遍历 B——而此时 B 的前面部分早就被驱逐出 Cache 了。
Cache Tiling(也叫 Loop Blocking)的思路很直接:把三层循环在 i、j、k 维度上都分块,每次只处理一个小块。
1for (int ii = 0; ii < M; ii += BM)
2 for (int jj = 0; jj < N; jj += BN)
3 for (int kk = 0; kk < K; kk += BK)
4 // BM×BN output block, BK-wide A and B tiles
5 for (int i = ii; i < ii+BM; i++)
6 for (int j = jj; j < jj+BN; j++)
7 for (int k = kk; k < kk+BK; k++)
8 C[i][j] += A[i][k] * B[k][j];选择合适的 BM、BN、BK,可以让每次内层小 GEMM 涉及的工作集(working set)控制在 Cache 容量以内。数据被读进 Cache 之后,在这个 block 的计算全部完成之前不会被驱逐。
换一种方式理解:Cache Tiling 缩短了数据的 reuse distance。从"上次使用这个元素"到"下次使用这个元素"之间经过的其他数据量变少了,少到不会把它从 Cache 里挤掉。
需要注意的是,Cache Tiling 并不直接控制硬件 Cache 具体缓存哪些行。Cache 的替换策略是硬件自动管理的。程序做的事情是调整访问顺序,使得在正常的 LRU 或类似策略下,需要复用的数据恰好还留在 Cache 里。
同一个思想,不同的层级#
把前面两节放在一起看。
从 DRAM 到 Register,每一级都更快、更小、离计算单元更近。Register Tiling 和 Cache Tiling 追求的是同一件事:让已经进入快速存储层级的数据,在被替换或覆盖之前参与尽可能多的计算。
但实现机制有区别。
Register Tiling 通过计算组织实现复用。程序员(或编译器)显式决定同时维护哪些累加器,决定哪些输入值在同一批计算中被共享。复用写在代码里。约束主要来自寄存器数量有限——tile 太大,寄存器不够用,编译器就会把一部分变量 spill 到栈上(也就是 L1 Cache 或更慢的地方),反而变慢。
Cache Tiling 通过访问顺序优化实现复用。程序调整循环的嵌套和步进方式,让工作集小到能被 Cache 容纳。复用是隐式的——代码不直接操作 Cache,只是通过控制访问模式来"引导"硬件的 Cache 管理行为。约束来自 Cache 容量和替换策略。
两者不互斥。实际的高性能 GEMM 实现(如 BLIS、OpenBLAS)通常同时使用多级 Tiling:外层 Loop Blocking 保证数据留在 L2 或 L3,中层针对 L1,内层 Register Tiling 保证寄存器级别的数据复用。每一级 Tiling 针对一个具体的存储层级做优化。
也不要把 Register Tiling 简化成"只优化 L1 到 Register 这一段"。Register Tiling 减少的 load 指令同时也减轻了 L1 Cache 的压力。类似地,Cache Tiling 减少 Cache Miss,但不能保证消除所有 DRAM 访问——如果矩阵大到连 L3 都放不下的部分,DRAM 访问终归不可避免。
大 Tile 的代价#
直觉上 tile 越大复用率越高——16×16 的 register tile 比 2×3 高得多。那为什么不直接用最大的 tile?
因为资源有限。
Register Tiling 方面:一个 \(m \times n\) 的 tile 需要 \(m \times n\) 个累加器寄存器,加上 \(m + n\) 个寄存器存放当前加载的输入。x86-64 有 16 个通用寄存器和 16-32 个 SIMD 寄存器。tile 太大,编译器被迫把一部分数据溢出到栈上——register spilling。溢出之后所谓的"寄存器级复用"实际变成了 L1 Cache 级别的访问,收益打折扣。
Cache Tiling 方面:如果 block 的工作集超过了目标 Cache 层级的容量,数据会在 block 还没算完的时候被驱逐,Miss 率反而增加。
在 GPU 上还有额外的约束:每个线程使用的寄存器数量直接影响 Occupancy——SM 上能同时调度多少个 warp。寄存器用多了,并行度下降,吞吐可能反而变差。
优化目标不是最大化 Tile Size,而是在硬件资源约束下平衡 data reuse、并行度和吞吐。
不只是 GEMM#
写到这里回头看,Register Tiling 和 Cache Tiling 的思想在 ML Systems 里到处都有。
GPU 上的高性能 GEMM kernel 在 Shared Memory 层面做 Tiling,把输入矩阵的 tile 搬到片上 SRAM,然后在片上完成计算——和 CPU 的 Cache Tiling 异曲同工,只是 GPU 的 Shared Memory 是程序显式管理的。Tensor Core 在更底层做 warp 级别的矩阵乘累加,可以看作一种硬件实现的 Register Tiling。
FlashAttention 把 Attention 的 \(S \times S\) Score Matrix 分块处理,每个 block 在片上计算、片上消费、片上丢弃——核心动机也是减少 HBM 搬运。Operator Fusion 把多个算子合并成一个 kernel,避免中间结果写回 HBM 再读回来——也是在缩短 reuse distance。
不过要注意区别。FlashAttention 不只是 Tiling——它还需要 Online Softmax 来打破跨 block 的全局归一化依赖,这是一次算法层面的等价重构。GPU Shared Memory Tiling 需要程序员显式管理数据搬运和同步,比 CPU 上改改循环顺序复杂得多。这些技术共享底层思想,但不是简单的套用。
最近几个月写了不少和 GPU 存储层级、Arithmetic Intensity、FlashAttention 相关的东西。这次从 CPU 的 GEMM 回过头来看 Tiling,反而觉得很多概念串起来了。
Tiling 的本质不是把计算切成小块,而是通过改变计算顺序,让每次昂贵的数据移动产生更多有用的计算。 这个思想从 CPU 的 Register 到 GPU 的 Shared Memory,从 BLAS 库到 FlashAttention,贯穿了整个高性能计算的优化逻辑。