读懂 CuTe 的 make_tiled_mma：一条 MMA 怎样铺成一块矩阵

读 CuTe 的矩阵乘代码时，经常会遇到这样的配置：

auto mma = make_tiled_mma(
    SM80_16x8x16_F16F16F16F16_TN{},
    Layout<Shape<_2,_2,_1>>{});

指令名里有 16×8×16，后面又跟着一个 2×2×1。它们相乘以后，是一条更大的 Tensor Core 指令，还是几组线程各算一部分？线程 32 和线程 64 各自算哪里？后面若再加一个 Tile<_64,_32,_16>，线程数量会不会继续增加？

先把第一个问题的结果画出来。上述配置使用四个 warp，共 128 个线程，在一个 K 分块中计算 32×16 的输出：

                        N = 0～7             N = 8～15
                   ┌──────────────────┬──────────────────┐
M = 0～15          │ warp 0，线程 0～31 │ warp 2，线程 64～95│
                   │ 一块 16×8 的输出   │ 一块 16×8 的输出  │
                   ├──────────────────┼──────────────────┤
M = 16～31         │ warp 1，线程 32～63│ warp 3，线程 96～127│
                   │ 一块 16×8 的输出   │ 一块 16×8 的输出  │
                   └──────────────────┴──────────────────┘

每个 warp 的这一步，都沿 K 方向使用 16 个元素。

每个 warp 执行的仍是原来的 m16n8k16 MMA。make_tiled_mma 把指令本身的 lane 分工、多个 warp 的排列和后续重复计算的坐标组织起来，让线程可以对 A、B、C 取得相互配合的访问视图。

下面先看一条 MMA 内部的分工，再沿源码扩展到四个 warp、较大的输入矩阵和第三个参数。配套程序会逐项核对文中的坐标。

先固定矩阵和方向

本文沿用前面几篇文章的指令：

using Op = SM80_16x8x16_F16F16F16F16_TN;

它用 FP16 的 A、B 计算 FP16 的结果，输出块为 16×8，一次沿 K 累加 16 项。本文讨论的是 SM80 这条由一个 warp 的 32 个 lane 协作执行的指令；其他架构的 MMA，其线程组和 fragment 规则要看相应 trait。

CuTe 对三个操作数的逻辑方向约定为：

A：(M,K)
B：(N,K)
C：(M,N)

D(m,n) = C(m,n) + Σ A(m,k) * B(n,k)

这里 B 的第一个坐标是 N。与通常把右矩阵写成 B_math(K,N) 的数学记法对应，有 B(n,k)=B_math(k,n)。这个坐标约定与实际内存怎样排列分开考虑。

指令名的 TN 对应底层 MMA 的操作数形式，不是给任意 Tensor 设置内存 stride 的开关。后面给 Tensor 配置按行连续或按列连续的 layout，才决定这些坐标落到哪里。

本文用 shape:stride 表示布局。例如 (2,2,1):(1,2,4) 将坐标 (wm,wn,wk) 映射为 wm+2*wn+4*wk。大小为 1 的维度只能取坐标 0，其 stride 不影响合法坐标的结果。CuTe 将一个维度称为 mode，mode 也可以继续嵌套。

一条 MMA 已经规定了 lane 持有哪些值

先只创建一个 warp 的配置：

auto mma_one = make_tiled_mma(Op{});

这时基本 tile 为 (16,8,16)。一条指令需要的逻辑元素数为：

操作数	矩阵大小	每 lane 的逻辑元素数
A	16×16=256	8 个 half
B	8×16=128	4 个 half
C / D	16×8=128	4 个 half

这些值的分布由 MMA_Traits<Op> 规定。构造 TiledMMA 时，不能像普通 copy 那样随意把“每线程的小矩形”传进去；底层指令已经要求特定 lane 持有特定值。

从 C 的四个值开始读

本例 C 的 Atom 布局是：

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

它将 (lane,value) 映射为 16×8 输出块的逻辑坐标编号。CuTe 按 shape 解码整数坐标时，第一个 mode 先变化，因此这个编号是：

q = m + 16*n

这还不是某个按行连续存储矩阵的内存偏移。

令 lane 编号为 l，值编号为 v。在线程 shape (4,8) 和值 shape (2,2) 中，标量坐标分别被拆成：

