读懂 CuTe 的 retile_D：给 MMA 寄存器换一种分组方式

读 CuTe 的 GEMM 代码时，经常会遇到这两行：

auto tCrA = thr_mma.partition_fragment_A(sA);
auto tXrA = thr_copy.retile_D(tCrA);

第一行已经准备好了 MMA 要用的寄存器，第二行为什么还要做一次 retile？

先看结果。假设当前线程有 8 个 half，分别放在 r[0] 到 r[7] 中。MMA 需要这 8 个值，但我们选用的加载指令每次只给这个线程写入 2 个 half。加载时便需要这样分组：

寄存器里的元素： r0 r1 r2 r3 r4 r5 r6 r7

Copy 使用的分组：[r0 r1] [r2 r3] [r4 r5] [r6 r7]
                   0       1       2       3

retile_D 为后一种访问方式创建了一个新视图。tCrA 和 tXrA 指向同一份存储：加载时通过 tXrA 写入，计算时通过 tCrA 读取。

下面先用这 8 个 A 值读一遍源码，再沿着同样的过程看看 B。最后把 B 的 K 方向扩大，解释为什么不能总是把 retile 简化成“每两个值分一组”。

先固定我们要看的例子

这里使用一条 SM80 MMA 指令，以及 ldmatrix.x1 加载指令：

auto mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{});

auto tiled_copy = make_tiled_copy_A(
    Copy_Atom<SM75_U32x1_LDSM_N, half_t>{}, mma);

half 是 16 位浮点数。GPU 上一个 warp 有 32 个线程，warp 内的线程编号叫 lane，取值为 0～31。这个例子只有一个 warp。

这两个名字里，MMA 规定计算所需的寄存器分工，Copy 则规定一次加载怎样填充寄存器。先把 MMA 的要求弄清楚。

一条 MMA 指令，为什么需要每线程的 fragment

MMA 是 Matrix Multiply Accumulate，也就是矩阵乘加。本文这条指令由整个 warp 协作计算：

D = A × B + C

A：16×16     B：16×8     C、D：16×8

这里按通常的数学顺序写 B 的形状。输出有 16×8=128 个元素，每个输出元素都要沿 K 方向累加 16 个乘积：

D(m,n) = C(m,n) + sum(k=0..15, A(m,k) * B(k,n))

这 128 个输出由 32 个 lane 共同完成。执行时，每个 lane 只持有 A、B、C 的一部分，指令按固定规则组合整个 warp 的输入，并把该 lane 负责的结果写回寄存器。

具体哪个元素由哪个 lane 提供，是这条指令的接口约定。CuTe 的 fragment 就是用来表达当前线程持有的那一份值的。

从 CuTe 的名字找到实际指令

在 mma_sm80.hpp 中，SM80_16x8x16_F32F16F16F32_TN::fma 封装的是下面这条 PTX 指令。PTX 是 NVIDIA GPU 的虚拟指令集，这里先看操作数的形式：

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
    {d0,d1,d2,d3},
    {a0,a1,a2,a3},
    {b0,b1},
    {c0,c1,c2,c3};

m16n8k16 就是刚才的 M、N、K 尺寸。末尾四种类型按 D、A、B、C 的顺序排列：A、B 使用 FP16，C、D 使用 FP32。

.row.col 表示这条指令采用的 A 行主序、B 列主序操作数形式。CuTe 类名的 TN 对应这里的约定。后文按 CuTe 的接口将 A 坐标写成 (M,K)、B 写成 (N,K)；其中 B 的坐标顺序与上面的数学写法相反。

这条 MMA 接口接收的都是寄存器操作数。shared memory 中的矩阵数据，需要先通过加载指令放进各个 lane 的寄存器，才能参与计算。

.sync.aligned 要求 warp 中的线程协作执行同一条指令。上面列出的寄存器是每个 lane 的参数，32 个 lane 各自提供一份。

源码把参数类型写得很直接：

using DRegisters = float[4];
using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];
using CRegisters = float[4];

A、B 的类型是 uint32_t，因为一个 32 位操作数打包了两个 half 的原始位。它们在乘法中仍按 FP16 解释。于是可以把单个 lane 的负担列成这张表：

操作数	warp 共同处理的矩阵	每个 lane 的逻辑值	每个 lane 的 PTX 操作数
A	`16×16`	8 个 half	4 个 `uint32_t`
B	`16×8`	4 个 half	2 个 `uint32_t`
C	`16×8`	4 个 float	4 个 float
D	`16×8`	4 个 float	4 个 float

下文的 r[i] 表示一个 half 的逻辑存储槽位。对 A 来说，r[0]、r[1] 一起组成 a0，r[2]、r[3] 一起组成 a1，依此类推。

这也解释了文章开头的两个数量：MMA 一次需要当前 lane 的 8 个 A 值，ldmatrix.x1 一次只填入其中两个值，也就是一个 32 位操作数。需要四次加载，才能准备好这一条 MMA 的 A 输入。

CuTe 怎样记录这套硬件约定

mma_traits_sm80.hpp 将指令的计算尺寸、线程数，以及各操作数的线程分工写成类型：

指令要求	CuTe 中的表达
计算尺寸 `16×8×16`	`Shape_MNK = Shape<_16,_8,_16>`
32 个 lane 协作	`ThrID = Layout<_32>`
每个 lane 的 A 值对应哪些矩阵位置	`ALayout`，在 MMA Atom 中称为 `LayoutA_TV`
每个 lane 的 B 值对应哪些矩阵位置	`BLayout`，在 MMA Atom 中称为 `LayoutB_TV`
每个 lane 的累加值对应哪些输出位置	`CLayout`，在 MMA Atom 中称为 `LayoutC_TV`

这些 TV 布局将 (线程编号, 值编号) 映射到矩阵位置。它们记录的是指令已经规定好的分工。

