从 atile 到 upcast：读懂 CuTe 的 A 广播与右逆布局

读 retile_D 的内部实现时，有一行容易让人困惑：

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

假设 A tile 是 32×16，只有 512 个矩阵元素，而 TiledNumThr*V 却等于 128×8=1024。为什么能做 upcast<1024>？为什么右逆布局中会出现 stride 512？为什么最后两个矩阵方向的范围都是 1？

这些问题的起点是 get_layoutA_TV() 中的 atile。它把 N 方向的线程广播接入 A 布局，使同一个 A 元素对应多个 TV，但保留完整线程域的编号。下面固定同一组配置，沿着源码把这件事算清楚。

先区分矩阵坐标与 TV 坐标

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

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

它的整个 MMA tile 是：

(M,N,K) = (16*2, 8*2, 16*1) = (32,16,16)

A 使用 (M,K) 坐标，所以 A tile 是 32×16。本文讨论恰好覆盖这个 tile 的 fragment，每线程第一维有 8 个 half 值。

符号	本文含义	取值或范围
lane	warp 内的线程编号	0…31
wm	M 方向的 warp 坐标	0、1
wn	N 方向的 warp 坐标	0、1
t	整个 TiledMMA 的线程编号	0…127
v	线程内一个 half 的值编号	0…7
q	A tile 的逻辑矩阵编号	q=m+32*k，范围 0…511
tv	展平后的线程和值编号	tv=t+128*v，范围 0…1023

其中：

t = lane + 32*wm + 64*wn

q 编码矩阵位置，tv 编码线程和值的位置。两者有不同的编号范围。本文的 q 也不是共享内存地址偏移；实际存储地址还要由源 tensor 的布局计算。

make_tiled_copy_A 直接采用 MMA 的 A 布局和 A tile 大小：

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

因此这一例子中：

TiledLayout_TV = mma.get_layoutA_TV()
TiledNumThr    = 128
V              = size<0>(tCrA) = 8
Tiler_MN       = (32,16)

Tiler_MN 是 Copy 实现的通用名字；在 A 操作数这里，两维实际是 M、K。V 取输入 fragment 的第一维大小，不是任意大 fragment 的总元素数。

atile 为何让 512 个 A 元素对应 1024 个 TV

get_layoutA_TV() 先创建参考矩阵布局：

auto ref_A = make_layout(
    make_shape(tile_size_mnk<0>(), tile_size_mnk<2>()));

本例中它是：

ref_A = (32,16):(1,32)
ref_A(m,k) = m+32*k

随后通过以下组合返回最终的 TV 布局：

return thrfrg_A(ref_A)
    .compose(atile, _)
    .compose(thridx_2_thrid, _);

先看 thrfrg_A(ref_A)。它将 MMA Atom 的 A 分工与 M、K 方向的线程排列组合起来，线程输入结构是：

(ThrV,(ThrM,ThrK))

这里 ThrV 是 Atom 内的线程坐标，本例就是 lane；它与线程内的值编号 v 不同。由于 A 只需要 M、K，当前线程域大小是：

32*2*1 = 64

配合每线程 8 个值，共有 64*8=512 个 TV，恰好覆盖 A 的 512 个不同元素。

但是，完整的 MMA 线程组织还包含 N 方向。源码用 atile 将它接入：

auto atile = make_tile(
    _,
    make_tile(
        make_layout(
            make_shape(size<1>(thr_layout_vmnk_),
                       size<2>(thr_layout_vmnk_)),
            make_stride(Int<1>{}, Int<0>{})),
        _));

本例中间那张布局是：

(ThrM,ThrN):(1,0) = (2,2):(1,0)

(wm,wn) → wm+0*wn → wm

atile 在求值时保留 wm、忽略 wn。将它接到 thrfrg_A(ref_A) 的输入端后，新布局可以接收包含 N 的完整线程坐标，而不同 wn 会落到相同的 A 位置。

阶段	线程域大小	TV 数量	不同 A 元素数量
thrfrg_A(ref_A)	32*2=64	64*8=512	512
接入 atile 后	32*2*2=128	128*8=1024	512

这就是 A 在 N 方向的广播：两个 N warp 分块计算不同的输出列，但使用相同的 A。广播表达重复的元素需求，没有将完整线程域压缩成 64。

最后的 thridx_2_thrid 将线程编号转换成逻辑线程坐标；本例它等价于 128:1。最终 F 的输入仍是全部 128 个线程及其值编号。

把 get_layoutA_TV 写成一个可以计算的公式

为了看清后面的 stride，先拆开 lane 和值编号：

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

这条 MMA Atom 的 A 布局给出：

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