lane  → (l%4, l/4)
value → (v%2, v/2)

代入 stride：

q = 32*(l%4) + l/4 + 16*(v%2) + 8*(v/2)
  = [l/4 + 8*(v/2)] + 16*[2*(l%4) + v%2]

所以：

m = l/4 + 8*(v/2)
n = 2*(l%4) + v%2

l = 0～31，v = 0～3

几个 lane 的四个输出位置为：

lane	值 0	值 1	值 2	值 3
0	C(0,0)	C(0,1)	C(8,0)	C(8,1)
1	C(0,2)	C(0,3)	C(8,2)	C(8,3)
4	C(1,0)	C(1,1)	C(9,0)	C(9,1)
31	C(7,6)	C(7,7)	C(15,6)	C(15,7)

每个 lane 持有相隔 8 行的两对元素。32 个 lane 合起来，正好覆盖整个 16×8 输出块。

A 的八个值和 B 的四个值

同一个 trait 中，A、B 的布局为：

ALayout = ((4,8),(2,2,2)):((32,1),(16,8,128))
BLayout = ((4,8),(2,2)):((16,1),(8,64))

A 的输出编号按 qA=m+16*k 解码，B 则按 qB=n+8*k 解码。与 C 同样展开，可得到：

A，v = 0～7：
  m = l/4 + 8*((v/2)%2)
  k = 2*(l%4) + v%2 + 8*(v/4)

B，v = 0～3：
  n = l/4
  k = 2*(l%4) + v%2 + 8*(v/2)

例如 lane 0 的需求是：

A 的值 0～7：
  (0,0), (0,1), (8,0), (8,1),
  (0,8), (0,9), (8,8), (8,9)

B 的值 0～3，坐标方向为 (N,K)：
  (0,0), (0,1), (0,8), (0,9)

这里的“值”是逻辑 half 元素。该指令的底层接口将两个 half 打包在一个 32 位寄存器参数中，因此每 lane 的 A 对应四个 uint32_t 参数，B 和 C/D 各对应两个。不能把“8 个 A 值”直接读成“8 个 32 位寄存器”。

make_tiled_mma 的构造函数做了什么

函数可以接收一个底层 Operation，也可以接收已经封装的 MMA_Atom。传入 Op{} 时，会先转到：

return make_tiled_mma(MMA_Atom<MMA_Op>{}, thr_layout, permutations);

接收 Atom 的重载主体是：

auto thr_layout_mnk  = append<3>(thr_layout, Layout<_1,_0>{});
auto permutation_mnk = append<3>(permutations, _);

return TiledMMA<MMA_Atom<MMA_Op>,
                decltype(thr_layout_mnk),
                decltype(permutation_mnk)>{mma_atom, thr_layout_mnk};

可以把三个参数的职责先列出来：

参数	本例	规定什么
Atom / Operation	Op{}	指令形状、参与线程，以及 A/B/C 的 lane-value 分工
线程布局	Layout<Shape<_2,_2,_1>>{}	Atom 的线程组沿 M、N、K 怎样排列、怎样编号
permutation / tile	默认 (_,_,_)	各方向参与分块的范围，以及范围内部的坐标顺序

append<3> 用来把不足三维的配置补到 M、N、K 三个方向。线程布局缺少的方向补成 1:0，permutation 缺少的方向补成 _。因此 Layout<Shape<_2,_2>>{} 在这里也可以表示沿 M、N 各放两组，K 方向只有一组。

构造过程没有接收 A、B、C 指针，也没有执行加载或乘加。它建立后续分区所使用的布局与类型。线程创建、block 的大小和 grid 的划分，仍然由调用它的 kernel 决定。

第二个参数怎样把一个 warp 扩展成四个

回到开头的配置：

auto mma = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>>{});

这里的 (2,2,1) 描述 Atom 线程组的排列。对于本文的 SM80 指令，一个 Atom 线程组恰好是一个 warp。默认紧凑布局让 M 方向先变化：

warp_id = wm + 2*wn

wm = 0～1，wn = 0～1，wk = 0

TiledMMA 内部将这个布局与 Atom 自己的线程 ID 布局组合：

using ThrLayoutVMNK = decltype(
    tiled_product(AtomThrID{}, AtomLayoutMNK{}));