MMA_Atom 把一条指令与这些 traits 放在一起，提供操作 fragment 的接口。TiledMMA 再将 MMA Atom 沿 M、N、K 方向组织起来，处理更大的矩阵块。本文开头没有额外指定重复配置，所以先从一条 MMA 的范围看起。

顺着这层关系，mma.get_slice(lane) 取得当前 lane 的分工，partition_fragment_A/B 再准备它要持有的值。计算时，mma_unpack 将 half fragment 通过 recast<uint32_t> 按位重解释为上面列出的打包操作数，再调用指令封装的 fma。

接下来看的 tCrA，就是这条调用链中的 A fragment。

tCrA 到底是什么

假设 sA 是 shared memory 中的一个 16×16 half tensor，按 (M,K) 索引，K 方向连续，行跨度为 16 个 half。实际加载时，共享内存基址按 16 字节对齐。当前线程先取得 MMA 的线程切片，再创建 A 的寄存器片段：

auto thr_mma = mma.get_slice(lane);
auto tCrA = thr_mma.partition_fragment_A(sA);

按这条 MMA 的分工，lane 0 的 A fragment 应当保存下面 8 个值：

r[0] 应当保存 A(0,0)    r[1] 应当保存 A(0,1)
r[2] 应当保存 A(8,0)    r[3] 应当保存 A(8,1)
r[4] 应当保存 A(0,8)    r[5] 应当保存 A(0,9)
r[6] 应当保存 A(8,8)    r[7] 应当保存 A(8,9)

这些矩阵位置由 MMA 的线程分工决定。创建 tCrA 时只是准备存储和索引方式，数据还要等后面的 copy 来加载。

partition_fragment_A 的源码也分成这两层：

return TiledMMA::make_fragment_A(partition_A(atensor));

partition_A 找出当前线程需要的元素；对于本例的 SM80 MMA，make_fragment_A 再通过 make_fragment_like 为它们创建寄存器片段。

实际打印 tCrA 的布局，会得到：

shape  = ((2,2,2),1,1)
stride = ((1,2,4),0,0)

读这对括号之前，先约定两个词。Shape 表示每个维度的大小，Stride 表示沿这个维度走一步时，存储位置增加多少。CuTe 把一个维度叫作 mode，mode 本身还可以由几个子维度组成。Layout 则把这些坐标换算成存储位置。

上面的三个顶层 mode 分别是：

(一条 MMA 内的 A 值, M 方向的 MMA 重复, K 方向的 MMA 重复)

第一个 mode 是 (2,2,2)，共有 2×2×2=8 个值。后两个 mode 都是 1，因为这个例子只计算一个 MMA tile，也就是一个矩阵块。最外层只有三个 mode，所以这个 tensor 的 rank 是 3；括号里的三个子维度不会让 rank 变成 5。

因此这两个操作的结果不同：

shape<0>(tCrA) = (2,2,2)
size<0>(tCrA)  = 8

源码里的 V 就是后者。矩阵块更大时，后两个 mode 也会大于 1，那时 V 只是一条 MMA 的每线程值数量，不是整个 tCrA 的元素总数。

Copy 与 MMA 怎样使用同一份寄存器

我们选的 ldmatrix.x1 每次给每个 lane 写入 2 个 half。于是 lane 0 的四次加载，应当依次写入四对位置：

Copy 分组	写入的槽位	lane 0 的 A 值	lane 1 的 A 值
0	`r[0],r[1]`	`A(0,0),A(0,1)`	`A(0,2),A(0,3)`
1	`r[2],r[3]`	`A(8,0),A(8,1)`	`A(8,2),A(8,3)`
2	`r[4],r[5]`	`A(0,8),A(0,9)`	`A(0,10),A(0,11)`
3	`r[6],r[7]`	`A(8,8),A(8,9)`	`A(8,10),A(8,11)`

注意，lane 0 始终写自己的寄存器，lane 1 也一样。表格右侧的差异是两个线程需要的矩阵元素不同，不是 retile 在两个线程之间交换了数据。

原来的 MMA 视图把这 8 个值组织为 (2,2,2)。Copy 视图要表达“每组 2 个值，共 4 组”，因此第一维改成 (2,4)：

tCrA 的 shape：((2,2,2),1,1)
tXrA 的 shape：((2,4),1,1)

将组内编号记为 a，分组编号记为 b，就有：

// a=0～1，b=0～3
tXrA(make_coord(a,b), 0, 0);
tCrA(a+2*b, 0, 0);

在本例中，这两种写法访问同一个元素。CuTe 允许用整数编号访问嵌套 mode，所以第二行可以用 0～7 访问第一维的 (2,2,2)。

外层 retile_D 的源码只有一个关键动作：

return make_tensor(
    static_cast<DTensor&&>(dtensor).data(),
    TiledCopy::retile(dtensor.layout()));

它用原来的 data()，配上算出来的新 layout，构造 tXrA。这就是两种视图共享存储的原因。

make_tiled_copy_A 已经用了 MMA 布局，为什么还要 retile

先看 make_tiled_copy_A 保存了什么：

return make_tiled_copy_impl(
    copy_atom,
    mma.get_layoutA_TV(),
    make_shape(tile_size<0>(mma), tile_size<2>(mma)));

TV 是 thread、value 的缩写。这里取得的正是前面介绍的 MMA A 分工：给定线程编号和该线程的值编号，找到它在整个 A tile 中的位置。

这套布局让 Copy 的参考分工与 MMA 保持一致。例如，参考布局中的 lane 0，第 0 个值就是 MMA lane 0 的第 0 个 A 值。

但是，copy_atom 自己还规定了“一次指令能处理多少值”。这两个数量同时保存在 TiledCopy 中：

TiledNumVal = 8   整个参考 tile 中，每线程的值数量
AtomNumVal  = 2   一个 Copy Atom 的参考布局中，每线程的值数量

AtomNumVal 取自 AtomLayoutRef：

