CuTe 中 A、B 如何加载到 MMA 寄存器：从 fragment 到 ldmatrix

读 CuTe 的 GEMM 内核时，常会遇到这样的调用：

auto tCrA = thr_mma.partition_fragment_A(sA);
auto tXrA = thr_copy_a.retile_D(tCrA);
copy(load_a, tXsA, tXrA);
gemm(mma, tCrA, tCrB, tCrC);

这几行代码连接了两件事：加载指令怎样写入数据，计算指令怎样读取数据。tCrA 准备线程局部的 A 片段，tXrA 为同一份存储提供 Copy 视图。真正的加载发生在 copy，随后 gemm 直接读取这份存储。

下面先用一个 warp 的小例子，把这条数据流走通，再回到官方 sgemm_sm80.cu 的完整配置。每个布局都可以用文末的程序亲手打印；暂不展开 swizzle 的地址运算和流水线调度。

源码对应 CUTLASS 提交 0b55a2f691d6。本文始终使用 FP16 输入、FP16 累加的 MMA；half 是一个 16 位浮点数。

1. 先固定一个 warp 的小例子

一个 warp 有 32 个线程，warp 内的线程编号叫 lane。我们先只计算一条 16×8×16 的 MMA：

auto mma = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{});
auto thr_mma = mma.get_slice(lane);

这条指令由 32 个 lane 共同计算：

D = A × B + C

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

CuTe 按 (M,K) 索引 A，按 (N,K) 索引 B。因此本文的 sA 是 16×16 tensor，sB 是 8×16 tensor；它们都让 K 方向连续：

sA layout = (16,16):(16,1)    地址偏移 = 16*m + k
sB layout = ( 8,16):(16,1)    地址偏移 = 16*n + k

这里用 shape:stride 表示 layout。Shape 是维度大小，stride 是沿该维度走一步时，元素偏移增加多少。下文的偏移以 half 为单位。

一条 MMA 中，每个 lane 只持有矩阵的一部分：

操作数	每 lane 的 half 数	打包后的 32 位寄存器数
A	8	4
B	4	2
C/D	4	2

这是所选指令固定的接口。A、B 合计使用 6 个 32 位寄存器；按原地累加、D 复用 C 的存储计算，加上累加器共 8 个。后面用 r[i] 表示一个 half 槽位，两个相邻 half 才组成一个 32 位操作数。

2. 先看 lane 0 的 A 片段里要放什么

当前线程通过下面两种调用，分别得到 shared memory 视图和线程局部片段：

auto tCsA_mma = thr_mma.partition_A(sA);
auto tCrA = thr_mma.partition_fragment_A(sA);

矩阵元素与操作数位置的关系，由所选 mma.sync 指令规定。CuTe 的 MMA_Traits 将这个硬件约定写成 ALayout/BLayout；对本文的指令，A 的定义为：

using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _2,_2,  _2>>,
                       Stride<Stride<_32,_1>,Stride<_16,_8,_128>>>;

它将 (lane,value) 映射为矩阵编号 q=m+16*k。展开嵌套坐标：

lane  t = t0 + 4*t1          t0=0…3，t1=0…7
value v = v0 + 2*v1 + 4*v2   v0、v1、v2 均为 0 或 1

q = 32*t0 + t1 + 16*v0 + 8*v1 + 128*v2
  = (t1 + 8*v1) + 16*(2*t0 + v0 + 8*v2)

m = t1 + 8*v1
k = 2*t0 + v0 + 8*v2

这就给出了“哪个 lane 的哪个值，应当保存哪个矩阵元素”。partition_A 在源码中通过 a_tensor.compose(AtomLayoutA_TV{},_) 将这张映射接入 A tensor，再选择当前线程；partition_fragment_A 据此准备同样值顺序的线程局部槽位。

取 lane 0，所需的 A 值以及最终进入的 32 位操作数为：

片段中的槽位	对应的 A 元素	shared memory 中的元素偏移	MMA 的 A 操作数：低 16 位、高 16 位
r[0]、r[1]	A(0,0)、A(0,1)	0、1	a0
r[2]、r[3]	A(8,0)、A(8,1)	128、129	a1
r[4]、r[5]	A(0,8)、A(0,9)	8、9	a2
r[6]、r[7]	A(8,8)、A(8,9)	136、137	a3

已知 Atom 内的 (m,k)，也可以反解出 lane 与值编号；下面 / 均为整数除法：

lane = 4*(m%8) + (k%8)/2
v    = k%2 + 2*(m/8) + 4*(k/8)

例如 A(0,9) 对应 t=0、v=5，因此使用 lane 0 的 half 槽位 5，最后进入 a2 的高 16 位。第 5 节会沿 mma_unpack 展开从 half 槽位到这些操作数的传参过程。

这些值在 shared memory 中并不连续。tCrA 则为它们安排了相邻的 8 个线程局部槽位：