本例 AtomThrID 是 32:1。组合后打印为：

ThrLayoutVMNK = (32,2,2,1):(1,32,64,0)

映射方向为：

(lane,wm,wn,wk) → 线程编号

t = lane + 32*wm + 64*wn

这里 VMNK 中的 V 表示 Atom 内部的线程维，也就是源码注释里的 ThrV。它与前面每 lane 的值编号 v 不是同一件事；“值”在 thrfrg_C 的注释中称为 FrgV。

第三个参数保持默认时，各方向的基本 tile 大小就是 Atom 尺寸乘上该方向的线程组数：

TileM = 16 × 2 = 32
TileN =  8 × 2 = 16
TileK = 16 × 1 = 16

总线程数为 32×2×2×1=128。注意，size(mma) 对应线程数量，tile_shape(mma) 才是 (32,16,16)。

get_slice 从线程编号反查位置

实际使用时，线程调用：

auto thr_mma = mma.get_slice(threadIdx.x);

源码先做：

auto thr_vmnk = thr_layout_vmnk_.get_flat_coord(thr_idx);

它根据线程布局反查当前线程在 Atom 内部及 M、N、K 方向的位置。对于本例：

lane = t%32
wm   = (t/32)%2
wn   = t/64
wk   = 0

于是前面的 C 坐标公式增加了 warp 所在块的偏移：

m = lane/4 + 8*(v/2) + 16*wm
n = 2*(lane%4) + v%2 + 8*wn

四个 warp 的 lane 0 各自持有：

线程	(wm,wn)	C 的四个值
0	(0,0)	(0,0), (0,1), (8,0), (8,1)
32	(1,0)	(16,0), (16,1), (24,0), (24,1)
64	(0,1)	(0,8), (0,9), (8,8), (8,9)
96	(1,1)	(16,8), (16,9), (24,8), (24,9)

这张表就对应开头那张四块输出的图。

换 stride，会换掉 warp 的排列

若将第二个参数改为：

Layout<Shape<_2,_2,_1>, Stride<_2,_1,_0>>{}

线程组的编号变成：

warp_id = 2*wm + wn
ThrLayoutVMNK = (32,2,2,1):(1,64,32,0)

此时 warp 1 在右上，warp 2 在左下。线程总数和 tile 尺寸都相同，改变的是哪个 warp 负责哪块矩阵。get_slice 会依据新的布局反查坐标；不能继续套用前一组 wm/wn 的拆解公式。

本文使用的线程布局都连续覆盖一段线程编号。对更一般的嵌套或带空洞布局，需要检查实际 ID 映射，再安排参与的线程。

partition_C 怎样把分工应用到一个 Tensor

现在准备一个按行连续存储的 32×16 输出 Tensor：

auto C = make_tensor(
    make_gmem_ptr(c_ptr),
    make_layout(make_shape(_32{}, _16{}),
                make_stride(_16{}, _1{})));

auto thr_mma = mma.get_slice(threadIdx.x);
auto tC = thr_mma.partition_C(C);

partition_C 保留 C 的存储，根据 thrfrg_C(C.layout()) 建立访问视图，再固定当前线程。下面按 thrfrg_C 的顺序，看看这些布局运算分别处理哪一层。

先按第三个参数划分 M、N

源码：

auto t_tile = make_tile(permutation_mnk<0>(),
                        permutation_mnk<1>());
auto t_tensor = logical_divide(ctensor, t_tile);

第三个参数为 _ 的方向，permutation_mnk<I>() 返回：

size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk())

所以本例的 t_tile 相当于 (32,16)。对恰好一个 tile 的输入：

t_tensor 的坐标结构：((32,1),(16,1))
                         M      N

每个方向先分为 tile 内坐标与 tile 外剩余坐标。本例的剩余大小都是 1。第三个参数若传入非平凡 Layout，元素的访问顺序也会在这一步改变，后面会单独举例。

再切出单条指令的输出块

源码按 Atom 的 M、N 尺寸进行 zipped_divide：

auto c_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})),
                        make_layout(size<1>(AtomShape_MNK{})));
auto c_tensor = zipped_divide(t_tensor, c_tile);

概念上的坐标结构变成：