using AtomNumVal = decltype(size<1>(AtomLayoutRef{}));

本例的 Copy traits 定义 RefLayout = DstLayout，参考布局就是目的端布局，因此这里的 2 也是一次加载后，每个线程收到的 half 数量。

make_tiled_copy_A 把这两层规则放在同一个对象里。它没有接收 tCrA，自然也没有把 tCrA 的 (2,2,2) 改成 (2,4)。把规则应用到已有寄存器上，是随后 retile_D(tCrA) 要做的事。

加载的完整过程便是：

auto thr_mma  = mma.get_slice(lane);
auto thr_copy = tiled_copy.get_slice(lane);

auto tCrA = thr_mma.partition_fragment_A(sA); // 准备 MMA 寄存器
auto tXsA = thr_copy.partition_S(sA);        // shared memory 的加载视图
auto tXrA = thr_copy.retile_D(tCrA);         // 寄存器的写入视图

copy(tiled_copy, tXsA, tXrA);
// 后续 MMA 通过 tCrA 读取已经加载的值。

其中，partition_S 处理 shared memory 的源端地址分工，retile_D 准备当前线程的目的寄存器视图。字母 D 在这里表示 Copy 的 destination。

这两个接口分开，是因为 ldmatrix 的地址提供者与数据接收者有不同的分工。下面把实际加载过程展开。

ldmatrix 怎样把行地址变成各线程的寄存器

SM75_U32x1_LDSM_N 在 copy_sm75.hpp 中封装的指令是：

ldmatrix.sync.aligned.x1.m8n8.shared.b16 {d}, [addr];

.shared 指定源数据位于 shared memory；.m8n8 表示一次处理一个 8×8 矩阵块；.b16 表示每个元素有 16 位，在本文中就是 half。.x1 表示加载一个这样的矩阵块。

整个块有 64 个 half，也就是 128 字节。指令执行后，每个 lane 得到一个 32 位的 d，里面装着两个 half。32 个 lane 合起来，刚好保存整个 8×8 块。

和上面的 MMA 一样，.sync.aligned 要求 warp 协作执行同一条指令。行地址还需要满足 16 字节对齐；本例在共享内存基址对齐后，32 字节的行跨度与 16 字节的子块偏移会保持这一对齐。

谁提供地址，谁接收值

先把当前 8×8 子块记为 X。对于 .x1，lane 0～7 各提供一行的起始地址，每行包含 8 个 half。

提供地址的 lane	地址指向的数据	接收这行数据的 lane
0	`X(0,0)` 开始的 8 个 half	0～3
1	`X(1,0)` 开始的 8 个 half	4～7
2	`X(2,0)` 开始的 8 个 half	8～11
…	…	…
7	`X(7,0)` 开始的 8 个 half	28～31

一行的 8 个 half 被分成四对，交给四个目的 lane。以第 0 行为例：

X(0,0) X(0,1) | X(0,2) X(0,3) | X(0,4) X(0,5) | X(0,6) X(0,7)
     lane 0   |      lane 1   |      lane 2   |      lane 3

一般地，令目的 lane 为 4*i+j，其中 i=0～7、j=0～3，它收到：

d 的低 16 位：X(i,2*j)
d 的高 16 位：X(i,2*j+1)

所以，lane 1 提供第 1 行的地址，但它自己的寄存器接收的是第 0 行中的第 2、3 个元素。这里的分发由 ldmatrix 完成。

copy_traits_sm75.hpp 用源、目的两个布局记录这套行为。对 .x1，源线程 mode 是 (8,4)，其中第二个子 mode 的 stride 为 0：CuTe 让 lane 0、8、16、24 使用同一行地址，其余行也各重复四次。有效的八行地址由 lane 0～7 提供。

因此，打印 partition_S 时可能看到不同 lane 有相同的地址。这与指令在整个 warp 上只加载一个 8×8 块相符。

为什么 Copy 的源端有 8 个值，目的端却只有 2 个

源码中的寄存器接口是：

using SRegisters = uint128_t[1];
using DRegisters = uint32_t[1];

copy 接收 uint128_t const& smem_src，随后把 &smem_src 转为 shared memory 地址，传给 PTX 的 [addr]。这个 128 位引用描述一行 8 个 half 的地址范围；真正返回给当前 lane 的结果只有一个 32 位 d。

因此，Copy_Atom<SM75_U32x1_LDSM_N, half_t> 的源端用 8 个 half 表示一行，目的端则用 2 个 half 表示当前 lane 收到的值。该 Atom 的 reference 采用目的布局，所以 AtomNumVal=2。

这两个数量描述加载过程的不同两端。它们通过 warp 内的分发规则对应起来。

四次 x1 怎样填满 A fragment

现在回到 16×16 的 A tile。它可以分成四个 8×8 子块，本例的 Copy 分组顺序为：

                  K=0～7     K=8～15
M=0～7             第 0 组     第 2 组
M=8～15            第 1 组     第 3 组

第 0 组加载左上角，lane 0 收到 A(0,0),A(0,1)；第 1 组加载左下角，它收到 A(8,0),A(8,1)。后两组同理，正好填入前面表格中的四对寄存器槽位。

partition_S 安排每次加载的行地址，tXrA 安排这次加载写入哪一对槽位。四条 ldmatrix.x1 完成后，每个 lane 的四个 32 位 A 操作数便准备好了，后续 MMA 按 tCrA 的方式读取它们。

x2 和 x4 改变了多少工作量

x1、x2、x4 统计的是一次指令加载的 8×8 块数：

形式	warp 加载的块数	提供各行地址的 lane	每个 lane 得到的 32 位结果	每个 lane 的 half 数
`x1`	1	0～7	1 个	2
`x2`	2	0～15，每 8 个 lane 对应一块	2 个	4
`x4`	4	0～31，每 8 个 lane 对应一块	4 个	8