逻辑值编号 v：          0   1    2    3   4   5    6    7
shared memory 偏移：    0   1  128  129   8   9  136  137
fragment 槽位偏移：     0   1    2    3   4   5    6    7

这时只是确定槽位与元素的对应关系，还没有从 shared memory 读取数据。 真正填入这些值，需要后面的 copy。

打印两者的 layout，可以得到以下结果。为了便于阅读，正文先省略静态整数的 _ 前缀：

tCsA_mma：((2,2,2),1,1):((1,128,8),0,0)
tCrA：   ((2,2,2),1,1):((1,  2,4),0,0)

最外层的三个 mode，也就是三个维度，分别表示：

(单条 MMA 中当前 lane 的 A 值, M 方向重复, K 方向重复)

第一个 mode 内部又拆成 (2,2,2)，共有 8 个值。后两个 mode 都是 1，因为我们只处理一条 MMA 的范围。CuTe 允许用整数 v=0…7 访问这个嵌套的第一个 mode：

tCrA(5, 0, 0);  // lane 0 中用于保存 A(0,9) 的 half 槽位

这里变化的是存储 stride；MMA 的逻辑元素分工没有改变。

3. make_fragment_like 怎样生成这 8 个槽位

顺着 partition_fragment_A 进入源码：

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

partition_A 先得到上一节的 shared memory 视图。对本文的 SM80 Atom，make_fragment_A 再调用：

return make_fragment_like<FrgTypeA>(atensor);

这里的 FrgTypeA 是 half 类型。Tensor 重载取出输入 layout，先生成 fragment layout，再调用 make_tensor<half_t> 创建存储。下面只看本例走到的“shape 静态、rank 大于 1”分支，将它拆成三步。

第一步：把 Atom 内的值排紧凑

第 0 个 mode 描述单条 MMA 的 A 值。源码先取出它的 shape、stride：

auto atom_shape = get<0>(layout.shape());
auto atom_stride = get<0>(layout.stride());

auto block = make_layout(
    atom_shape,
    compact_major<LayoutLeft>(filter_zeros(atom_stride, atom_shape)));

本例的输入是：

atom_shape  = (2,2,2)
atom_stride = (1,128,8)

compact_major<LayoutLeft> 按左侧子维度优先的顺序，重新生成紧凑 stride：

第一个子维度：1
第二个子维度：2
第三个子维度：2×2 = 4

block = (2,2,2):(1,2,4)

因此，三个子坐标 (v0,v1,v2) 对应的槽位偏移为：

v = v0 + 2*v1 + 4*v2

8 个值占据槽位 0…7，不再沿用 shared memory 中的间隔 128、8。

为什么生成新 stride，还要传入原 stride

因为 shape 只说明有多少个逻辑位置，不能说明这些位置是否共享同一个存储槽位。原 stride 为 0 的维度表示广播：沿这一维改变坐标，仍然访问同一个值。生成 fragment 时，需要保留这种关系。

将这一行拆开，原 stride 的用途就清楚了：

auto compact_shape = filter_zeros(atom_stride, atom_shape);
auto new_stride = compact_major<LayoutLeft>(compact_shape);
auto block = make_layout(atom_shape, new_stride);

filter_zeros(a,b) 在源码中的规则是：递归检查 a，将与静态 Int<0> 对应的 b 项替换为 Int<1>，其余项保留。因此它返回的是用于生成紧凑 stride 的辅助 shape。本例没有零 stride：

原 atom_shape：                    (2,2,2)
原 atom_stride：                   (1,128,8)
filter_zeros 后的 compact_shape：  (2,2,2)
compact_major 生成的 new_stride：  (1,2,4)

原 stride 的 128、8 不参与大小排序，也不参与新 stride 的乘积；它们只需要通过“是否为静态 0”的检查。Atom 内的紧凑顺序由 LayoutLeft 决定。

再看一个带广播的示意布局。设第二个子维度的 stride 是静态 _0：

原 atom_shape：                    (2,2,2)
原 atom_stride：                   (1,0,8)
filter_zeros 后的 compact_shape：  (2,1,2)
compact_major 生成的 new_stride：  (1,0,2)
最终 block：                       (2,2,2):(1,0,2)

为什么中间的 shape 改成 1，最后又保留 2？辅助 shape 中的 1 表示这一维不增加独立存储。compact_major 遇到静态大小 1 时，生成 stride 0，并保持当前累计跨度不变；所以第三维的 stride 是 2。最终 make_layout 仍使用原 atom_shape，保留全部逻辑坐标：

block(0,0,0) = 0
block(0,1,0) = 0    // 第二维变化，仍引用同一槽位

逻辑位置共 2×2×2 = 8 个，实际只需 4 个 half 槽位。

如果直接对原 shape 调用 compact_major<LayoutLeft>(atom_shape)，会得到 (1,2,4)，上面两个位置将分别映射到槽位 0、2，广播关系就丢失了。