((AtomM,AtomN),(剩余M,剩余N))
      (16,8)       (2,2)

本例这一步会消去上一层多余的 size-1 结构，实际打印的 layout 为：

((16,8),(2,2)):((16,1),(256,8))

第一个组合 mode 表示一条 16×8 MMA 内部的位置，第二个表示这个 Atom 块放在哪里。沿 M 跨一个 Atom 块要走 16 行，因此 stride 是 16×16=256；沿 N 跨一个 Atom 块要走 8 列，因此 stride 是 8。

用 Atom 的 TV 布局替换块内坐标

接着：

auto tv_tensor = c_tensor.compose(AtomLayoutC_TV{}, _);

AtomLayoutC_TV 就是前面已经展开的 CLayout。它把 (lane,value) 映射到 Atom 块内的 (m,n) 编号。与 c_tensor 复合以后，得到：

((Atom内线程,Atom内值),(剩余M,剩余N)) → C 的存储偏移

第二个参数 _ 表示保留 Atom 块之外的坐标。

例如 Atom 内 lane 0 的值 2，对应块内 (8,0)。C 每行跨度为 16，所以这里的存储偏移是 8×16=128。前面 CLayout 单独返回的逻辑编号却是 8：复合到真实 Tensor 以后，才换算成实际存储偏移。

把剩余块分给 M、N 方向的线程组

最后一次 zipped_divide 使用线程组数 (2,2)：

auto thr_tile = make_tile(_,
    make_tile(make_layout(size<1>(thr_layout_vmnk_)),
              make_layout(size<2>(thr_layout_vmnk_))));
auto thr_tensor = zipped_divide(tv_tensor, thr_tile);

结果的输入结构可读成：

((lane,(wm,wn)),(Atom内值,(RestM,RestN)))

前一半都用于选择线程：Atom 内部选 lane，外部选 M、N 方向的线程组。后一半是该线程需要持有的值和重复计算的部分。

ThrMMA::partition_C 从 get_slice 保存的坐标中取出 (lane,wm,wn)，固定前一半，再将值和剩余方向展开为顶层 mode。对于当前 32×16 输入，返回：

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

第一维大小是 2×2=4，后两维各为 1。线程 0 的四个存储偏移为 0、1、128、129，对应 C(0,0)、C(0,1)、C(8,0)、C(8,1)。

线程 32 的返回 layout 相同，但视图起点推进到 C(16,0)；线程 64 的起点则是 C(0,8)。只比较 layout，不看数据起点，不能判断两个线程是否访问同一份数据。

A、B 怎样与四块输出配合

A 的分区沿 (M,K) 进行，B 沿 (N,K) 进行。源码与 C 的步骤相同，只是换成对应方向和对应 Atom TV 布局。线程坐标也各取所需的两项：

partition_A 选择 (lane,wm,wk)
partition_B 选择 (lane,wn,wk)
partition_C 选择 (lane,wm,wn)

这三个选择解释了数据为什么会重复。

在 2×2×1 的配置中，warp 0 和 warp 2 的 M、K 坐标一样。它们计算不同的输出列，但需要同一块 A。warp 0 和 warp 1 的 N、K 坐标一样，则需要同一块 B。

               A 的 M 块         B 的 N 块       输出块
warp 0           上半              左半           左上
warp 1           下半              左半           左下
warp 2           上半              右半           右上
warp 3           下半              右半           右下

因此对 A(32,16)、B(16,16)、C(32,16)：

操作数	每 lane 的值数	全部线程持有的值数	不同矩阵元素数
A	8	128×8=1024	32×16=512
B	4	128×4=512	16×16=256
C	4	128×4=512	32×16=512

A、B 每个逻辑元素在不同 warp 中各出现两次，C 则覆盖一次。这里的重复描述的是各线程 fragment 的需求；是否能在某一层缓存或共享加载结果，取决于内核的数据搬运实现。布局构造本身没有生成一次跨 warp 广播。

输入扩大以后，后两个 mode 怎样工作

保持 mma 不变，将输入扩大成：

A：(64,32)
B：(32,32)
C：(64,32)

也就是 M=64、N=32、K=32。

基本 tile 仍为 (32,16,16)。同样 128 个线程继续分区，会得到：

