跳到主要内容

从 partition_S 到 copy_unpack：ldmatrix 怎样取得行地址、写入 MMA 寄存器

读 CuTe 的共享内存加载代码时，经常会看到这一组操作：

auto load_a = Copy_Atom<SM75_U32x2_LDSM_N, half_t>{};
auto tiled_copy = make_tiled_copy_A(load_a, mma);
auto thr_copy = tiled_copy.get_slice(threadIdx.x);

auto tXsA = thr_copy.partition_S(sA);
auto tXrA = thr_copy.retile_D(tCrA);
copy(tiled_copy, tXsA, tXrA);

tXsA 是当前线程的源视图，tXrA 是当前线程的目的视图。但打印它们对应的矩阵元素，会发现同一 lane 的两端并不逐项相同。这次加载怎样保证正确？源视图里又没有一个叫“行首”的字段，指令如何取得每一行的地址？

先固定一个观察目标：假设共享内存中 A(8,0)=23。第一次加载时，它经过的路线是：

lane 8 提供 &A(8,0)
        ↓ warp 协作的 ldmatrix.x2
lane 0 的第二个 32 位输出，低 half
        ↓ 输出引用绑定
lane 0 的 tCrA 槽位 2 = 23

partition_S 计算各线程提供的源地址；copy_unpack 将源、目的 Tensor 展开成指令参数；ldmatrix 按硬件规则完成 warp 内的加载与分发。 本文就沿着这条路线，逐步接上源码。

先固定矩阵和线程分工​

本文使用 SM80 的 16×8×16 MMA，沿 M、N 方向各安排两个 warp：

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

一个 warp 有 32 个线程，warp 内编号叫 lane，范围是 0…31。这里共有 128 个线程，MMA tile 的 (M,N,K) 大小是 (32,16,16)。A 按 (M,K) 索引，因此共享内存中的 A tile 为 32×16：

__shared__ __align__(16) half_t smem_a[32*16];

auto sA = make_tensor(make_smem_ptr(smem_a),
    Layout<Shape<_32,_16>, Stride<_16,_1>>{});

(32,16):(16,1) 表示 M、K 的大小分别为 32、16，沿 M 前进一步增加 16 个 half，沿 K 前进一步增加 1 个 half。其物理地址为：

half 偏移 = 16*m+k
字节偏移  = 2*(16*m+k)

假设数据已经写入共享内存，并且对参与加载的线程可见。为当前线程创建 MMA fragment 和 Copy 视图：

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

auto load_a = Copy_Atom<SM75_U32x2_LDSM_N, half_t>{};
auto tiled_copy = make_tiled_copy_A(load_a, mma);
auto thr_copy = tiled_copy.get_slice(threadIdx.x);
auto tXsA = thr_copy.partition_S(sA);
auto tXrA = thr_copy.retile_D(tCrA);

创建 tCrA 只是准备当前线程的寄存器存储，数据在后面的 copy 中写入。这里每个 lane 的 A fragment 包含 8 个 half。

为了先看清 warp 内的路线，下面主要观察 warp 0，即线程 0…31。其他 warp 的地址变化会在源地址公式之后接上。本文不引入 swizzle 和流水线。

warp 协作，为什么每个 lane 都要执行 copy​

SM75_U32x2_LDSM_N 封装的指令是：

ldmatrix.sync.aligned.x2.m8n8.shared.b16 {d0,d1}, [addr];

.m8n8 表示每个矩阵块有 8 行、8 列，.b16 表示每个元素占 16 位，.x2 表示一次加载两个这样的矩阵块。两个块共有 128 个 half，由整个 warp 接收。

这条指令的 addr、d0、d1 都是每个 lane 自己的操作数。各 lane 执行同一条指令，硬件收集地址，并将结果分发到各 lane 的寄存器中。

NVIDIA 的 PTX 官方文档：ldmatrix 明确规定：

The mandatory .aligned qualifier indicates that all threads in the warp must execute the same ldmatrix instruction.

所以，每个线程取得自己的 thr_copy 并执行 copy，正是完成 warp 协作的方式。对 .x2，两端分工为：

职责对应 lane
提供第一个 8×8 块的 8 个行首地址0…7
提供第二个 8×8 块的 8 个行首地址8…15
接收结果，每 lane 两个 32 位寄存器全部 32 个 lane

官方对这些源地址的表述是：

Each address corresponds to the start of a matrix row.