这里的判断限于编译期零 stride。若 stride 的类型是运行时 int，即使当前数值为 0，filter_zeros 也不会将它作为静态广播维处理。配套程序会打印这两种情况，便于对照源码。

第二步：给外层重复生成布局

源码将第 0 个 mode 以外的部分交给 make_ordered_layout：

auto rest = make_ordered_layout(
    take<1,R>(layout.shape()),
    take<1,R>(layout.stride()));

我们的小例子没有 M、K 重复，所以：

rest_shape = (1,1)
rest       = (1,1):(0,0)

处理更大的 tile 时，rest 才会安排多份 Atom 片段的顺序。先记住它的作用，文末的源码补充再用官方完整配置展开。

第三步：接起来，再创建存储

auto fragment_layout = tiled_product(block, rest);

这一步将 Atom 内布局与外层重复接起来。本例的外层大小都是 1，结果就是：

fragment_layout = ((2,2,2),1,1):((1,2,4),0,0)

外层包装随后调用：

make_tensor<half_t>(fragment_layout);

它创建拥有存储的 Tensor，内部使用 ArrayEngine<T, cosize_v<Layout>>。这里 cosize 为 8，因此准备 8 个 half 槽位。它没有复用 shared memory 的指针，也没有加载原数据；这是静态大小的线程局部存储表达，最终由编译器分配寄存器。

至此，make_fragment_like 完成了两件事：保留当前线程所需值的 shape，为这些值重新安排紧凑的存储布局。

4. 让 ldmatrix 直接写入这份片段

现在已经有 8 个 A 槽位，接下来选择每 lane 输出 8 个 half 的 ldmatrix.x4：

auto load_a = Copy_Atom<SM75_U32x4_LDSM_N, half_t>{};
auto copy_a = make_tiled_copy_A(load_a, mma);
auto thr_copy_a = copy_a.get_slice(lane);

auto tXsA = thr_copy_a.partition_S(sA);
auto tXrA = thr_copy_a.retile_D(tCrA);

make_tiled_copy_A 接入 mma.get_layoutA_TV()，以 MMA 的 A 值需求构造 Copy 分工。所选 ldmatrix 的 reference 等于目的端，因此加载后的目的分工可以与 MMA 的需求对齐。

partition_S 确定源地址，retile_D 确定目的槽位

tXsA 是 shared memory 的源视图，告诉当前 lane 应提供哪些地址。tXrA 是线程局部的目的视图，告诉加载指令应写入哪些槽位。

对于 ldmatrix，两端不能按普通的逐线程复制理解：源 lane 提供行地址，硬件再将数据分发到各个目的 lane 的寄存器。 例如源 lane 0 提供 A(0,0…7) 所在的行地址，其中 A(0,2) 会进入目的 lane 1 的片段；目的 lane 0 收到的则是第 2 节列出的那 8 个值。

为什么 Copy 的分工恰好能对上 MMA

这里的对应关系来自构造过程。tCrA、TiledCopy 的线程和值分工、tXrA 的存储布局之间，存在明确的依赖链：

同一个 mma
  ├─ partition_fragment_A(sA) → tCrA：为各 lane 的 MMA A 值准备槽位
  └─ get_layoutA_TV()         → TiledCopy 的 reference TV
                                      │
tCrA.layout() ─────────────────────────┤
                                      ▼
                           retile_D(tCrA) → tXrA

先看 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)));

它将 MMA 的 A 布局直接作为 TiledLayout_TV。因此，TiledCopy 的 reference 线程和值编号，一开始就按这个 MMA 所需的 A 元素来定义。

这里还需要解释：MMA 与 Copy 各有自己的 Atom TV 布局，为什么能够衔接？先区分它们输出的编号：

布局	输入	输出
mma.get_layoutA_TV()	整个 TiledMMA 的 (thread,value)	A tile 的矩阵编号 q
Copy 的 AtomLayoutRef	一个 Copy Atom 的 reference (thread,value)	Copy 内部元素编号 e
Copy 的 AtomLayoutSrc	一个 Copy Atom 的源 (thread,value)	同一个 Copy 内部元素编号 e
Copy 的 AtomLayoutDst	一个 Copy Atom 的目的 (thread,value)	同一个 Copy 内部元素编号 e

q 与 e 的含义不同，数值不需要相等。right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{}) 只在 Copy Atom 内部消去 e，得到“目的 TV → reference TV”。它的输出随后接到 MMA 布局的 TV 输入侧。

源码 tile2thrfrg 为此先做切分，再接入转换：

auto atom_layout_TV = zipped_divide(
    TiledLayout_TV{}, make_shape(AtomNumThr{}, AtomNumVal{}));

auto dst_to_ref = right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{});
auto trg_layout_TV = atom_layout_TV.compose(dst_to_ref, _);