partition_A：((2,2,2),2,2)   // (Atom值, M重复, K重复)
partition_B：((2,2),  2,2)   // (Atom值, N重复, K重复)
partition_C：((2,2),  2,2)   // (Atom值, M重复, N重复)

第一维是 Atom 自己规定的值分组。后面两个大小为 2 的 mode，分别来自各方向多出来的工作。

以 C 为例，将重复坐标记为 rm、rn，则：

m = lane/4 + 8*(v/2) + 16*wm + 32*rm
n = 2*(lane%4) + v%2 + 8*wn + 16*rn

v = 0～3，rm = 0～1，rn = 0～1

线程 0 的四组输出是：

(rm,rn)	四个 C 坐标
(0,0)	(0,0), (0,1), (8,0), (8,1)
(1,0)	(32,0), (32,1), (40,0), (40,1)
(0,1)	(0,16), (0,17), (8,16), (8,17)
(1,1)	(32,16), (32,17), (40,16), (40,17)

对于按行连续存储的 C，完整分区 layout 为：

((2,2),2,2):((1,256),1024,16)

其中 Atom 内跨 8 行的 stride 是 8×32=256，M 方向重复一次跨 32 行，即 32×32=1024；N 方向重复一次跨 16 列，即 stride 16。

同样，A 的坐标公式增加 32*rm 和 16*rk，B 增加 16*rn 和 16*rk。每 lane 此时需要 32 个 A 值、16 个 B 值和 16 个累加器值。

这些重复工作仍由原来的四个 warp 完成。输入 Tensor 变大，不会让 partition_A/B/C 创建更多线程。

partition 与 partition_fragment 分别返回什么

典型调用会接着创建 fragment：

auto tA = thr_mma.partition_A(A);
auto tB = thr_mma.partition_B(B);
auto tC = thr_mma.partition_C(C);

auto rA = thr_mma.partition_fragment_A(A);
auto rB = thr_mma.partition_fragment_B(B);
auto rC = thr_mma.partition_fragment_C(C);

tA/tB/tC 是原 Tensor 的线程访问视图。对于本文的 SM80 Atom，rA/rB/rC 则提供该线程用于 MMA 的值存储，设备代码通常将其保存在寄存器中。它们刚创建时还没有加载 A、B，也没有获得可用的初始累加值。

对于前面的较大输入，程序打印：

fragment_A = ((2,2,2),2,2):((1,2,4),8,16)
fragment_B = ((2,2),  2,2):((1,2),  4,8)
fragment_C = ((2,2),  2,2):((1,2),  4,8)

它们保留分区的逻辑 shape，但存储布局已经不同。以 C 为例，原 Tensor 视图的 stride 为 ((1,256),1024,16)，fragment 的 stride 则为 ((1,2),4,8)：需要累加的 16 个值在本线程的存储中紧凑排列。

因此，不能拿 tC 的真实矩阵偏移去索引 rC。两者通过相同的逻辑坐标对应元素，底层存储各有自己的布局。

这里的结论限定于本文采用的寄存器 MMA。MMA_Atom::make_fragment_A/B 也支持由 trait 指定的其他 fragment 类型，有些会创建特定视图，不能把所有架构的 partition_fragment 都概括为“分配一组普通寄存器”。

第三个参数：tile 范围与范围内的重排

构造函数将第三个参数称为 permutations。它的每个方向可以是 _、静态整数或静态 Layout。默认值是 (_,_,_)。

源码中的两个函数把它的行为直接联系起来：

// 省略 static_assert
auto permutation_mnk() const {
  auto perm = get<I>(PermutationMNK{});
  return conditional_return(
      is_underscore<decltype(perm)>{},
      size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk()),
      perm);
}

auto tile_size_mnk() const {
  return size(permutation_mnk<I>());
}

_ 使用 Atom 尺寸与线程组数的乘积；显式配置则直接交给前面 logical_divide 的第一步。tile_size 取该配置的 size。

只传更大的整数范围

例如：

auto mma_large_tile = make_tiled_mma(
    Op{},
    Layout<Shape<_2,_2,_1>>{},
    Tile<_64,_32,_16>{});

这里没有指定非平凡的坐标置换。M、N 方向的 tile 范围扩大为 64 和 32，线程布局仍是：

