ldmatrix 怎样对上 MMA 的寄存器：从目的槽位反推源地址

读 CuTe 的 GEMM 代码时，经常会看到这样的加载过程：

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

copy(tiled_copy, tXsA, tXrA);

MMA 的布局规定 lane 0 的 A 值槽位 5 应当保存 A(0,9)。Copy Atom 的目的布局又把同一个 (lane=0,value=5) 映射成内部编号 129。这两张布局的输出不同，怎样保证最终取到正确的矩阵元素，并写进正确的槽位？

先给这个元素一个具体数值：

shared memory：A(0,9) 中存放数值 37
加载完成后：  lane 0 的 tCrA(5,0,0) 中也应当是 37

MMA 布局确定槽位需要什么元素，Copy 布局确定哪个源位置会送到这个槽位。CuTe 将两套规则组合，计算源地址，并将指令输出绑定到已有 fragment 存储。 真正执行时，ldmatrix 按固定路由把数值 37 搬过去。

下面先从目的槽位反查源位置，再看源码怎样算出实际地址，最后走一遍加载和 MMA 的寄存器参数。全文始终沿用 A(0,9)，源码对应 CUTLASS 提交 0b55a2f691d6。

先固定一个能手算的例子

我们只看一个 warp，使用一条 SM80 MMA 和非转置的 ldmatrix.x4：

auto mma = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{});
auto load_a = Copy_Atom<SM75_U32x4_LDSM_N, half_t>{};
auto tiled_copy = make_tiled_copy_A(load_a, mma);

一个 warp 有 32 个线程，warp 内的线程编号叫 lane，取值为 0…31。half 是 16 位浮点数，两个 half 可以打包到一个 32 位寄存器操作数中。

MMA 的计算尺寸为 M=16、N=8、K=16。本文只看 A：它是一个 16×16 矩阵，按 (m,k) 索引；每个 lane 提供 8 个 half，共同组成这 256 个 A 元素。

假设数据已经写入 shared memory，并且对参与加载的线程可见。基址满足 16 字节对齐，K 方向连续，每行 16 个 half：

// smem_a 指向已经准备好的 shared memory 数据。
auto sA = make_tensor(
    make_smem_ptr(smem_a),
    Layout<Shape<_16,_16>, Stride<_16,_1>>{});

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

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

本文用 shape:stride 表示布局。Shape 给出维度大小，stride 给出该维度走一步增加的元素偏移，因此：

sA.layout() = (16,16):(16,1)
元素偏移     = 16*m+k
字节偏移     = 2*(16*m+k)

后面用 0x1000 作为示例 shared memory 基址，便于手算地址。所有元素编号都从 0 开始，除法取整数；暂不展开 swizzle 和流水线。

第一步：明确目的槽位需要哪个 A 元素

在 mma_traits_sm80.hpp 中，这条 MMA 的 A 布局为：

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

TV 是 thread、value 的缩写。输入的 T 是线程编号，本例就是 lane；V 是该线程内的逻辑值编号，以 half 为单位。输出将 A tile 的矩阵坐标按 q=m+16*k 编码。

取 lane 0、值编号 5：

lane 0 在 shape (4,8) 中展开为 (0,0)
value 5 在 shape (2,2,2) 中展开为 (1,0,1)

q = 32*0 + 1*0 + 16*1 + 8*0 + 128*1
  = 144

m = q%16 = 0
k = q/16 = 9

因此，lane 0 的 A 值编号 5 应当保存 A(0,9)。partition_fragment_A(sA) 根据 MMA 的分工准备线程局部存储，本例中值编号 5 就对应 half 槽位 5。创建 fragment 时只准备槽位，真正读取数据要等到 copy。

现在可以把目标写成一个具体赋值：

// 只表达这个槽位的加载结果，实际搬运由 warp 协作的 ldmatrix 完成。
// 以下 tCrA 属于 lane 0。
tCrA(5,0,0) = sA(0,9);

这里同时出现了三个不同的数：

数值	含义
q=144	矩阵坐标 (0,9) 的逻辑编码
offset=9	sA 中该元素的 half 存储偏移：16*0+9
37	该位置实际存放的数据

在示例基址下，待取数据的地址是 0x1000+2*9=0x1012。接下来要确定：哪一个源 lane 应当提供包含这个元素的行地址？

第二步：从目的槽位反查 Copy 的源位置

SM75_U32x4_LDSM_N 描述的是一条 ldmatrix.x4。它一次加载四个 8×8 的 16 位元素矩阵，每个目的 lane 得到四个 32 位输出，共 8 个 half。