第一步切的是 (thread,value) 的编号范围，切分大小取自 Copy Atom。它不要求 MMA 的 Atom 嵌套 shape 与 Copy 的嵌套 shape 完全相同。例如，省略大小为 1 的重复 mode 后，本例 MMA 的线程 shape 是 (4,8)、值 shape 是 (2,2,2)；它们分别可以用整数 t=0…31、v=0…7 编号。ldmatrix.x4 的 reference shape 则是 (32,(2,4))，同样有 32 个线程、每线程 8 个值。CuTe 可以在这些嵌套坐标和整数编号之间转换。

忽略嵌套结构，只按各 mode 的总大小看，当前 A 可以这样切：

Copy Atom	一个 Atom 的 (T,V)	A 的 (32,8) 切分后剩余组数 (rest_T,rest_V)
ldmatrix.x4	(32,8)	(1,1)
ldmatrix.x2	(32,4)	(1,2)
ldmatrix.x1	(32,2)	(1,4)

切分后的输入形式为：

((atom_thread,atom_value),(rest_thread,rest_value))

t = atom_thread + AtomNumThr*rest_thread
v = atom_value  + AtomNumVal*rest_value

原 MMA 映射到矩阵位置的规则仍然保留。第二步 compose(dst_to_ref,_) 只将第一个 mode 换成 Copy 的目的 TV，外层重复组保持不变。整条映射是：

Copy 目的局部 TV
  → AtomLayoutDst：Copy 内部元素 e
  → right_inverse(AtomLayoutRef)：reference 局部 TV
  → 接回所在的 rest 组，恢复整个 TV
  → TiledLayout_TV：矩阵位置 q

目的分工由此确定，源地址也必须接入同一套矩阵坐标。tidfrg_S 使用的是：

return tile2thrfrg(
    zipped_divide(stensor, Tiler_MN{}),
    right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{}));

Copy Atom 自身的 e 是内部元素编号；源 tensor 中读取哪个矩阵坐标，由上述转换再接入 TiledLayout_TV 决定。 copy_traits.hpp 的注释也明确说明，源、目的、reference 共享的抽象 bit 编号可以任意选择，它只用于构造各端 TV 之间的映射。

先看本例：一个 Copy Atom 恰好覆盖全部 reference TV，可以省略外层重复组。记 F=TiledLayout_TV、S=AtomLayoutSrc、D=AtomLayoutDst、R=AtomLayoutRef，两端最终得到的矩阵编号是：

源侧：qS = F(right_inverse(R)(S(src_tv)))
目的：qD = F(right_inverse(R)(D(dst_tv)))

Copy 指令将同一个内部元素从源传到目的：
S(src_tv) = D(dst_tv) = e

因此：qS = F(right_inverse(R)(e)) = qD

这里保证相等的是最终的矩阵坐标。Copy 内部编号 e 与矩阵编号 q 可以不同，因为源、目的都经过同一个 right_inverse(R) 和 F。对当前 reference 等于目的端的 ldmatrix，目的 TV 还直接保持为 MMA 所需的 TV；retile_D 再把这些值接到已有 fragment 槽位。

用 A(0,9) 走完这条路径。转换成 half 单位后，本例 ldmatrix.x4 的 Atom 布局为：

AtomLayoutSrc = (32,8):(8,1)
AtomLayoutRef = AtomLayoutDst = (32,(2,4)):(2,(1,64))

源 (tS,vS)=(16,1)：S(16,1) = 8*16+1 = 129
目的 (tD,vD)=(0,5)：D(0,5) = 2*0+1+64*2 = 129

共同的 e=129
  → right_inverse(R)：reference (t,v)=(0,5)，tv=160
  → F：矩阵编号 q=144=0+16*9，即 A(0,9)

所以 partition_S 为源 lane 16 安排的这一行是 A(0,8…15)；行内编号为 1 的 half 就是 A(0,9)，编号从 0 开始。按本文 shared memory 布局，行首偏移为 8，元素偏移为 8+1=9。ldmatrix 执行时，lane 16 提供这一行的起始地址，硬件将其中的值分发给相应目的 lane；源值编号 1 描述行内元素，并不要求 lane 16 另传一个地址。

完整的数据路径因此是：

shared memory：A(0,9)，元素偏移 9
  → 源 lane 16 所提供行内编号为 1 的 half
  → ldmatrix 按 e=129 对应的硬件路由分发
  → 目的 lane 0 的 half 槽位 5
  → retile_D 对应的 tCrA(5,0,0)
  → MMA 的 a2 高 16 位，硬件约定其含义为 A(0,9)

这条路径同时出现了 e=129、tv=160、q=144 和内存偏移 9。它们分别用于 Copy 内部元素匹配、reference TV 编号、矩阵坐标编号和 shared memory 寻址；各步的布局组合将它们连接起来。

切分也有明确约束。TiledCopy 在源码中检查：

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