SM75_U32x1_LDSM_N、SM75_U32x2_LDSM_N、SM75_U32x4_LDSM_N 分别封装这三种非转置加载。带 .trans 的形式在 CuTe 中使用另一组 Copy 类型，本文先沿用非转置版本。

这里的 8×8 是 Copy 指令的块形状，而 MMA 指令的计算尺寸是 16×8×16。两条指令各有自己的粒度，所以 CuTe 分别提供 Copy_Atom 和 MMA_Atom，再用 TiledCopy 的分工与寄存器视图将它们接起来。

带着 8 和 2，读一遍 retile 的源码

内部 retile 在这条调用路径上接收的是 tCrA.layout()。名字虽然叫 tensor，这里实际处理的是 Layout：

template <class Tensor>
CUTE_HOST_DEVICE constexpr static
auto retile(Tensor&& tensor)
{
  constexpr int R = remove_cvref_t<Tensor>::rank;
  auto V = size<0>(tensor);

  auto frg_layout_mn = upcast<TiledNumThr{} * V>(
      right_inverse(TiledLayout_TV{})
          .with_shape(shape(Tiler_MN{})));

  auto frg_layout_v = zipped_divide(
      logical_product(make_layout(V), right_inverse(frg_layout_mn)),
      make_layout(AtomNumVal{}));

  auto t_tensor = zipped_divide(
      tensor,
      prepend(product_each(shape(frg_layout_mn)), V));

  auto v_tensor = t_tensor.compose(frg_layout_v, _);

  return v_tensor(_, append<R>(Int<0>{}, _));
}

这一段的难点在于，它同时算两件事：Copy 应当访问哪些逻辑坐标，以及这些坐标怎样落到原寄存器上。下面分开看。

先确定要从哪些 fragment 位置取值

本例中，几个参数的值为：

R           = 3
V           = 8
TiledNumThr = 32
AtomNumVal  = 2
Tiler_MN    = (16,16)

Tiler_MN 是 Copy 实现里的通用名字；对 A 操作数，这两个方向实际是 M、K。

TiledLayout_TV 从线程和值编号找到矩阵位置。right_inverse 沿相反方向，找到矩阵位置对应的 TV 编号。线程维度在前，TV 编号按 thread+32*value 展平。例如：

A(0,0) -> 0  = lane 0 的值 0
A(0,1) -> 32 = lane 0 的值 1
A(8,0) -> 64 = lane 0 的值 2
A(0,2) -> 1  = lane 1 的值 0

with_shape 把这个逆映射的输入组织成 A tile 的 (M,K) 两维。接着，upcast<32*8> 按 256 个元素为一个单位压缩布局。

对于当前例子，所有线程的前 8 个值已经覆盖完整的 16×16 A tile，因此压缩后没有更多需要收集的位置：

product_each(shape(frg_layout_mn)) = (1,1)
right_inverse(frg_layout_mn)       = 1:0

本文用 shape:stride 表示一个布局。1:0 只有一个合法坐标 0，它也映射到位置 0；(1,1) 则说明，在 M、K 两个方向各取一个 fragment 位置就够了。

建立 Copy 的分组编号

现在看这行最容易让人停住的代码：

auto frg_layout_v = zipped_divide(
    logical_product(make_layout(V), right_inverse(frg_layout_mn)),
    make_layout(AtomNumVal{}));

先把内层展开：

make_layout(V)              = 8:1
right_inverse(frg_layout_mn) = 1:0

logical_product(8:1,1:0) = (8,1):(1,0)

8:1 为当前 fragment 的值建立编号 0～7。另一个布局指定还要从哪些 fragment 位置取值。本例只有一个位置，product 的第二个 mode 大小便是 1。

这里创建的是一张编号表，没有分配新的寄存器。接着按每组 2 个值划分：

zipped_divide((8,1):(1,0),2:1) = (2,4):(1,2)

结果读作：第一维大小是 2，第二维大小是 4，坐标 (a,b) 映射到 a+2*b。这正好是前面四对值的分组方式。

此时仍要区分两个布局：

布局	回答的问题
`tCrA.layout()`	给定 MMA fragment 坐标，访问哪个寄存器槽位？
`frg_layout_v`	给定 Copy 分组坐标，访问输入 tile 的哪个逻辑编号？

frg_layout_v 的输出暂时不是寄存器地址。它还要经过输入布局，才能找到实际存储位置。

把编号接到原来的存储上

下一行先把输入划分为“本次处理的 tile”和“外层剩余部分”：

auto t_tensor = zipped_divide(
    tensor,
    prepend(product_each(shape(frg_layout_mn)), V));

代入 (1,1) 和 V=8，tiler 为 (8,1,1)。于是本次取出全部 8 个值，外层各维度都只剩一个位置：

t_tensor 的 shape = ((8,1,1),(1,1,1))
                       本次 tile  剩余部分

再将刚才的 Copy 分组布局接到 tile 上：

auto v_tensor = t_tensor.compose(frg_layout_v, _);

compose 可以按函数嵌套来读：先用 frg_layout_v 算编号，再交给 t_tensor。_ 表示保留第二个 mode，也就是剩余部分。

Copy 坐标 (a,b)
      │ frg_layout_v
      ▼
输入 tile 的逻辑编号 a+2*b
      │ t_tensor
      ▼
原寄存器槽位 r[a+2*b]

组合后的 shape 是 ((2,4),(1,1,1))。

去掉一个多余的大小为 1 的维度

最后一行：

return v_tensor(_, append<R>(Int<0>{}, _));

本例 R=3，append 生成 (0,_,_)。这行固定第二个 mode 中的第一个坐标为 0，保留后面的 M、K 剩余维度：

v_tensor(_, make_coord(Int<0>{}, _, _));

最终得到：

shape  = ((2,4),1,1)
stride = ((1,2),0,0)

外层 retile_D 再把这个布局和原来的 data() 放在一起。到这里，tXrA 的每个坐标就都能找到 tCrA 中对应的槽位了。