对每个 8×8 块，一行 8 个 half 被分成四对，分给四个连续的目的 lane。官方的 fragment 分布图 给出了这个规则。

以一个块 X 的第 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 的 d0，第二个块的数据写入各 lane 的 d1。因此，lane 8 可以提供第二个块第 0 行的地址，而 lane 0 接收这行中的前两个 half。

这就解释了同一 lane 的源、目的视图为什么可以不同：当前 lane 提供的地址与它自己接收的片段，是指令的两套分工。

partition_S 怎样让每个线程取得一行​

先把源、reference 和矩阵坐标接起来​

make_tiled_copy_A 用 MMA 的 A 布局作为 Copy 的 reference 分配：

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

本例 Tiler_MN=(32,16)，两个方向实际是 M、K。mma.get_layoutA_TV() 规定每个线程的每个 A 槽位需要哪个矩阵元素。

接下来，partition_S 通过 tidfrg_S 生成源端的线程和值布局，再选出当前线程：

auto thr_tensor = make_tensor(
    stensor.data(),
    TiledCopy::tidfrg_S(stensor.layout()));

return thr_tensor(thr_idx_, _, repeat<rank_v<STensor>>(_));

其中 tidfrg_S 的关键转换是：

right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{})

令 S=AtomLayoutSrc、D=AtomLayoutDst、R=AtomLayoutRef。在 .x2 中 R=D。将源码中的 bit 布局换算成 half 单位后，得到：

S = ((16,2),8):((8,0),1)
D = (32,(2,2)):(2,(1,64))
R = D

源端每个 lane 对应 8 个 half；目的端每个 lane 对应 4 个 half。令源 lane 为 l、行内 half 编号为 p=0…7，目的 lane 为 td、本次加载的目的 half 编号为 a=0…3，两个 Atom 布局是：

S(l,p)  = 8*(l%16)+p
D(td,a) = 2*td+(a%2)+64*(a/2)

它们的输出是同一套 Copy 内部元素编号 e。这些编号用来连接源、目的位置，不是共享内存地址，也不是 A 的矩阵编号。

例如，源 lane 8 的第 0 个 half：

源 Atom 位置 (8,0)
    ↓ S
e=64
    ↓ R⁻¹
reference Atom 位置 (0,2)

因为 D(0,2)=64，这表示硬件会把这个源位置的元素送到目的 lane 0、本次 half 槽位 2。

第一次加载时，reference 值 2 就是完整 MMA fragment 的值 2。MMA 规定 lane 0 的这个槽位需要 A(8,0)，因此 partition_S 会为源 lane 8 选择 A(8,0)。

一般地，将加载组记为 g=0,1。tile2thrfrg 先把 reference TV 按 (32,4) 分成 Atom，再接入上面的源到 reference 转换。对 warp 0，最终关系为：

源位置 (l,p,g)
    ↓ S，再通过 R⁻¹ 找到 reference Atom TV
reference 位置 (td,a)
    ↓ 加入加载组
完整 reference TV (td,a+4*g)
    ↓ MMA 的 A 分配
A(l%16, 8*g+p)

因此，源坐标不是由 copy_unpack 临时决定的；它已经在 partition_S 中，按指令路线和 MMA 的目的需求计算好了。

相同的布局，不同的线程基址​

本例实际打印 tXsA，会得到：

tXsA = ((8,2),1,1):((1,8),0,0)

第一维的 (8,2) 表示每次 Atom 使用 8 个 half，共两组。行内 stride 为 1，组间 stride 为 8。后两个大小为 1 的 mode 表示当前没有额外的 M、K 重复。

不同 lane 的这张布局相同，但线程切片后的存储起点不同。warp 0 中：

lane 0：tXsA.data() 指向 A(0,0)
lane 1：tXsA.data() 指向 A(1,0)
lane 8：tXsA.data() 指向 A(8,0)

因此访问关系是：

// p=0…7，g=0,1；两种访问指向同一位置。
tXsA(make_coord(p,g),0,0);
sA(lane%16,8*g+p);

每组的第 0 个 half，就是该子矩阵的行首：

lane，r=0…7第一次加载，g=0第二次加载，g=1
r&A(r,0)&A(r,8)
8+r&A(8+r,0)&A(8+r,8)
16+r重复 &A(r,0)重复 &A(r,8)
24+r重复 &A(8+r,0)重复 &A(8+r,8)