所以并非任意搭配都能切分。单条 16×8×16 MMA 的 B 每 lane 只有 4 个值，直接搭配每 lane 输出 8 个 half 的 ldmatrix.x4，会因 4 % 8 != 0 无法通过该检查。这也是前面小例子的 B 选择 x2，而官方配置扩展 N 方向的参考 tile 后才可以使用 x4 的原因。这里检查的是整个 TiledLayout_TV 的值数，不要求单条 MMA Atom 的值数一定等于 Copy Atom 的值数。

此外，源码还假定一个 Copy Atom 的线程在 TiledCopy 线程编号中连续且对应一致。真正执行 ldmatrix 时，推导出的源行还需要满足连续性和 16 字节地址对齐；数量能整除只解决编号范围的切分问题。本文选用的 A 布局满足这些条件，配套程序可核对源行与目的坐标。

最后，retile_D(tCrA) 要将这些目的值落到已有槽位。它直接使用 tCrA.layout()，在源码中通过 zipped_divide 划分原布局，再通过 compose 接上 Copy 分组到原 fragment 坐标的映射。将原 fragment 布局记为 L_frag，坐标转换记为 P，新布局就是：

Copy 坐标 c → P(c)：原 fragment 坐标 → L_frag(P(c))：原存储偏移

tXrA 的 layout = L_frag ∘ P

所以新布局访问的位置由原布局推导而来，并不需要独立生成一套存储排列，再碰巧与 tCrA 一致。当前 A 例子中，Copy 与 MMA 每 lane 都是 8 个值，转换只需将 (8,1) 对应到 (2,2,2)；后文 B 的数量不同，就需要跨两个 MMA 片段重新分组。

这条构造链依赖正确的输入语义。retile 的源码明确假设：输入第 0 个 mode 的 V 个值，对应 TiledLayout_TV 中最前面的 V 个 value ID。这里使用同一个 mma 创建 fragment 和 TiledCopy，正是为了使这些 value ID 的含义一致。

retile_D 怎样复用已经准备好的槽位

先区分 Tensor 的两个部分：engine 决定数据存放在哪里，layout 决定怎样用坐标访问数据。对本例的 half Tensor，元素访问最终执行的是：

return data()[layout()(coord)];

上一小节说明了两种分工如何对应。存储层面，retile_D 再将新的 layout 接到 tCrA 的数据起点，其源码为：

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

代入当前变量，可以拆成：

auto copy_layout = decltype(copy_a)::retile(tCrA.layout());
auto tXrA = make_tensor(tCrA.data(), copy_layout);

retile 在这里接收的是 tCrA.layout()，负责算出新的坐标到偏移的映射；存储起点则直接取自 tCrA.data()。make_tensor(iterator, layout) 走源码中的视图构造分支，用 ViewEngine 保存这个迭代器。

回顾上一节，make_tensor<half_t>(fragment_layout) 创建的是带 ArrayEngine 的存储。两次构造的区别是：

Tensor	engine 保存什么	作用
tCrA	ArrayEngine 内的 8 个 half 槽位	持有当前 lane 的 A 数据，按 MMA 规则访问
tXrA	指向上述槽位的迭代器	为这些槽位提供 Copy 所需的索引方式

所以 auto tXrA = ... 虽然创建了一个新的 Tensor 对象，复制进去的是存储起点和布局信息，并未再创建 8 个 half 数据槽位。这里的共享也发生在同一个 lane 内部：每个 lane 都为自己的 tCrA 创建对应视图。

“第一维的分组”具体改了什么

对当前的 A 与 ldmatrix.x4，实际布局是：

MMA 视图 tCrA：((2,2,2),1,1):((1,2,4),0,0)
Copy视图 tXrA：((8,1),  1,1):((1,0),  0,0)

两个 Tensor 的外层 rank 都是 3。这里说的“第一维”，是嵌套的第 0 个 mode：MMA 使用 (2,2,2) 描述一份 A 片段，Copy 使用 (8,1) 描述“一次 Copy 的 8 个值，共 1 组”。它们的大小都为 8；后面的 M、K 重复维在这个小例子中都为 1。

展开各自的偏移计算：

MMA 第一维坐标 (a,b,c)：offset = a + 2*b + 4*c
Copy第一维坐标 (v,g)：  offset = v + 0*g，且 g 只能为 0

令 v=a+2*b+4*c，两边就访问相同的 half 槽位。例如：

MMA 的第一维坐标	Copy 的第一维坐标	共同访问的槽位
(0,0,0)	(0,0)	r[0]
(0,1,0)	(2,0)	r[2]
(1,0,1)	(5,0)	r[5]

最后一行可以直接写成：

&tCrA(make_coord(1,0,1), 0, 0)
    == &tXrA(make_coord(5,0), 0, 0);

CuTe 也支持用整数访问嵌套 mode。将第一维展平后，对任意 v=0…7：

&tXrA(v,0,0) == &tCrA(v,0,0);