Copy_Atom<...,half_t> 将 traits 中以 bit 为单位的布局转换成 half 单位。当前 Atom 的布局为：

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

源布局的值编号 p=0…7 表示一个源行内的 half；目的值编号 v=0…7 的 (2,4) 表示“每个输出两个 half，共四个输出”。两张布局输出同一套 Copy 内部元素编号 e：

S(tS,p) = 8*tS+p
D(tD,v) = 2*tD + v%2 + 64*(v/2)

先从目的 lane 0、值 5 出发：

D(0,5) = 2*0 + 5%2 + 64*(5/2)
       = 129

再反查源位置。令 S(tS,p)=129，其中 p 取 0…7：

8*tS+p = 129

tS = 16
p  = 1

这就是下面这个转换的含义：

right_inverse(S).compose(D)

Copy 目的 (0,5)
 → 内部元素 e=129
 → Copy 源 (16,1)

要填入 lane 0 的 half 槽位 5，必须由源 lane 16 所提供行内编号 1 的元素送来。 这一结论来自 Copy 指令的固定路由。e=129 在这里用于匹配源和目的位置。

接下来还需要把 A 的坐标赋给这个源位置。不能直接用 129 作为 sA 的元素下标。

第三步：将 Copy 路由接到 MMA 的矩阵分工

make_tiled_copy_A 将 MMA 的 A 布局作为参数传入：

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

TiledCopy 把这个布局保存为 TiledLayout_TV。该接口的语义是“reference TV 到矩阵位置”，因此可以记为：

F = TiledLayout_TV = mma.get_layoutA_TV()

F(reference_thread, reference_value) → A tile 的矩阵编号 q

当前 R=D，Copy 目的 TV 转成 reference TV 后，线程和值编号保持不变。所以 Copy 目的 (0,5) 的最终矩阵含义是：

目的 (0,5)
 → D：内部编号 129
 → right_inverse(R)：reference (0,5)
 → F：矩阵编号 144，即 A(0,9)

这就是前面两套编号的连接点。D 先将目的 TV 转成 Copy 内部编号，再通过 right_inverse(R) 回到 reference TV，最后才交给 MMA 的矩阵映射 F。

将三个布局并列，可以看到“输出不同”发生在哪一层：

布局	输入	输出
MMA 的 F	(0,5)	矩阵编号 q=144
Copy Atom 的 D	(0,5)	内部元素编号 e=129
Copy 最终目的映射 F ∘ right_inverse(R) ∘ D	(0,5)	矩阵编号 q=144

Copy Atom 的内部编号与矩阵编号含义不同；组合后的 Copy 目的映射与 MMA 的矩阵分工则一致。

从这个对应关系出发，第二步反查到的源位置 (16,1) 就必须读取 A(0,9)。下面看 partition_S 如何落实这个要求。

第四步：partition_S 将源位置变成实际地址

前面为了便于理解，从目的位置反查源位置。源码生成源视图时，则从源 TV 正向计算矩阵位置。tidfrg_S 传入的转换为：

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

tile2thrfrg 再把这个转换接到 TiledLayout_TV：

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

auto trg_layout_TV = atom_layout_TV.compose(ref2trg, _);

本例整个 reference TV 和一个 Copy Atom 都是 32 个线程、每线程 8 个值，切分后剩余组大小为 1。省略这些大小为 1 的 mode，源到矩阵的映射就是：

F ∘ right_inverse(R) ∘ S

将源 (16,1) 代入：

S(16,1) = 129

right_inverse(R)(129)
  → reference TV 的展平编号 tv=160
  → (thread,value)=(0,5)，因为 160=0+32*5

F(0,5) = 144
  → (m,k)=(0,9)

源位置经过组合后，得到的也是 A(0,9)。更一般地，Copy 路由中相互对应的源、目的位置满足 S(src_tv)=D(dst_tv)=e，所以两端接入同一个 right_inverse(R) 和 F 后，会得到相同的矩阵坐标。

但现在仍只是求出了坐标，还要让源 tensor 真正指向该元素。tile2thrfrg 整理线程和值的 mode 后执行：

auto tv_tensor = tensor.compose(thrval2mn, _);

这一步将刚才的逻辑映射接到 sA 原来的布局上。对 (16,1)，结果为：

源 TV=(16,1)
 → 矩阵坐标 (0,9)
 → sA.layout()(0,9)=9
 → sA.data() 中的第 9 个 half
 → 地址 0x1012，里面存放数值 37