M 方向的 warp 排列再增加 16*wm；N 方向经过 atile 广播，不增加 A 的矩阵坐标。因此：

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

代入 q=m+32*k：

F(t,v) = q
       = 64*t0+t1+16*wm+0*wn+32*v0+8*v1+256*v2

省略大小为 1 的 mode，合并后的布局正是：

子轴：      t0  t1  wm  wn  v0  v1   v2
shape：      4   8   2   2   2   2    2
q stride：  64   1  16   0  32   8  256

wn 的 stride 为 0，但 shape 仍为 2。这一个子轴同时说明了两件事：输入 TV 数量翻倍，输出的不同矩阵元素数量不变。

例如，两个 N warp 中的 lane 0、值 4 都需要 A(0,8)：

F(t=0, v=4) = 256
F(t=64,v=4) = 256

它们的原 TV 编号分别是：

0 +128*4 = 512
64+128*4 = 576

right_inverse 选一个代表，但保留原来的 TV 编号

令：

auto F = mma.get_layoutA_TV();
auto G = right_inverse(F);

G 需要满足：

F(G(q)) = q

也就是每个矩阵元素反查出一个可用的 TV。它不必还原所有重复的 TV，因此不要求 G(F(tv))=tv。

源码先对 F 的完整 shape 计算前缀积，作为原输入展平编号的权重：

子轴：      t0  t1  wm  wn   v0   v1   v2
shape：      4   8   2   2    2    2    2
TV 权重：    1   4  32  64  128  256  512

这些权重也可以直接从 TV 编号公式中看出：

tv = t+128*v
   = t0+4*t1+32*wm+64*wn+128*v0+256*v1+512*v2

接着，right_inverse 按 q stride 排序，构造连续覆盖。wn 的零 stride 无法扩展连续覆盖，因此被跳过，本例相当于选 wn=0 的代表。

跳过广播维度不会重新计算原 TV 编号的权重。 v2 的权重仍包含完整线程域中的 ThrN=2：

v2 的 TV 权重 = 4*8*2*2*2*2 = 512

所以：

G(256) = 512

F(512)    = 256
G(F(576)) = 512，而不是 576

这里 F 接收单个整数时，将它按自身 shape 解读为展平的 TV 坐标。G 选中了 TV 512，不再返回同一元素的另一个代表 TV 576。

实际得到的右逆布局为：

G = (8,2,2,2,4,2):(4,256,32,128,1,512)
size(G) = 512

size(G)=512 限制的是输入数量，输出仍然是原来 0…1023 空间中的 TV 编号。例如一个简单布局 2:512 也只有两个输入，却分别输出 0 和 512；shape 不会限制 stride 必须小于元素数量。

with_shape 把右逆的输入组织成 M、K

现在看：

auto H = right_inverse(F)
             .with_shape(make_shape(_32{}, _16{}));

with_shape 的实现是：

return composition(*this, make_layout(shape));

因此 H 的含义为：

H(m,k) = G(m+32*k)

它将矩阵的二维坐标接到 G 的输入端，输出仍是 G 选中的 TV 编号。实际布局为：

H = ((8,2,2),(2,4,2)):((4,256,32),(128,1,512))

把它的两个顶层轴拆开：

m = m0+8*m1+16*m2           shape=(8,2,2)
k = k0+2*k1+8*k2            shape=(2,4,2)

与前面的坐标对应：

(m0,m1,m2) = (t1,v1,wm)
(k0,k1,k2) = (v0,t0,v2)

因此输出 TV 编号是：

H(m,k) = 4*m0+256*m1+32*m2+128*k0+k1+512*k2

现在 stride 512 的意义就很具体了：k2 从 0 变成 1，矩阵 K 增加 8，对应线程内 v 增加 4，原 TV 编号增加 128*4=512。

输入矩阵坐标	q	选中的 (t,v)	H 输出的 tv
(0,0)	0	(0,0)	0
(0,8)	256	(0,4)	512
(0,9)	288	(0,5)	640

H 虽然只有 512 个输入，输出却没有压紧为 0…511。 选中 wn=0 后，每个 v 对应的 TV 编号是 128*v…128*v+63，后面的 64 个位置仍然留给未被右逆选中的 N 广播副本。

with_shape 没有凭空产生 stride 512。它将右逆已有的 TV 权重组织到了 K 的对应子轴上。

upcast 为什么使用 1024，而不是 512

回到目标表达式：

auto frg_layout_mn = upcast<128*8>(H);

128*8 来自完整线程数和输入 fragment 的第一维值数。它表示全部线程各自前 8 个值所对应的 TV 范围。

这一范围大小为 1024。虽然右逆只选了其中 512 个代表，它们的编号仍散布在这个范围中，例如：