第一维的分组决定了 Copy Atom 每次取得哪些目的槽位。retile 内部按 AtomNumVal 划分值；本例 reference 等于目的端，ldmatrix.x4 的 AtomNumVal 是 8，因此分成 (8,1)。如果改用每 lane 输出 2 个 half 的 ldmatrix.x1，同一份 A 存储就会分成 (2,4)：

x4 的 Copy 视图：一次取 r[0…7]
x1 的 Copy 视图：依次取 r[0…1]、r[2…3]、r[4…5]、r[6…7]

这也说明分组的用途：Copy 根据自己的指令粒度逐组写入，MMA 再按自己的片段规则读取。以槽位 5 为例，在配套主机程序中可以直接验证：

tXrA(make_coord(5,0), 0, 0) = half_t(7);
// 此时 tCrA(make_coord(1,0,1), 0, 0) 读到的也是 7。

真正加载时，copy(load_a,tXsA,tXrA) 直接写入这些槽位，gemm 通过 tCrA 读取它们。两种坐标形式之间的对应关系已经由 layout 确定，数据无需在两个 Tensor 之间再搬一次。

“只改第一维分组”描述的是当前 A 例子的结果。后文官方配置的 B 会将两个 MMA 片段合成一组 Copy，届时外层 N 重复次数也会随之变化。

5. 把 A、B 加载与 MMA 接起来

小例子的 B 每 lane 只需要 4 个 half，所以选用 ldmatrix.x2，正好准备一份 B 片段：

auto tCrB = thr_mma.partition_fragment_B(sB);

auto load_b = Copy_Atom<SM75_U32x2_LDSM_N, half_t>{};
auto copy_b = make_tiled_copy_B(load_b, mma);
auto thr_copy_b = copy_b.get_slice(lane);

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

累加器则由当前线程的 C 分区创建，再初始化为 0：

auto tCgC = thr_mma.partition_C(gC);
auto tCrC = thr_mma.make_fragment_C(tCgC);
clear(tCrC);

make_fragment_C 使用 C 分区的 shape 创建累加器，不会读取 global C；当前实现也不沿用 global C 的 stride。这里从零开始累加乘积。

假设 shared memory 中的 A、B 已经准备好，warp 中的所有 lane 共同执行：

// shared memory → 各 lane 的寄存器
copy(load_a, tXsA, tXrA);
copy(load_b, tXsB, tXrB);

// 直接使用刚才写入的同一份寄存器
gemm(mma, tCrA, tCrB, tCrC);

每个 lane 都参加 A 的加载，也都参加 B 的加载。加载完成后的操作数关系为：

ldmatrix.x4 → A 的 8 个 half → 4 个 32 位操作数
ldmatrix.x2 → B 的 4 个 half → 2 个 32 位操作数
                                  ↓
                        mma.sync 使用 A、B、C

tXrA 和 tCrA 之间没有额外一次复制。Copy 通过前者写入，MMA 通过后者读取。

fragment 的槽位怎样成为 MMA 的寄存器操作数

gemm 遍历 M、N、K 的重复片段，将一份 A、B 交给 MMA Atom。Atom 的 call 随后调用 mma_unpack。对本文指令，其寄存器类型来自架构封装：

using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];

mma_unpack 取出数组的元素类型，再执行：

Tensor rA = recast<RegTypeA>(A);  // 本例 RegTypeA = uint32_t
Tensor rB = recast<RegTypeB>(B);  // 本例 RegTypeB = uint32_t

这里是按位重解释。每两个相邻 half 组成一个 32 位操作数，原来的 FP16 位模式保留。A 的 8 个 half 因而成为 4 个操作数，B 的 4 个 half 成为 2 个操作数。detail::explode 按序展开后，传给架构封装的参数相当于：

MMA_Op::fma(d0, d1,
            rA(0), rA(1), rA(2), rA(3),
            rB(0), rB(1),
            c0, c1);

架构封装再将这些参数依次绑定到 mma.sync 的 A、B 操作数列表。第 2 节的 lane 0 因此按 r[0:1] → a0、r[2:3] → a1、r[4:5] → a2、r[6:7] → a3 传入。这里的 r[i:j] 表示两个 half 槽位。

B 采用相同的过程，只是其硬件布局不同。在 CuTe 的 (n,k) 索引下，令 t=lane、v=0…3：

n = t/4
k = 2*(t%4) + v%2 + 8*(v/2)

反解：lane = 4*n + (k%8)/2，v = k%2 + 2*(k/8)

所以 lane 0 的 B 操作数为：

B 操作数	低 16 位	高 16 位
b0	B(0,0)	B(0,1)
b1	B(0,8)	B(0,9)

a2、b1 表示当前 lane 在本条指令中的操作数位置；具体物理寄存器由编译器分配。指令规定的是“lane 编号、操作数序号、低/高 16 位”分别表示哪个矩阵元素。执行时，mma.sync 只接收这些寄存器操作数，不接收矩阵坐标或 CuTe layout；调用方必须事先把正确的值放到规定位置。