后 16 个 lane 的重复来自 S 在线程子 mode 上的零 stride。官方文档还特别说明，在 sm_75 上所有线程都需要有效地址，.x1/.x2 可以把低编号线程的地址复制到高编号线程。这里的源布局也表达了这种重复。

这里的“行首”是 8×8 子矩阵的行首。第二组从 K=8 开始，因此它使用 &A(m,8)，不是整张 A 行的 &A(m,0)。

其他三个 warp 的地址怎样变化​

令 M、N 方向的 warp 坐标为 wm、wn：

threadIdx.x = lane + 32*wm + 64*wn

将前面的地址公式扩展到所有 128 个线程：

m = 16*wm + lane%16

tXsA((p,g),0,0) 对应 A(m,8*g+p)

M 方向的另一个 warp 从第 16 行开始；N 方向的另一个 warp 使用相同的 A。所以线程 32 的两组行首是 A(16,0)、A(16,8)，线程 64 则与线程 0 一样使用 A(0,0)、A(0,8)。

每个 warp 独立执行自己的 ldmatrix，分发发生在各自 warp 内。

copy 怎样把两组视图拆成两次 Atom 调用​

源视图已经有了，目的端也要按同样的加载组组织。

本例每线程的 MMA fragment 与 Copy 目的视图为：

tCrA = ((2,2,2),1,1):((1,2,4),0,0)
tXrA = ((4,2),1,1):((1,4),0,0)

retile_D 保留 tCrA.data()，将 8 个 half 改为两组，每组 4 个。将组内目的编号记为 a=0…3：

// 两种访问引用同一份存储。
tXrA(make_coord(a,g),0,0);
tCrA(a+4*g,0,0);

这里只需要这个绑定结果。内部 retile 如何推导分组，可以另读retile_D 的逐步解析。

现在调用：

copy(tiled_copy, tXsA, tXrA);

copy(TiledCopy,...) 先转到 copy(Copy_Atom,...)。外层大小为 1 的 mode 处理后，Copy_Atom::call 看到的仍是一个嵌套的单 mode：

源：单 mode 的 shape 为 (8,2)，总共 16 个 half
目的：单 mode 的 shape 为 (4,2)，总共 8 个 half

而一次 .x2 只需要源 8 个 half、目的 4 个 half，所以这时还不能直接执行一次指令。Copy_Atom::call 会展开这个嵌套 mode：

return copy(*this, tensor<0>(src), tensor<0>(dst));

tensor<0> 取出第一个 mode 的布局，将 (8,2) 和 (4,2) 分别暴露为 rank 为 2 的视图。copy 再循环第二个 mode，也就是这里的加载组 g。每次留下源 8 个 half、目的 4 个 half，再进入 copy_unpack。

省略通用模板中的外层处理后，效果相当于：

for (int g=0; g<2; ++g) {
  auto src_atom = tXsA(make_coord(_,g),0,0);  // 8 个 half
  auto dst_atom = tXrA(make_coord(_,g),0,0);  // 4 个 half

  copy_unpack(Copy_Traits<SM75_U32x2_LDSM_N>{}, src_atom, dst_atom);
}

因此，从每个 warp 的指令流看，这里执行两条 .x2。每条都由 warp 的全部 lane 共同执行。

带着源 8、目的 4，读一遍 copy_unpack​

copy_unpack 将 Tensor 视图适配成底层指令封装的参数。它的核心代码是：

using CopyOp       = typename CPY_Op<AnyCPYTraits>::type;
using RegistersSrc = typename CopyOp::SRegisters;
using RegistersDst = typename CopyOp::DRegisters;
using RegTypeSrc   = typename remove_extent<RegistersSrc>::type;
using RegTypeDst   = typename remove_extent<RegistersDst>::type;
constexpr int RegNumSrc = extent<RegistersSrc>::value;
constexpr int RegNumDst = extent<RegistersDst>::value;

Tensor rS = recast<RegTypeSrc>(src);
Tensor rD = recast<RegTypeDst>(dst);

CUTE_STATIC_ASSERT_V(size(rS) == Int<RegNumSrc>{});
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{});

detail::explode(detail::CallCOPY<CopyOp>{},
                rS, make_int_sequence<RegNumSrc>{},
                rD, make_int_sequence<RegNumDst>{});

它依次做三件事：读取操作数类型、建立相应的存储视图、将视图展开为参数。下面仍然只代入 .x2。

读取指令需要的操作数形式​