partition_S 保留原 tensor 的存储起点，并用组合后的布局构造源视图。因此，tXsA 的元素访问最终能够引用这个真实的数据位置。源码中的 Tiler_MN 是通用名字；在当前 A tensor 上，它的方向实际是 M、K。

指令需要的是行首地址

ldmatrix 的每个源 lane 提供一行的起始地址。因此还要计算 lane 16 的源值 0：

源 (16,0)
 → S：128
 → reference (0,4)
 → F：128，即 A(0,8)
 → sA 元素偏移 8
 → 行首地址 0x1010

将这行的前两个值放在一起：

源位置	对应 A 元素	half 偏移	示例地址
(16,0)	A(0,8)	8	0x1010
(16,1)	A(0,9)	9	0x1012

继续计算源值 2…7，得到 A(0,10)…A(0,15)。这些元素在当前存储布局中连续，因此 lane 16 最终向指令提供 &A(0,8)，行内编号 1 就是待搬运的 A(0,9)。

到这里，读取哪一份数据已经确定。接下来要让指令输出引用正确的 fragment 槽位。

第五步：将 Copy 输出绑定到 tCrA 的槽位

retile_D(tCrA) 使用 fragment 原来的数据起点，并从原布局推导 Copy 的目的访问方式：

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

本例打印的布局为：

tCrA：((2,2,2),1,1):((1,2,4),0,0)
tXrA：((8,1),  1,1):((1,0),  0,0)

第 0 个 mode 分别表示“一条 MMA 当前 lane 的 8 个 A 值”和“一次 Copy 的 8 个 half，共 1 组”。后两个 mode 是 M、K 方向的重复，本例都为 1。

对值编号 5，两种布局计算的存储偏移为：

tCrA：5 展开为 (1,0,1)，偏移 = 1+2*0+4*1 = 5
tXrA：5 展开为 (5,0)，  偏移 = 5+0*0 = 5

两者又保留同一个存储起点，因此：

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

随后，copy_unpack 将传入的目的视图按指令接口重解释：

Tensor rD = recast<RegTypeDst>(dst);  // 本例 RegTypeDst=uint32_t

每两个相邻 half 成为一个 32 位输出。detail::explode 按顺序传参，展开后的调用形式为：

CopyOp::copy(rS(0), rD(0), rD(1), rD(2), rD(3));

因此，当前 lane 的输出引用分别绑定为：

指令输出	fragment 的 half 槽位	后续 MMA 的 A 操作数
dst0，即 rD(0)	0、1	a0
dst1，即 rD(1)	2、3	a1
dst2，即 rD(2)	4、5	a2
dst3，即 rD(3)	6、7	a3

lane 0 的 dst2 高 16 位已经绑定到 tCrA 的 half 槽位 5。 这是输出参数与 fragment 存储的对应关系，加载指令执行后会直接写入该位置。

第六步：执行 ldmatrix，搬运数值 37

源行地址和目的输出都准备好后，warp 共同执行：

ldmatrix.sync.aligned.x4.m8n8.shared.b16
    {dst0,dst1,dst2,dst3}, [addr];

这条指令一次处理四个 8×8 矩阵，编号为 b=0…3。各组 lane 提供的行地址，经前面的布局推导，指向下面这些 A 区域：

矩阵编号 b	提供行地址的 lane	对应 A 的范围
0	0…7	A(0…7,0…7)
1	8…15	A(8…15,0…7)
2	16…23	A(0…7,8…15)
3	24…31	A(8…15,8…15)

对非转置的 m8n8.x4，矩阵内 row、col 都取 0…7，硬件分发规则是：

提供这一行地址的 lane = 8*b+row
接收该元素的 lane     = 4*row+col/2
目的 32 位输出编号    = b
输出中的 half 编号    = col%2       // 0 为低位，1 为高位

目标元素 A(0,9) 在第 b=2 个矩阵的 row=0、col=1 处。因此：

lane 16 提供行首地址 0x1010
  → 行内编号 1，地址 0x1012，读到数值 37
  → 接收 lane = 4*0+1/2 = 0
  → 输出编号 2，即 dst2
  → half 编号 1，即高 16 位
  → 直接写入 lane 0 的 tCrA 槽位 5

整个 warp 的线程各自提供地址和输出参数，由硬件完成跨 lane 的分发。源 lane 16 提供这一行的地址，目的 lane 0 接收其中相应的值。

copy_sm75.hpp 的封装将源参数的 shared memory 地址传入指令，并将四个输出引用绑定到 PTX 操作数：

uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src);

asm volatile(
  "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
  : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3)
  : "r"(smem_int_ptr));