ThrLayoutVMNK = (32,2,2,1):(1,32,64,0)
tile_shape   = (64,32,16)

还是四个 warp、128 个线程。一份 64×32 输出 tile 内，每个 warp 需要处理 2×2 个 Atom 输出块，每 lane 有 16 个 C 值。对相配的 A、B、C，分区 shape 为：

A：((2,2,2),2,1)
B：((2,2),  2,1)
C：((2,2),  2,2)

将同样大小的 Tensor 交给前面默认第三参数的 mma，也会得到这些线程工作量。区别是 tile_shape 以及 get_layoutA/B/C_TV() 所描述的参考范围：前者仍报告基本 (32,16,16)，这里则将额外重复也包含在显式的 (64,32,16) tile 内。

所以第三个参数不等同于“把第二个参数再乘一遍”。它能扩展每线程在一个所声明 tile 内的工作，而不增加参与线程。

用 Layout 改变 M 方向的坐标顺序

再看一个可以手算的重排：

using PermM = Layout<Shape<_8,_4>, Stride<_4,_1>>;

auto mma_perm = make_tiled_mma(
    Op{},
    Layout<Shape<_2,_2,_1>>{},
    Tile<PermM, Underscore, Underscore>{});

PermM 一共覆盖 32 个位置。对输入整数坐标 i=0～31，它按 shape (8,4) 解码，再计算：

p(i) = 4*(i%8) + i/8

几个位置的结果为：

原坐标 i	0	1	7	8	15	16	24	31
重排后 p(i)	0	4	28	1	29	2	3	31

它是对 0～31 的一一重排，tile 大小仍为 (32,16,16)，线程数量仍为 128。可以先用原来的公式求出虚拟 M 坐标，再将它代入 p：

i = lane/4 + 8*(v/2) + 16*wm
m = p(i)
n = 2*(lane%4) + v%2 + 8*wn

于是四个 warp 的 lane 0 输出变为：

线程	原来的 C 坐标	重排后的 C 坐标
0	(0,0), (0,1), (8,0), (8,1)	(0,0), (0,1), (1,0), (1,1)
32	(16,0), (16,1), (24,0), (24,1)	(2,0), (2,1), (3,0), (3,1)
64	(0,8), (0,9), (8,8), (8,9)	(0,8), (0,9), (1,8), (1,9)
96	(16,8), (16,9), (24,8), (24,9)	(2,8), (2,9), (3,8), (3,9)

A 的 M 方向也应用相同的 p，所以所需输入与输出行仍然对应；B 没有 M 方向，保持原来的分工。

此时开头那张“每个 warp 占据一个连续 16×8 矩形”的图已经不适用。warp 0 负责余数为 0、1 的那些行，warp 1 负责余数为 2、3 的那些行，余数按 4 计算。完整输出仍被四个 warp 覆盖一次。

变化的是各个 lane 的 fragment 槽位对应哪些矩阵坐标。底层执行的仍是同一条指令，Tensor 的原始数据也没有因为构造 mma_perm 而被搬动。

本文选的范围能被 Atom 与线程组的分块整除，重排也是完整双射。配置一般 Layout 时，需要检查它是否适合相应的划分，不能将任意 shape、stride 都当成有效 permutation。

K 方向的线程组，与 K 方向的循环

前面较大输入的 K=32，使用的线程组数仍为 (2,2,1)。每个 warp 沿 K 处理两个 16 项分块，并在自己的同一组累加器中继续累加。

如果改的是第二个参数：

auto mma_split_k = make_tiled_mma(
    Op{}, Layout<Shape<_1,_1,_2>>{});

情况就不同了。现在有两个 warp，分别处在 K 方向的两组：

ThrLayoutVMNK = (32,1,1,2):(1,0,0,32)
tile_shape   = (16,8,32)

warp 0：使用 k = 0～15
warp 1：使用 k = 16～31

它们读取不同的 A、B 分块，却对应同一个 16×8 输出范围。前面源码中 partition_C 只选择 (lane,wm,wn)，没有选择 wk；因此线程 0 和线程 32 的 C 视图指向相同的四个逻辑输出位置。

如果两个 warp 分别从零累加，它们得到：

warp 0：P0(m,n) = Σ(k= 0..15) A(m,k)*B(n,k)
warp 1：P1(m,n) = Σ(k=16..31) A(m,k)*B(n,k)