SM75_U32x2_LDSM_N 声明：

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

remove_extent 去掉数组外壳，extent 取得数组长度。因此：

参数源端目的端
操作数类型uint128_tuint32_t
操作数数量12

这里 uint128_t 占 16 字节。它在这个加载接口中用于引用共享内存的一行，后面实际传给 PTX 的是其地址。

recast 同时改变元素类型和布局单位​

当前一次 Atom 调用接收到的两个视图都是连续的：

src：half，8:1
dst：half，4:1

执行：

Tensor rS = recast<uint128_t>(src);
Tensor rD = recast<uint32_t>(dst);

结果为：

视图原类型与布局新类型与布局
源half，8:1uint128_t，1:1
目的half，4:1uint32_t，2:1

recast 沿用原存储，并换算 shape、stride。源端元素类型扩大 8 倍，布局调用 upcast<8>；目的端扩大 2 倍，布局调用 upcast<2>。本例的分组完整、地址对齐，得到上表的结果。

源端 rS(0) 引用的起始位置，就是当前组 src(0) 的位置。目的端两个 32 位引用分别覆盖当前组的 half 槽位 0、1 和 2、3。

这些操作保留原有位模式，不进行 half 到整数的数值转换，也不搬运数据。

随后两个断言核对：重新解释后的源操作数数量是不是 1、目的操作数数量是不是 2。它们验证当前视图能否按该指令的参数数量展开。矩阵元素应当放在哪里，则由前面的布局组合和硬件路线确定。

explode 将视图展开成引用参数​

代入操作数数量后，最后一段是：

detail::explode(detail::CallCOPY<CopyOp>{},
                rS, make_int_sequence<1>{},
                rD, make_int_sequence<2>{});

explode 对这两个参数序列的展开方式是：

return fn(s[Is]..., d[Id]...);

所以本例展开为：

CopyOp::copy(rS[0], rD[0], rD[1]);

CallCOPY 使用转发引用把这些元素交给 CopyOp::copy，因此目的参数仍然引用原 fragment 的存储。

在 lane 0 中，两次调用的绑定是：

加载组源参数引用的行首rD[0] 绑定的原槽位rD[1] 绑定的原槽位
0&A(0,0)tCrA 的 0、1tCrA 的 2、3
1&A(0,8)tCrA 的 4、5tCrA 的 6、7

第二次调用的局部输出仍叫 dst0、dst1，但它们已经绑定到另一组槽位。第二次写 dst0 会写入原槽位 4、5，不会覆盖第一次的槽位 0、1。

指令怎样取得行首地址，再跨 lane 写入结果​

现在可以接上 SM75_U32x2_LDSM_N::copy。省略架构开关后，源码为：

static void copy(uint128_t const& smem_src,
                 uint32_t& dst0, uint32_t& dst1)
{
  uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src);

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

传进来的 smem_src 就是上一节的 rS[0]。封装取 &smem_src，转换成 shared memory 地址后，放进 PTX 的地址操作数。源端在这里使用引用，并没有先复制成一个 128 位的临时数据值。

对第 g 组：

&smem_src
  = &rS[0]
  = &src_atom(0)
  = &A(m,8*g)

这就是“行首地址”最终进入指令的位置。它不需要一个额外的行首标记：partition_S 已经让每组的起点对应正确的矩阵坐标，copy_unpack 与指令封装沿用这个起点。

本例的行首字节偏移为：

2*(16*m+8*g) = 32*m+16*g

共享内存基址按 16 字节对齐，所以每个行首也都按 16 字节对齐。这里依赖本文的 K 连续存储和所选 CopyAtom；partition_S 并不会替任意输入布局重新排列共享内存数据。

A(8,0)=23 为什么会进入 lane 0​

第一次加载时，两个 8×8 块分别覆盖 M=0…7 和 M=8…15，K 都为 0…7。lane 8 提供第二个块第 0 行的地址：

lane 8：&smem_src = &A(8,0)

硬件把这行的前两个 half 写入目的 lane 0 的 dst1。lane 0 的这个输出引用，又已经绑定到 tCrA 槽位 2、3：

A(8,0)=23，位于共享内存 half 偏移 128
    ↓ lane 8 提供行首，ldmatrix 读取
lane 0 的 dst1 低 half
    ↓ dst1 引用 tCrA 的槽位 2、3
lane 0 的 tCrA(2,0,0)=23