这个 A 例子能不能直接 divide

可以。这里没有额外的 fragment 收集，8 个值的顺序也正好符合 Copy 的分组要求。

直接写：

auto direct_zip = zipped_divide(tCrA, make_layout(_2{}));

会得到 (2,4):(1,2)，访问的槽位相同，但返回的是两维 tensor。若只划分第一个 mode，同时保留另外两个 mode：

auto direct_mode = logical_divide(
    tCrA, make_tile(make_layout(_2{})));

结果就是 ((2,4),1,1):((1,2),0,0)，与本例 retile_D 的结果完全相同。

这也是阅读这段源码时需要留意的地方：简单例子里的 logical_product 看起来只是多加了一个大小为 1 的 mode。它的作用，要到 Copy 需要跨 fragment 取值时才明显。

把 B 操作数也走一遍

先沿用同一条 MMA 和同一个 ldmatrix.x1，只把操作数从 A 换成 B：

auto mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{});
auto tiled_copy_b = make_tiled_copy_B(
    Copy_Atom<SM75_U32x1_LDSM_N, half_t>{}, mma);

这条 MMA 的 A tile 是 16×16，B tile 则是 8×16。32 个线程分担 128 个 B 元素，每线程只需要 4 个 half。因此 B 的加载是“两组，每组 2 个值”：

当前线程的 B 值： r0 r1 r2 r3
Copy 使用的分组：[r0 r1] [r2 r3]
                   0       1

先把 B 的坐标顺序说清楚

通常写矩阵乘法时，B 是一个 K×N 矩阵，元素写作 B(k,n)。CuTe 的 MMA 接口把 B tensor 按 (N,K) 顺序组织。

下文的 B(n,k) 都采用 CuTe 的顺序。例如，文中的 B(0,8) 是 n=0、k=8，也就是普通矩阵写法中的 B(8,0)。

本例的 shared memory 视图如下：

// sB 的坐标顺序为 (N,K)，K 方向连续。
auto sB = make_tensor(
    make_smem_ptr(storage),
    make_layout(make_shape(_8{}, _16{}),
                make_stride(_16{}, _1{})));

这是为布局推导选定的存储方式。后面的 retile_D 处理寄存器视图；shared memory 中的坐标怎样对应地址，由 sB 的布局决定。

tCrB 保存哪四个值

创建 B fragment 的方法与 A 对称：

auto thr_mma = mma.get_slice(lane);
auto tCrB = thr_mma.partition_fragment_B(sB);

源码里对应的是：

return TiledMMA::make_fragment_B(partition_B(btensor));

partition_B 确定当前线程的 B 元素，make_fragment_B 在这个 SM80 配置中通过 make_fragment_like 创建寄存器片段。创建时尚未加载数据。

把前两个 lane 的元素列出来：

槽位	lane 0 应当保存的值	lane 1 应当保存的值	Copy 分组
`r[0]`	`B(0,0)`	`B(0,2)`	0
`r[1]`	`B(0,1)`	`B(0,3)`	0
`r[2]`	`B(0,8)`	`B(0,10)`	1
`r[3]`	`B(0,9)`	`B(0,11)`	1

注意，lane 1 的 N 坐标仍是 0，只是 K 坐标向后移动了两个位置。到 lane 4 时，才轮到 N 坐标为 1 的那一行。

这个分工可以从 mma_traits_sm80.hpp 的 BLayout 中读出来。本文使用的 FP32 累加指令继承了下面这套 B 布局：

using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_2, _2>>,
                       Stride<Stride<_16,_1>,Stride<_8,_64>>>;

线程 mode 的 (4,8) 把 lane 拆成 (lane%4,lane/4)；值 mode 的 (2,2) 则拆成组内编号 a 和分组编号 b。按照 stride，输出为：

offset = 16*(lane%4) + lane/4 + 8*a + 64*b

这个 offset 按 n+8*k 编码 B tile 内的逻辑位置。令 a=0～1、b=0～1，就能读出任意 lane 的对应关系：

r[a+2*b] 应当保存 B(lane/4, 2*(lane%4)+a+8*b)

这里的除法是整数除法。这个关系来自 MMA 的 B 线程分工，Copy 需要把值放到这些指定位置。

实际的 tCrB 布局是：

shape  = ((2,2),1,1)
stride = ((1,2),0,0)

三个顶层 mode 依次表示：

(一条 MMA 内的 B 值, N 方向的 MMA 重复, K 方向的 MMA 重复)

因此，size<0>(tCrB)=4，R=3。第一维里的两个 2 分别区分一对 half 中的两个值，以及 K 方向相隔 8 的两对值。

make_tiled_copy_B 怎样采用 MMA 的分工

它的实现只比 A 版本换了两处：

return make_tiled_copy_impl(
    copy_atom,
    mma.get_layoutB_TV(),
    make_shape(tile_size<1>(mma), tile_size<2>(mma)));

get_layoutB_TV() 给出 B 的线程和值布局；tile_size<1> 和 tile_size<2> 取 N、K 尺寸，所以参考 tile 是 (8,16)。

和 A 一样，这一步先保存规则，之后 retile_D(tCrB) 才把规则应用到寄存器视图上。本例的几个数量是：

V           = 4   tCrB 第 0 个 mode 的大小
TiledNumVal = 4   参考 B tile 中每线程的值数量
AtomNumVal  = 2   一次 ldmatrix.x1 的参考值数量
TiledNumThr = 32
Tiler_MN    = (8,16)

这里的 Tiler_MN 仍是 Copy 源码里的变量名，对 B 应当按 (N,K) 理解。下面也把局部变量 frg_layout_mn 称为 frg_layout_nk，方便对照两个方向；运算没有改变。

将 B 的数字代入 retile

先做 TV 布局的逆映射。TV 编号仍按 lane+32*value 展平，这次对应的是 B 的矩阵位置：