最终结果需要把 P0 和 P1 加起来。make_tiled_mma 和本文使用的线程局部 gemm 不会自动完成跨 warp 归约，不能让两个 warp 都把各自部分和直接覆盖写回同一个 C。若原算式还有初始 C，也要在归约过程中安排它只被加入一次。

因此，单个 warp 遍历两个 K 分块，与两个 warp 分别处理 K 的一半，是两种不同的执行方式。配套程序对这个例子检查的是 A、B 各覆盖一次、C 覆盖两次的坐标关系，不把它当成已经完成归约的矩阵乘。

真正的 MMA 在哪里执行

对于已经准备好共享内存数据的 kernel，本文四 warp 配置的典型使用过程可写成：

// kernel 内：128 个线程参与，sA/sB 已经写好并完成所需同步。
// sA 的逻辑方向为 (M,K)，sB 为 (N,K)，gC 为 (M,N)。
auto mma = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>>{});
auto thr_mma = mma.get_slice(threadIdx.x);

auto rA = thr_mma.partition_fragment_A(sA);
auto rB = thr_mma.partition_fragment_B(sB);
auto rC = thr_mma.partition_fragment_C(gC);

auto copyA = make_tiled_copy_A(
    Copy_Atom<SM75_U32x4_LDSM_N, half_t>{}, mma);
auto copyB = make_tiled_copy_B(
    Copy_Atom<SM75_U32x2_LDSM_N, half_t>{}, mma);
auto thr_copyA = copyA.get_slice(threadIdx.x);
auto thr_copyB = copyB.get_slice(threadIdx.x);

auto tAsA = thr_copyA.partition_S(sA);
auto tBsB = thr_copyB.partition_S(sB);
auto tArA = thr_copyA.retile_D(rA);
auto tBrB = thr_copyB.retile_D(rB);

copy(copyA, tAsA, tArA);
copy(copyB, tBsB, tBrB);
clear(rC);
gemm(mma, rA, rB, rC);
copy(rC, thr_mma.partition_C(gC));

这是一段说明调用关系的代码片段，计算的是从零开始的 A×B。完整内核还需要准备指针、共享内存及同步，并选择满足 ldmatrix 地址、对齐和矩阵方向要求的共享内存布局。A、B 的加载路线可接着看从共享内存到 MMA：读懂 A、B 的 ldmatrix 与 retile_D。

make_tiled_copy_A/B 从 MMA 取得操作数的 TV 分工，retile_D 将已有的 MMA fragment 组织成 Copy 所需的目的视图。加载完成以后，gemm 才消费这些值。

对于本文的寄存器 fragment，gemm 把形状看成：

A：(V_A, M重复, K重复)
B：(V_B, N重复, K重复)
C：(V_C, M重复, N重复)

它沿 K 重复维循环，再遍历 M、N 的输出组合，最终把各 Atom 内部的值交给 MMA_Atom::call。该函数经过 mma_unpack，调用底层 Operation 的 fma，执行 mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16。

对于 M=64、N=32、K=32 的例子，每个 warp 有 2×2×2=8 个 Atom 乘加步骤；每一步由该 warp 的 32 个 lane 共同参与。这是八次 warp 级指令执行，不是 32 条互不相关的标量矩阵乘。

这些形状也给出了复用方式：同一份 rA(_,rm,rk) 可以参与不同 rn 的输出，同一份 rB(_,rn,rk) 可以参与不同 rm 的输出。MMA 的输入和输出通过分区坐标对应起来。

与 make_tiled_copy 放在一起看

前一篇 make_tiled_copy 从用户给定的线程布局和值布局出发，定义一套搬运分工。本篇的 Atom 则已经规定一条指令需要哪些 lane-value，构造函数在这个基础上排列线程组，并组织更多矩阵工作。

要决定的事情	make_tiled_copy	make_tiled_mma
底层操作	Copy Atom	MMA Atom
线程怎样排列	第二个参数给出线程布局	第二个参数给出 Atom 线程组的 M/N/K 布局
每线程有哪些值	值布局结合 Copy Atom 组织	Atom trait 加上矩阵重复与 permutation 决定
第三个参数	本接口中的 val_layout	各方向的 tile 范围或重排
取得原 Tensor 视图	partition_S/D	partition_A/B/C
执行操作	copy	gemm 最终调用 MMA Atom