v=0：  0…63
v=1：128…191
v=2：256…319
...
v=7：896…959

upcast 改变布局的编号粒度，并同步缩小输入 shape。它不要求输入坐标数量大于或等于参数 N；这里也没有 Tensor 元素类型的转换。

对本例正的静态 stride，源码逐个子轴计算：

new_shape  = ceil_div(old_shape, ceil_div(N, old_stride));
new_stride = ceil_div(old_stride, N);

本例 N 与各 stride 都满足源码的整除条件，所以下面可以用“每组覆盖多少个原坐标位置”来理解内层除法。

先看 K 的最后一个子轴 (shape,stride)=(2,512)。它有两个坐标，对输出 TV 编号的贡献分别为 0 和 512。以 1024 为一组，这两个位置在同一组中：

每组覆盖的原坐标数量 = ceil_div(1024,512) = 2
所需组数             = ceil_div(2,2) = 1

其他子轴也一样：

子轴	原 shape	原 stride	每组覆盖的原坐标数	新 shape
M₀	8	4	256	ceil_div(8,256)=1
M₁	2	256	4	ceil_div(2,4)=1
M₂	2	32	32	ceil_div(2,32)=1
K₀	2	128	8	ceil_div(2,8)=1
K₁	4	1	1024	ceil_div(4,1024)=1
K₂	2	512	2	ceil_div(2,2)=1

ceil_div 是向上取整的除法。例如 8/256=0.03125，向上取整就是 1：这一子轴只有 8 个位置，虽然不足一整组，仍需要保留一组。

所有原 stride 都大于 0、小于 1024，所以 ceil_div(old_stride,1024) 也都等于 1。实际结果是：

frg_layout_mn
  = ((1,1,1),(1,1,1)):((1,1,1),(1,1,1))

product_each(shape(frg_layout_mn)) = (1,1)

作为核对，H 的最大输出是：

7*4 + 1*256 + 1*32 + 1*128 + 3*1 + 1*512 = 959

全部选定 TV 都在第一个 1024 范围内，与上面的布局计算一致。返回的新布局已经缩小了输入域，不能把它当成保留全部原 (m,k) 输入、仅将每个输出做整数除法。

如果对同一个 H 改用 upcast<128*4>，K 最后一个子轴就会留下：

new_shape = ceil_div(2, ceil_div(512,512)) = 2

此时 product_each(shape(...))=(1,2)，对应 K 的前后两段。这说明压缩粒度取决于所考虑的每线程值范围，而不是简单取矩阵的不同元素数量。

放回 retile_D，结果为 1 表示什么

retile 随后会使用这份布局计算 fragment 收集范围：

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

本例 product_each(...)=(1,1)，因此后一个 tiler 是 (8,1,1)：第一维保留每线程的 8 个值，M、K 方向各取一个 fragment 位置。

这 8 个 half 没有消失。后面的 logical_product(make_layout(V), ...) 继续保留它们，再按 Copy Atom 每次需要的值数分组。例如 A 使用 ldmatrix.x2 时，每次写入 4 个 half，最终仍会将 8 个值组织成 (4,2)。

右逆选择 wn=0 的代表，也不会使其他 N warp 停止加载。这里计算的是分组布局；ThrCopy::retile_D 最终为每个线程使用该线程自身 fragment 的 data() 建立视图，各 warp 仍然准备自己需要的 A 寄存器。

这条调用链中，atile 建立 A 的 N 广播，right_inverse 选择保留原编号的 TV 代表，with_shape 接入 M、K 坐标，upcast 则提取按当前 V 范围划分的 fragment 收集布局。

对照源码

本文依据 CUTLASS 提交 0b55a2f691d6。布局数值已用主机端程序核对，包括全部 1024 个 TV 的组合结果，以及全部 512 个矩阵编号的 F(G(q))=q；无需启动 GPU kernel 即可验证这些布局关系。

内容	源码
MMA Atom 的 A 布局	mma_traits_sm80.hpp
thrfrg_A 的 M、K 线程划分	mma_atom.hpp
get_layoutA_TV 与 atile	mma_atom.hpp
make_tiled_copy_A 选择布局和 tile	copy_atom.hpp
right_inverse 的前缀积和连续覆盖	layout.hpp
with_shape 的组合	layout.hpp
upcast 的 shape、stride 公式	layout.hpp
内部 retile 与外层 retile_D	copy_atom.hpp、copy_atom.hpp

若要继续对照实际加载的源地址、寄存器槽位和 A、B 的不同分组，可以阅读从共享内存到 MMA：读懂 A、B 的 ldmatrix 与 retile_D。