B(0,0) -> 0  = lane 0 的值 0
B(0,1) -> 32 = lane 0 的值 1
B(0,8) -> 64 = lane 0 的值 2
B(0,2) -> 1  = lane 1 的值 0

with_shape 把逆映射的输入恢复为 (N,K) 两维：

right_inverse(TiledLayout_TV).with_shape((8,16))
    = (8,(2,4,2)):(4,(32,1,64))

接着的 upcast 参数是 32*4=128。所有线程的前 4 个值，已经覆盖整个 8×16 B tile，因此压缩后，在 N、K 方向各取一个 fragment 位置就够了：

product_each(shape(frg_layout_nk)) = (1,1)
right_inverse(frg_layout_nk)       = 1:0

再建立值编号，按每组 2 个值划分：

logical_product(make_layout(4),1:0) = (4,1):(1,0)
zipped_divide((4,1):(1,0),2:1)      = (2,2):(1,2)

所以 frg_layout_v(a,b)=a+2*b。它表示两组 Copy，每组访问两个逻辑编号。

随后，从原布局中取出本次处理的 tile。把 V=4 放在 (1,1) 前面，得到 tiler (4,1,1)：

t_tensor = zipped_divide(tCrB.layout(),(4,1,1))
         = ((4,1,1),(1,1,1)):((1,0,0),(0,0,0))

再用 compose 把 Copy 编号接到这个 tile 上：

v_tensor = t_tensor.compose(frg_layout_v,_)
         = ((2,2),(1,1,1)):((1,2),(0,0,0))

最后，append<3>(0,_) 生成 (0,_,_)，固定剩余 V mode 的坐标，保留 N、K 的剩余 mode：

tXrB 的 layout = ((2,2),1,1):((1,2),0,0)

此时，输入与输出的布局相同。这个 B fragment 原来就把 4 个值组织成两对，恰好符合 ldmatrix.x1 的要求。retile_D 算出的 Copy 分组便与原布局重合。

从 shared memory 加载，再交给 MMA

完整调用过程是：

auto thr_mma    = mma.get_slice(lane);
auto thr_copy_b = tiled_copy_b.get_slice(lane);

auto tCrB = thr_mma.partition_fragment_B(sB);
auto tXsB = thr_copy_b.partition_S(sB);
auto tXrB = thr_copy_b.retile_D(tCrB);

copy(tiled_copy_b, tXsB, tXrB);
// 后续 MMA 通过 tCrB 读取 B 值。

tXrB 与 tCrB 共享存储。在当前配置中：

// a=0～1，b=0～1
tXrB(make_coord(a,b),0,0);
tCrB(a+2*b,0,0);

两行访问同一槽位。第一组 Copy 写入 r[0],r[1]，第二组写入 r[2],r[3]。

源端的线程分工也与 A 的例子类似。第一组中，源端 lane 0 提供 B(0,0～7) 的地址，这 8 个 half 分发给目的端 lane 0～3。第二组则提供 B(0,8～15)，填入这几个 lane 的另一对寄存器。

而源端 lane 1 提供的是 B(1,0～7) 的地址，其数据发给目的端 lane 4～7。它自己的目的寄存器收到的却是表格中的 B(0,2),B(0,3)。源端提供地址的线程，与目的端接收值的线程，依然不能混为一谈。

在这个例子里，直接划分也成立：

zipped_divide(tCrB,2:1)
    -> (2,2):(1,2)

logical_divide(tCrB,make_tile(2:1))
    -> ((2,2),1,1):((1,2),0,0)

前者返回两维视图，后者保留原来的三个顶层 mode。

把 Copy 换成 ldmatrix.x2 后，一次加载会给每线程写入 4 个 half，刚好装下这一条 MMA 的全部 B 值：

auto copy_b_x2 = make_tiled_copy_B(
    Copy_Atom<SM75_U32x2_LDSM_N, half_t>{}, mma);
auto tXrB_x2 = copy_b_x2.get_slice(lane).retile_D(tCrB);

此时 AtomNumVal=4，同一份 tCrB 会得到另一种分组：

ldmatrix.x1：((2,2),1,1)  两组，每组 2 个值
ldmatrix.x2：((4,1),1,1)  一组，每组 4 个值

寄存器还是那四个 half 槽位，变化的是一次 Copy 使用多少个槽位。

两份 fragment 最后怎样交给 MMA

把前面 A、B 的加载放在同一个 kernel 中，copy 完成后，tCrA 和 tCrB 中就有了 MMA 需要的输入。

假设 gC 是全局内存中的 16×8 输出 tensor，这个小例子从 C=0 开始计算：

auto tCrC = thr_mma.partition_fragment_C(gC);
clear(tCrC);

gemm(mma, tCrA, tCrB, tCrC); // 在原来的累加器中计算 C = A*B+C
copy(tCrC, thr_mma.partition_C(gC));

在当前单个 tile 的配置中，gemm 最终调用前面那条 mma.sync.aligned.m16n8k16。每个 lane 提供 4 个打包的 A 操作数、2 个打包的 B 操作数，以及 4 个 float 累加值，结果仍放进这 4 个 float 中。

整个过程可以连起来看：

shared memory 中的 A、B
         │ ldmatrix：按 warp 分发数据
         ▼
每个 lane 的输入寄存器
  Copy 通过 tXrA、tXrB 写入
  MMA  通过 tCrA、tCrB 读取
         │ mma.sync：矩阵乘加
         ▼
每个 lane 的 4 个 float 累加值 tCrC
         │ copy：写回当前 lane 负责的输出位置
         ▼
全局内存中的输出矩阵

partition_fragment 准备存储，retile_D 建立加载视图，copy 和 gemm 执行加载与计算。把它们对应到实际指令后，每一层的用途就比较容易分清了。

当一次 Copy 要合并两个 MMA fragment

