为什么 cutlass_cute_tutorial_zh.md 中要用 retile_D
本文解释 cutlass_cute_tutorial_zh.md 中 gemm_smem() 的这两行:
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
tCrA/tCrB 按 MMA 的操作数规则组织,而 shared memory → 寄存器的加载按 Copy Atom 的规则组织。retile_D 给同一份寄存器创建适合加载的视图,让 copy() 写入后,gemm() 可以直接读取。
1. 先看这两个寄存器片段从哪里来
教程在调用 gemm_smem() 前,已经创建了 MMA 所需的寄存器片段:
Tensor tSrQ = thr_mma.partition_fragment_A(sQ); // (MMA, MMA_M, MMA_K)
Tensor tSrK = thr_mma.partition_fragment_B(sK); // (MMA, MMA_N, MMA_K)
在 score = Q @ K.T 中,Q 是 MMA 的 A 操作数,K 是 B 操作数。它们传入 gemm_smem() 后,分别对应形参 tCrA 和 tCrB。
这两个 tensor 已经属于当前线程,其布局围绕 MMA 的需求组织:
- 第一维表示一条 MMA 指令所需的操作数值。
- 第二维表示 M 或 N 方向上的重复。
- 第三维表示 K 方向上的重复。
它们准备好的是寄存器存储和索引方式。真正从 shared memory 加载数据,还需要后面的 copy()。
2. 为什么 MMA 的寄存器视图不能直接用于加载
教程对应的 kernel_traits.h 在 SM80、FP16 路径下使用:
// MMA 计算指令
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>
// Shared memory → 寄存器加载指令
Copy_Atom<SM75_U32x4_LDSM_N, elem_type>
SM75_U32x4_LDSM_N 对应 ldmatrix.x4。这里名称中的 SM75 表示该指令封装的架构来源,SM80 也使用它。
计算与加载的粒度不同:
| 操作 | 每线程涉及的 FP16 值 |
|---|---|
| 一条 MMA 的 A 操作数 | 8 个 half |
| 一条 MMA 的 B 操作数 | 4 个 half |
一条 ldmatrix.x4 的输出 | 8 个 half,即 4 个 32 位寄存器 |
B 操作数最能说明问题:一次加载写入 8 个 half,但一次 MMA 只使用其中 4 个 half。因此 Copy 视图需要把两个 MMA B 片段组合成一次加载的目的片段。
这里的“每线程输出 8 个 half”描述的是 warp 协作加载后分配给每个线程的寄存器结果;源端的 shared memory 地址分工由 ldmatrix 的源布局规定。
3. 用实际布局看 retile_D 改了什么
取对应源码中的 BlockM=64、BlockN=64、4 warps,并令 head_dim=64。MMA 的构造为:
auto tiled_mma = make_tiled_mma(
SM80_16x8x16_F32F16F16F32_TN{},
Layout<Shape<_4, _1, _1>>{},
Tile<_64, _16, _16>{}
);
按该配置实际打印 B 寄存器的 layout,可以得到:
MMA 视图:
shape = ((2, 2), 8, 4)
stride = ((1, 2), 4, 32)
Copy 视图:
shape = ((8, 1), 4, 4)
stride = ((1, 0), 8, 32)
为了理解分组,将第一维内部的嵌套展平:
| 视图 | Shape | 指令内的值 | N 方向重复 | K 方向重复 |
|---|---|---|---|---|
tCrB | (4, 8, 4) | 4 | 8 | 4 |
tCrB_copy_view | (8, 4, 4) | 8 | 4 | 4 |
两者都有 128 个 half,元素总数相同。改变的是如何把这些元素分组。
固定一个 K 分块,前两个 MMA B 片段对应:
MMA 视图:
第 0 个 N 分组 → r[0..3]
第 1 个 N 分组 → r[4..7]
Copy 视图:
第 0 个 N 分组 → r[0..7]
于是,一条 ldmatrix.x4 可以填好两个 MMA B 片段。随后 gemm() 仍通过原来的 MMA 视图,分别取出所需的 4 个 half。
A 操作数的一条 MMA 本来就需要 8 个 half,但第一维的嵌套结构仍需适配 Copy Atom。同一配置下,A 的第一维从 (2, 2, 2) 变成 (8, 1),外层 M、K 重复次数保持一致。代码统一为 A、B 创建对应的 Copy 视图。
4. copy() 和 gemm() 如何使用同一份寄存器
教程中的关键关系是:
// 为 MMA 寄存器创建 Copy 目的视图
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
// 加载时使用 Copy 视图
cute::copy(smem_tiled_copy_A, tCsA(_, _, i), tCrA_copy_view(_, _, i));
cute::copy(smem_tiled_copy_B, tCsB(_, _, i), tCrB_copy_view(_, _, i));
// 计算时使用 MMA 视图
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
tCrB_copy_view 与 tCrB 引用同一份寄存器存储。写入前者,就填好了后者。
因此,整个路径可以表示为:
shared memory
│
│ copy:通过 Copy 视图写入
▼
同一份寄存器存储
│
│ gemm:通过 MMA 视图读取
▼
accumulator
前面的 make_tiled_copy_A/B(atom, tiled_mma) 也很关键:它以 MMA 的线程和值分工为依据构造 TiledCopy。retile_D 再按照这个 TiledCopy 的规则重新组织已有寄存器片段。
5. 为什么用 retile_D,而不是 partition_D
两者处理的输入不同:
| API | 输入 | 作用 |
|---|---|---|
partition_D | 完整的目的 tensor | 按 Copy 的线程分工,选出当前线程负责的部分 |
retile_D | 已经属于当前线程的 fragment | 将已有片段重新组织成 Copy 所需的视图 |
tCrA/tCrB 已经由 partition_fragment_A/B() 按 MMA 的线程分工创建好了。这里需要重新组织线程内部的值,因此使用 retile_D。
D 表示这次 Copy 的 destination。在 shared memory → 寄存器的路径中,目的端就是寄存器,它与 GEMM 输出矩阵 D 无关。
6. retile_D 本身会搬运或重排数据吗
retile_D 的实现 是:
return make_tensor(
static_cast<DTensor&&>(dtensor).data(),
TiledCopy::retile(dtensor.layout())
);
它保留原来的 data(),只变换 layout,创建一个新视图。因此:
- 不创建第二份寄存器数据。
- 不执行 shared memory 加载。
- 不执行线程间的数据交换。
- 不负责传算交叠或流水线同步。
实际数据加载发生在 cute::copy() 中。
教程中“加载期间执行 retile”的注释,更准确的表述应当是:加载前,为已有的 MMA 寄存器片段创建 Copy 目的视图。