当数据将进入 MMA 时，直接采用 MMA 的 A/B TV 布局来构造 TiledCopy，便能把加载结果接到指令要求的槽位上。两套配置就在 make_tiled_copy_A/B 处连接。

打印布局，核对一次

配套程序位于 cute_make_tiled_mma_trace.cu。在 CUTLASS 根目录运行：

nvcc -std=c++17 -arch=sm_80 -Iinclude \
  cute_make_tiled_mma_trace.cu -o /tmp/cute_make_tiled_mma_trace
/tmp/cute_make_tiled_mma_trace

也可以用主机编译器运行布局检查：

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

程序使用 identity tensor，让分区结果直接返回矩阵坐标，再独立按本文的 lane、warp 和重复坐标公式核对。它遍历所有参与线程，检查七组配置：单 warp、四 warp、较大输入、交换 warp 顺序、显式 tile 范围、M 方向置换和 K 方向线程分工。

总计检查 20,992 个 A/B/C 坐标，检查每个矩阵元素出现的次数；当输入恰好等于声明的 tile 范围时，还逐项核对 get_layoutA/B/C_TV()。另外会打印 C 分区的中间布局，以及原 Tensor 视图与 fragment 的布局。

四 warp 配置的关键输出为：

ThrLayoutVMNK = (_32,_2,_2,_1):(_1,_32,_64,_0)
tile MNK = (_32,_16,_16)
partition_A shape = ((_2,_2,_2),_1,_1)
partition_B shape = ((_2,_2),_1,_1)
partition_C shape = ((_2,_2),_1,_1)
Verified 2048 coordinates; A/B/C multiplicity = 2/2/1.

较大输入的关键输出为：

partition_A shape = ((_2,_2,_2),_2,_2)
partition_B shape = ((_2,_2),_2,_2)
partition_C shape = ((_2,_2),_2,_2)
fragment_A layout = ((_2,_2,_2),_2,_2):((_1,_2,_4),_8,_16)
fragment_B layout = ((_2,_2),_2,_2):((_1,_2),_4,_8)
fragment_C layout = ((_2,_2),_2,_2):((_1,_2),_4,_8)
Verified 8192 coordinates; A/B/C multiplicity = 2/2/1.

CuTe 打印静态整数时带 _ 前缀，正文推导省略了这个前缀。第一维的嵌套结构保留了 Atom 的值分组，不能只凭外层逗号数量来数值个数。

这些检查在主机上完成，不启动 GPU kernel，也不执行 MMA 数值计算。它们验证的是构造、线程分工、坐标覆盖以及 TV 与 Tensor 分区的对应关系。

本文尺寸均能被相应分块整除。make_tiled_mma 不会自动处理边界、准备输入或选择累加器初始化方式；面对不能整除的矩阵，需要在外层切块与数据准备中安排有效范围和补零等处理。

对照源码

• mma_atom.hpp：make_tiled_mma：Operation 重载、补齐三维配置与 TiledMMA 的构造。
• mma_atom.hpp：TiledMMA：ThrLayoutVMNK、thrfrg_C/A/B 和 get_slice。
• mma_atom.hpp：permutation 与 tile size：默认 _ 与显式配置的解释；同文件还有 get_layoutA/B/C_TV。
• mma_atom.hpp：ThrMMA：三个操作数如何选择线程坐标，以及 partition_fragment_A/B/C。
• mma_traits_sm80.hpp：本文使用的 Atom shape、ThrID 和 A/B/C 的 TV 布局。
• mma_sm80.hpp：底层寄存器参数与 mma.sync 指令。
• gemm.hpp：分区后的寄存器 fragment 怎样沿 K 循环，并继续派发到各个 Atom。
• copy_atom.hpp：make_tiled_copy_A/B 如何采用 MMA 已有的操作数分工。
• 读懂 CuTe 的 make_tiled_copy：对照线程布局、值布局与 Copy Atom 的职责。

本文依据当前工作区的 CUTLASS/CuTe C++ 实现，源码提交为 0b55a2f691d69981583568fd9eb69687b1f0de8a。