上面的 B 例子中，x1 把 4 个值拆成两组，x2 把 4 个值放进一组。如果继续换成 ldmatrix.x4，Copy 目的端每线程就需要 8 个 half，而一条 MMA 只能提供 4 个 B 值的寄存器位置，需要从两个 MMA fragment 中收集。

选择下面的配置，让 Copy 在 K 方向收集这两个 fragment：

auto mma_k = make_tiled_mma(
    SM80_16x8x16_F32F16F16F32_TN{},
    Layout<Shape<_1,_1,_1>>{},
    Tile<_16,_8,_32>{});

auto copy_k = make_tiled_copy_B(
    Copy_Atom<SM75_U32x4_LDSM_N,half_t>{}, mma_k);

这里的 tile 让参考布局覆盖 K 方向的 32 个元素，而单条 MMA 的 K 尺寸仍是 16。

对每个 N 位置，这段 K=32 的计算要用两条 K=16 的 MMA，依次累加到同一份 C。一次 ldmatrix.x4 可以提前装好这两条指令要用的 B 寄存器，之后两次 MMA 分别读取各自的 4 个 half。

于是几个数量变成了：

V           = 4   一个 MMA B fragment 的每线程值数量
TiledNumVal = 8   参考 tile 跨两个 K fragment，每线程有 8 个值
AtomNumVal  = 8   一次 ldmatrix.x4 的每线程目的值数量
Tiler_MN    = (8,32)

这里 TiledNumVal 和 V 已经不同。前者描述整个参考 tile，后者只取 tCrB 第 0 个 mode 的大小。

令 shared memory 中的 B tensor 按 CuTe 的 (N,K) 顺序组织，大小为 32×32，布局为 (32,32):(32,1)。创建 tCrB 后，将每个顶层 mode 展平，寄存器布局是：

shape  = (4,4,2)   // (MMA 内部值, N 重复, K 重复)
stride = (1,8,4)

tCrB(v,n,k) -> r[v+8*n+4*k]

两个 K fragment 怎样连成一组

这次 upcast 仍按 32*4=128 个值压缩。但参考 tile 有 8×32=256 个元素，压缩后会留下 K 方向的两个 fragment 位置：

inverse with NK shape = (8,(2,4,4)):(4,(32,1,64))
frg_layout_nk         = (1,(1,1,2)):(1,(1,1,1))

product_each(shape(frg_layout_nk)) = (1,2)
G = right_inverse(frg_layout_nk)   = 2:1

先把每个 fragment 的 4 个值与这两个位置组合，再按每组 8 个值划分：

logical_product(4:1,2:1) = (4,2):(1,4)
frg_layout_v            = (8,1):(1,0)

也就是一个 Copy 分组，组内有 8 个值。输入的收集 tile 则为 (4,1,2)：值维度取 4 个，N 方向取一个 fragment，K 方向取两个。

代入源码的最后几行：

t_tensor = zipped_divide(tCrB.layout(),(4,1,2))
         = ((4,1,2),(1,4,1)):((1,0,4),(0,8,0))

v_tensor = t_tensor.compose(frg_layout_v,_)
         = ((8,1),(1,4,1)):((1,0),(0,8,0))

result   = v_tensor(_,(0,_,_))
         = ((8,1),4,1):((1,0),8,0)

结果中的 4 保留了 N 方向的四个剩余位置；K 方向的两个 fragment 已经收进第一个 mode，外层 K mode 便只剩 1。

对于 a=0～7、n=0～3，两个视图的对应关系是：

tXrB(make_coord(a,0),n,0);
tCrB(a%4,n,a/4);

两者访问的槽位都是 r[a+8*n]。以 lane 0 的 n=0 分组为例，它写入：

第一个 MMA fragment：r[0] B(0,0)    r[1] B(0,1)
                    r[2] B(0,8)    r[3] B(0,9)

第二个 MMA fragment：r[4] B(0,16)   r[5] B(0,17)
                    r[6] B(0,24)   r[7] B(0,25)

现在，MMA 仍然按每个 fragment 的 4 个值来读，Copy 则一次写入跨两个 fragment 的 8 个值。

具体到上面这 8 个槽位，r[0..3] 组成第一条 MMA 的两个 32 位 B 操作数，r[4..7] 则组成第二条 MMA 的两个 B 操作数。Copy 的一次分组可以服务多条 MMA，这正是此处需要跨 fragment 收集的原因。

为什么直接 divide 会取错方向

第一个 Copy 分组需要从 K 方向取两个 fragment。对比直接划分原 tensor 的结果：

做法	第一个 fragment	第二个 fragment	最终槽位
`retile_D`	`(v=0～3,n=0,k=0)`	`(v=0～3,n=0,k=1)`	`r[0..7]`
`zipped_divide(tCrB,8:1)`	`(v=0～3,n=0,k=0)`	`(v=0～3,n=1,k=0)`	`r[0..3]` 和 `r[8..11]`

直接划分按原 tensor 的逻辑坐标顺序走，值维度之后是 N 维度，所以跨到了 N 方向。它取到了 8 个值，但这 8 个值不是当前 Copy 分组需要的那一组。

前面的 retile 推导则指定了 (4,1,2) 的收集 tile。正是这个 tile，让逻辑编号按需要的方向解码。同样是编号 7，在不同 shape 中有不同含义：

编号在哪个 shape 中解码	对应的 `(v,n,k)`	寄存器槽位
本次收集 tile `(4,1,2)`	`(3,0,1)`	7
原 tensor `(4,4,2)`	`(3,1,0)`	11

现在就能看出 logical_product 和 compose 为什么分开了：前者按参考布局建立要访问的逻辑编号，后者让这些编号经过选定的输入 tile，落到原来的存储上。

如果自己构造一个包含相同收集规则的 tiler，也能改写这段运算。只给普通 zipped_divide 一个 8，它没有足够的信息判断应当沿 N 还是 K 方向收集。