与此同时，lane 0 提供的 &A(0,0) 服务于另一条路线。第一组中，两端分别处理：

lane 0 的视图对应元素
源 tXsAA(0,0)…A(0,7)
目的 tXrAA(0,0)、A(0,1)、A(8,0)、A(8,1)

其中源行的后 6 个 half 会进入 lanes 1…3；目的的后两个 half 则来自 lane 8 提供的行地址。

源端每 lane 的视图描述 16 字节，目的端每 lane 接收 8 字节，这两个数量并不冲突：一次 .x2 实际使用 16 个不同的行地址，共读取 16*16=256 字节，分发给 32 个 lane，每个 lane 8 字节，也正好是 256 字节。

A(0,9)=37 怎样进入 lane 0 的槽位 5​

第二次加载从 K=8 开始。此时 lane 0 提供 &A(0,8)，指令从行内第 1 个 half 取得 A(0,9)=37：

第二组的源 lane 0：&A(0,8)
    ↓ 行内 half 1
A(0,9)=37，位于共享内存 half 偏移 9
    ↓ 第二组中，第一个 8×8 块的第 0 行
目的 lane 0 的 dst0 高 half
    ↓ 第二组 dst0 绑定 tCrA 的槽位 4、5
lane 0 的 tCrA(5,0,0)=37

这个例子中源 lane 与目的 lane 恰好相同；上一个例子则从源 lane 8 到目的 lane 0。两条路线都遵循同一套指令规则。

加载完成后，MMA 通过 tCrA 读取已经填好的槽位。retile_D 负责槽位绑定，copy_unpack 负责参数展开，数据的读取与 warp 内分发由 ldmatrix 完成。

用一个小程序核对地址和参数绑定​

配套程序为 cute_partition_s_copy_unpack_trace.cu。它使用本文相同的 2×2 warp 配置和 32×16 A tile。

主机端除了核对 partition_S 的坐标和连续性，还使用一个具有相同参数类型的捕获操作，调用真正的 copy_unpack。捕获操作核对源、目的引用的地址，验证 explode 最终交给底层函数的参数是否对应行首和正确的寄存器槽位。

随后程序按非转置 ldmatrix 的固定规则，在主机模拟分发，核对原 fragment 中的 half 位模式。它不在主机执行 GPU 指令。

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

/tmp/cute_partition_s_copy_unpack_trace

关键输出为：

tXsA = ((8,2),1,1):((1,8),0,0)
tXrA = ((4,2),1,1):((1,4),0,0)
lane0 slot2 = 23; slot5 = 37

thread0 row starts: A(0,0), A(0,8)
thread1 row starts: A(1,0), A(1,8)
thread8 row starts: A(8,0), A(8,8)
thread16 row starts: A(0,0), A(0,8)
thread32 row starts: A(16,0), A(16,8)
thread64 row starts: A(0,0), A(0,8)

PASS: 256 source rows, 256 copy_unpack calls, 1024 destination half slots, 512 packed outputs; host only.

还可以编译配套文件中的真实 GPU kernel，核对模板实例化后的指令：

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

该 kernel 将加载后的 fragment 写到输出，生成的 PTX 包含两条 ldmatrix.x2。它们分别对应本文的两组加载，每条在各 warp 内协作执行。本次只检查编译结果，没有在 GPU 上启动 kernel。

对照源码与官方文档​

CuTe 源码固定到 CUTLASS 提交 0b55a2f691d6。官方 PTX 文档核对日期为 2026-10-29。

要核对的内容链接
warp 协作、.x2 地址提供者与目的寄存器PTX ISA：ldmatrix
一个 8×8 块的目的 fragment 分布PTX 官方分布图
Copy 内部编号的含义copy_traits.hpp
.x2 的 Src、Dst、Ref bit 布局copy_traits_sm75.hpp
bit 布局换算到 half 布局copy_atom.hpp
make_tiled_copy_Acopy_atom.hpp
partition_S 与源到 reference 的转换copy_atom.hpp、线程切片
TiledCopy 转到 CopyAtom、按剩余 mode 分组copy.hpp、TiledCopy 重载
CopyAtom 展开嵌套 modecopy_atom.hpp、tensor<0>
copy_unpack 的完整实现copy_traits.hpp
recast 与类型比例换算tensor_impl.hpp、recast_layout
CallCOPY 转发与 explode 参数展开arch/util.hpp
.x2 取共享内存地址并执行 PTXcopy_sm75.hpp