因此，从矩阵到计算的对应关系完整地连接为：

硬件约定 → MMA_Traits 的 ALayout/BLayout
        → partition_A/B 确定每个逻辑值对应的矩阵坐标
        → fragment 提供槽位，Copy 按该对应关系填入值
        → mma_unpack 按位打包并依次传参
        → mma.sync 按硬件约定解释各 lane 的操作数

加载完成与共享内存同步

需要事先保证 shared memory 数据可见。官方内核在读取相应 stage 前使用 cp_async_wait 和 __syncthreads()；ldmatrix.sync 不替代这一步共享内存同步。A、B 两次加载之间则不需要额外的 CTA 屏障，后续 MMA 的寄存器依赖保证它使用已就绪的输入。

6. 再回到官方的完整 CTA

小例子走通后，再看 sgemm_sm80.cu 的配置：

// CTA tile：128×128×64
// 4 个 warp 沿 M、N 排成 2×2
auto mma = make_tiled_mma(
    SM80_16x8x16_F16F16F16F16_TN{},
    Layout<Shape<_2,_2>>{},
    Tile<_32,_32,_16>{});

单条 MMA 的接口没有改变，线程需要保存的片段份数增加了：

MMA_M = 128 / (16×2) = 4
MMA_N = 128 / ( 8×2) = 8
MMA_K =  64 / 16     = 4

把各 mode 内部的嵌套大小相乘，每线程的 fragment 大小为：

片段	mode 含义	各 mode 的大小	half 总数
tCrA	(MMA,MMA_M,MMA_K)	(8,4,4)	128
tCrB	(MMA,MMA_N,MMA_K)	(4,8,4)	128
tCrC	(MMA,MMA_M,MMA_N)	(4,4,8)	128

源码使用 sA(_,_,0)、sB(_,_,0) 创建 A、B fragment，是取一个 stage 的二维形状与布局来推导存储，并没有在此读取 stage 0，也没有限定以后只能加载它。

为什么官方 B 也选 ldmatrix.x4

官方对 A、B 都使用 ldmatrix.x4。B 沿 N 方向有多份片段，可以让一次加载准备两份：

MMA 视图：每份 4 个值，N 方向 8 份
Copy视图：每组 8 个值，N 方向 4 组

retile_D 建立的对应关系为：

tXrB(v,j,k) ↔ tCrB(v%4, 2*j+v/4, k);

例如固定 j=k=0：

Copy 第 0 组： [r0 r1 r2 r3 r4 r5 r6 r7]
MMA 的 B 片段：[r0 r1 r2 r3] [r4 r5 r6 r7]
                    N=0           N=1

这里的 N=0、1 是 B 片段编号。一次 ldmatrix.x4 填好两份片段，随后两条 MMA 分别使用自己的 4 个 B 值，并复用同一份 A。

固定一个 K 分块，完整调用会展开成：

调用	每个 warp 的指令次数	完成的工作
A 的 copy	4 次 ldmatrix.x4	准备 4 份 A
B 的 copy	4 次 ldmatrix.x4	准备 8 份 B
gemm	32 次 mma.sync	组合 4 份 A 与 8 份 B，更新对应 C 片段

因此，代码里的一次 A Copy、一次 B Copy，不等于总共只有两条加载指令。copy 会遍历片段的剩余分组，逐组调用 Copy Atom。

原始内核还将 k_block_next 的预取与 k_block 的计算交错执行。本文按同一 K 分块的数据依赖来说明；预取调度不改变“Copy 写入、MMA 读取同一份 fragment”的关系。

每线程三份完整 fragment 各含 128 个 half。按两个 half 打包到一个 32 位寄存器，其数据容量合计为 64+64+64=192 个寄存器。这与单条 MMA 的 8 个操作数寄存器是不同口径；编译器实际占用还取决于存活期、复用、地址与临时变量及 spill，应查看编译报告。

7. 亲手打印一次布局

配套程序：cute_mma_ldmatrix_registers_trace.cu。它包含上面的小例子与官方完整配置，在主机上打印布局、核对逻辑坐标和存储别名，不执行 GPU 加载或 MMA 指令。

在 CUTLASS 根目录编译运行，CUDA 头文件路径按安装位置调整：

g++ -x c++ -std=c++17 \
  -Iinclude -I/usr/local/cuda/include -I/usr/local/cuda/include/cccl \
  cute_mma_ldmatrix_registers_trace.cu \
  -o /tmp/cute_mma_ldmatrix_registers_trace

/tmp/cute_mma_ldmatrix_registers_trace

可以先对照小例子的几行输出：

partition_A layout = ((_2,_2,_2),_1,_1):((_1,_128,_8),_0,_0)
block = (_2,_2,_2):(_1,_2,_4)
rest = (_1,_1):(_0,_0)
fragment layout = ((_2,_2,_2),_1,_1):((_1,_2,_4),_0,_0)