对源 lane 16，传入的是行首 0x1010，行内元素 1 由硬件读取。指令参数中没有 129、144，也没有 (m,k)；这些逻辑编号已经用于生成地址和访问方式。

后续 gemm 将 tCrA 的片段交给 MMA Atom，mma_unpack 同样通过 recast<uint32_t> 将相邻 half 组成操作数，并按顺序传入 a0…a3。槽位 5 因而成为 a2 的高 16 位，MMA 按硬件约定将该位置作为 A(0,9) 使用。

recast 保留 half 的原始位模式。此时装在那个位置里的数据是数值 37，Copy 输出和 MMA 输入引用同一份 fragment 存储。

把这次加载对应到两端的数据访问

对于 lane 0 的目标槽位，这次加载的语义可以写成：

q = F(0,5);             // 144
m = q % 16;            // 0
k = q / 16;            // 9

tCrA(5,0,0) = sA(m,k); // 将数值 37 写入槽位 5

这是对结果的描述。实现时，源地址由 partition_S 根据 Copy 路由计算，目的存储由 retile_D 和 copy_unpack 绑定，warp 的 ldmatrix 完成搬运。

可以把本例的两端访问并列起来：

读：sA.data()   + sA.layout()(0,9)    → shared memory 中的 half 偏移 9
写：tCrA.data() + tCrA.layout()(5,0,0) → 当前 lane fragment 的 half 槽位 5

MMA 布局提供的 (lane,value) → 矩阵坐标 关系，正是连接这两个访问位置的依据。它本身不执行数据搬运；数据读取和输出写入发生在加载指令执行时。

本文中 R=D，两边都是 32 个 lane、每 lane 8 个 half，线程编号对应一致，所以目的 TV 可以直接按 MMA 的值编号解释。更大的 tile 会按 Copy Atom 的大小切分 TV，再在各组内应用相同的转换；逐槽位相等的具体形式也要随配置推导。

打印布局，核对一次

配套程序：cute_ldmatrix_mma_coordinate_trace.cu。在 CUTLASS 根目录编译运行：

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

/tmp/cute_ldmatrix_mma_coordinate_trace

CUDA 头文件路径按本机安装位置调整。程序在主机上使用布局和带 shared memory 标签的数组视图，模拟上述路由，不执行 GPU 的 ldmatrix 或 MMA。示例地址 0x1000 仅用于打印，不会被解引用。

程序先用 A(m,k)=16*m+k 填入值标签，再将 A(0,9) 改为 37。这样可以直接看到，搬运的数据与矩阵编号、内存偏移是不同的量。

首先打印目的需求和反查结果：

MMA target (0,5): q=144; Copy final destination q=144
Copy route: destination (0,5) -> e=129 -> source (16,1)

再核对源地址的正向推导：

source TV | e | reference TV | tv | q | A(m,k) | half offset | example byte address
(16,0) | 128 | (0,4) | 128 | 128 | A(0,8) | 8 | 0x1010
(16,1) | 129 | (0,5) | 160 | 144 | A(0,9) | 9 | 0x1012

最后，lane 0 的 fragment 应当得到：

dst0 / a0: low=0, high=1
dst1 / a1: low=128, high=129
dst2 / a2: low=8, high=37
dst3 / a3: low=136, high=137

程序枚举全部源行和行内元素，用独立写出的 m8n8.x4 分发公式对照 Copy traits，并核对源矩阵坐标、目的 MMA 坐标、fragment 存储别名和操作数中的 half 位模式。最后应输出：

PASS: 256 hardware routes, 256 fragment slots, 128 packed operands verified (host only).

对照源码

下面的链接固定到本文使用的提交，便于逐行核对：

• MMA 的 ALayout 与 BLayout：矩阵坐标与线程、值编号的对应关系。
• MMA 对 A 的分区 与 get_layoutA_TV：将 Atom 分工接入 tensor 和 TiledMMA。
• Copy traits 对内部编号的约定：Src、Dst、Ref 的输出用于构造 TV 转换。
• ldmatrix.x4 的 Src、Dst、Ref：一条加载指令内的分发关系。
• Copy_Atom 的 bit 到 value 转换：从 bit 单位转换为 half 单位。
• partition_S 使用的转换与 tile2thrfrg：源 TV 如何接到 reference TV 和输入 tensor。
• retile_D 与 make_tiled_copy_A：已有 fragment 的目的视图，以及 MMA 分工的接入位置。
• copy_unpack 与 ldmatrix.x4 指令封装：源地址和输出寄存器怎样传入 PTX。
• mma_unpack 与 MMA 指令封装：fragment 槽位怎样成为 MMA 的 A 操作数。