换配置时，要检查什么

本文的 A 例子使用 make_tiled_copy_A，而 ldmatrix 的 reference 等于 destination。它们保证了这里的输入 MMA 值顺序与 Copy 参考布局兼容。

源码里还明确假定，输入第 0 个 mode 的前 V 个值，对应 TiledLayout_TV 中前 V 个 value ID。retile_D 不像 partition_D 那样额外做 Dst 到 Ref 的转换。因此，换一套 Copy Atom 或参考线程布局后，需要重新检查这个对应关系。

数量上，TiledCopy 检查整个 tile 的线程数和值数量能被 Atom 整除：

TiledNumThr % AtomNumThr == 0;
TiledNumVal % AtomNumVal == 0;

没有要求 V == AtomNumVal。A 例子是 8 拆成四组 2；B 配 x1 是 4 拆成两组 2，配 x2 是 4 个值组成一组，配 x4 则从两个 fragment 中各取 4 个值。真正要对齐的是参考值顺序和指令需要的分组。

打印布局，核对一次

文中的布局都用当前 CuTe 源码构造并检查过。完整程序单独放在三个文件中，方便复制运行：

cute_retile_trace.cu：打印 A 例子的中间布局，检查全部 32 个 lane 的寄存器对应关系，以及两种直接划分的等价性。
cute_retile_b_trace.cu：打印单个 MMA 的 B 布局和逐步 retile 结果，检查全部 32 个 lane 的源端、目的端对应关系，以及 x1、x2 的分组。
cute_retile_gather_trace.cu：打印扩展 B 例子的中间布局，检查沿 K 收集的结果，并对比直接划分得到的槽位。

在 CUTLASS 根目录编译：

nvcc -std=c++17 -I include cute_retile_trace.cu -o /tmp/retile_trace
/tmp/retile_trace

nvcc -std=c++17 -I include cute_retile_b_trace.cu -o /tmp/retile_b_trace
/tmp/retile_b_trace

nvcc -std=c++17 -I include cute_retile_gather_trace.cu -o /tmp/retile_gather_trace
/tmp/retile_gather_trace

三个程序在主机上检查坐标和布局，不执行 GPU 上的 MMA 或 ldmatrix 指令。指令形式可对照下方链接中的封装源码。源端检查的计数包含不同 lane 重复提供的地址，统计的是逻辑对应项。

A 例子会打印：

tCrA: ((_2,_2,_2),_1,_1):((_1,_2,_4),_0,_0)
tXrA: ((_2,_4),_1,_1):((_1,_2),_0,_0)
direct_zip: (_2,_4):(_1,_2)
direct_mode: ((_2,_4),_1,_1):((_1,_2),_0,_0)

CuTe 打印静态整数时带有 _ 前缀，如 _2；正文中的布局省略了这个前缀。

单个 MMA 的 B 例子会打印：

tCrB: ((_2,_2),_1,_1):((_1,_2),_0,_0)
tXrB: ((_2,_2),_1,_1):((_1,_2),_0,_0)
tXrB_x2: ((_4,_1),_1,_1):((_1,_0),_0,_0)
Verified 128 B destination matches and 512 source references; 128 x2 destination matches.

扩展 B 例子的关键输出是：

K group slots: 0 1 2 3 4 5 6 7
Direct group slots: 0 1 2 3 8 9 10 11

看到这两行，再回看 frg_layout_v 和 t_tensor.compose(...)，就能把“该取哪些值”和“这些值存在哪里”分别对上。

对照源码

mma_sm80.hpp：MMA 的 PTX 指令、每个 lane 的输入与输出寄存器接口。
copy_sm75.hpp：ldmatrix.x1/x2/x4 及转置版本的指令封装。
copy_atom.hpp：make_tiled_copy_A/B、TiledCopy::retile、ThrCopy::retile_D。
mma_atom.hpp：partition_fragment_A/B 和 make_fragment_A/B。
mma_traits_sm80.hpp：本文使用的 SM80 MMA 的 A、B 线程和值布局。
mma_traits.hpp：mma_unpack 如何把 fragment 按位重解释为指令操作数，并调用 fma。
copy_traits_sm75.hpp：ldmatrix 的 Src、Dst、Ref 布局。
gemm.hpp：如何沿 fragment 的 M、N、K mode 调用 MMA Atom。
layout.hpp：right_inverse、logical_product、zipped_divide、upcast。
tensor_impl.hpp：Tensor 的 logical_divide 与 zipped_divide。

先固定我们要看的例子​

一条 MMA 指令，为什么需要每线程的 fragment​

从 CuTe 的名字找到实际指令​

CuTe 怎样记录这套硬件约定​

tCrA 到底是什么​

Copy 与 MMA 怎样使用同一份寄存器​

make_tiled_copy_A 已经用了 MMA 布局，为什么还要 retile​

ldmatrix 怎样把行地址变成各线程的寄存器​

谁提供地址，谁接收值​

为什么 Copy 的源端有 8 个值，目的端却只有 2 个​

四次 x1 怎样填满 A fragment​

x2 和 x4 改变了多少工作量​

带着 8 和 2，读一遍 retile 的源码​

先确定要从哪些 fragment 位置取值​

建立 Copy 的分组编号​

把编号接到原来的存储上​

去掉一个多余的大小为 1 的维度​

这个 A 例子能不能直接 divide​

把 B 操作数也走一遍​

先把 B 的坐标顺序说清楚​

tCrB 保存哪四个值​

make_tiled_copy_B 怎样采用 MMA 的分工​

将 B 的数字代入 retile​

从 shared memory 加载，再交给 MMA​

两份 fragment 最后怎样交给 MMA​

当一次 Copy 要合并两个 MMA fragment​

两个 K fragment 怎样连成一组​

为什么直接 divide 会取错方向​

换配置时，要检查什么​

打印布局，核对一次​

对照源码​