观察到的变化应与第 2、3 节一致：shape 保留，shared memory 的间隔被替换为紧凑的片段 stride。程序还打印静态零 stride 与运行时零 stride 的处理结果，核对 32 个 lane 的 Copy 目的坐标与 MMA 坐标，以及两种寄存器视图的别名关系。它会通过 Copy 视图写入槽位 5，再通过 MMA 视图读回，并验证 ldmatrix.x1 将同一份 A 存储分成四对槽位。

针对第 4 节的 Atom 衔接问题，程序会打印 A 的 TV 按 x4/x2/x1 粒度切分后的布局，以及 Copy 内部元素编号、reference TV 编号和矩阵编号的转换；还会核对本例 x4 源行的连续性与 16 字节对齐。对全部 256 个 A 元素，程序验证相同 Copy 内部编号在源、目的两端得到相同矩阵坐标，并在主机上按这条路由读取源值、写入对应 fragment 槽位；它也会打印上面的 A(0,9) 路径。

针对第 5 节的操作数绑定，程序还会核对 A、B 坐标到 lane/value 的反解，以及 recast<uint32_t> 后每个操作数的低、高 16 位。这里验证的是主机上的布局与位模式，不执行 GPU 指令。

接着可以做两个小练习：

1. 在小例子中，将 A 的行跨度从 16 改为 24，同时将占位数组扩大为 16×24。重新运行后，partition_A 中的 stride 128 会变为 192，tCrA 的布局仍然相同。
2. 在官方完整配置中，用 tXrB(v,j,k) 与 tCrB 的对应公式，检查 (v,j,k)=(6,1,0)。它应访问 tCrB(2,3,0)，也就是 B fragment 中偏移为 14 的 half 槽位。

源码补充：外层重复的 stride 怎样生成

理解了主线，再展开第 3 节中尚未使用的外层重复。取官方配置的线程 0，实际打印结果为：

partition_A layout：
((_2,_2,_2),_4,(_2,_2)):((_1,_512,72),_2048,(144,288))

这里特意保留 _：_2048 是编译期常量，144、288 是运行时值。我们只分析该布局如何生成 fragment，不展开 shared memory 地址中的 swizzle 运算。

Atom 内的处理与小例子相同：

block = (2,2,2):(1,2,4)

外层输入为：

rest_shape = (4,(2,2))
order      = (_2048,(144,288))

make_ordered_layout 使用原 stride 提供的顺序信息生成紧凑 stride，不保留原来的数值。源码规定：静态 order 按数值比较，动态 order 视为更大，并按出现顺序排列。本例于是先安排 M，再安排 K 的两个子维度：

rest = (4,(2,2)):(1,(4,8))

不能只看打印出的 144 < 2048 就判断 K 先变化，二者的类型也参与决定顺序。

接着，tiled_product(block,rest) 让外层重复跨过完整的 Atom 片段。每份 block 占 8 个 half，外层 stride 因而乘以 8：

M stride：1×8 = 8
K stride：(4×8,8×8) = (32,64)

tCrA = ((2,2,2),4,(2,2)):((1,2,4),8,(32,64))

用整数 k=0…3 访问嵌套 K mode，可以写出 A、B 的槽位偏移：

offset_A(v,m,k) = v + 8*m + 32*k
offset_B(v,n,k) = v + 4*n + 32*k

make_fragment_like 单独优先处理第 0 个 mode，也区别于一般的 make_tensor_like。后者会尽可能参考所有维度的原始顺序；对这里的 A 分区，它得到的 stride 是 ((1,2,16),4,(32,64))，Atom 内的 8 个值分散在槽位 0…3、16…19。make_fragment_like 优先将 Atom 内值按 LayoutLeft 紧凑排列，再组织外层重复，便于按指令片段访问与向量化。

源码索引

• 官方 sgemm_sm80 的 fragment、Copy 视图与主循环
• MMA Atom 的固定寄存器接口
• 所选 MMA 的 ALayout/BLayout
• partition_A/B 如何接入 MMA 的 TV 布局
• mma_unpack 如何重解释并展开操作数
• make_fragment_A/B 如何调用 make_fragment_like
• make_fragment_like 的 Tensor 包装与存储创建
• Tensor 如何根据 layout 访问数据
• retile 如何重新组织 fragment 的值坐标
• TiledCopy 的整除检查、Atom 切分与端点转换
• Copy 内部元素编号与 reference 的含义
• ldmatrix.x4 的源、目的与 reference Atom 布局
• make_fragment_like 的布局计算
• filter_zeros 如何将静态零 stride 对应的 shape 改为 1
• compact_major 如何生成紧凑 stride 并处理静态大小 1
• compact_order 如何处理动态顺序信息
• retile_D 与 make_tiled_copy_A/B