跳到主要内容

为什么 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)484
tCrB_copy_view(8, 4, 4)844

两者都有 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 目的视图。

源码位置​