G2 讲解：CuTe 布局代数

① 先看全貌：Layout = 函数

Layout(shape, stride)，记作 shape:stride。它把坐标映射成偏移：offset = Σ coord_i × stride_i。

两个 layout 的 shape 一样，只是 stride 不同，于是“第 5 个元素”在内存里的位置不同：列主序是 5，行主序是 3。

② 下标 → 坐标：列字典序

drills_idx2crd（26–35 行） · drills_inner / call（38–52 行）

整数下标先按列字典序（第一个模式变化最快）拆成坐标，再和 stride 做点积。

下标 5，shape (4, 2)：5 % 4 = 1，5 // 4 = 1，1 % 2 = 1，坐标 (1, 1)。

• 列主序 (4,2):(1,4)：1·1 + 1·4 = 5
• 行主序 (4,2):(2,1)：1·2 + 1·1 = 3

嵌套的 shape 按嵌套结构递归拆：((2,2),3):((1,2),4) 的下标 7 → 7 % 4 = 3 交给内层 (2,2) 拆成 (1,1)，7 // 4 = 1 给外层，坐标是 ((1,1),1)，偏移 1·1 + 1·2 + 1·4 = 7。嵌套让一个“逻辑上的维度”由几个物理维度组成，比如“一行 4 个元素，其实是 2 个线程各拿 2 个”。

和 G1 的区别：G1 用 order 指定哪一维最快；CuTe 固定第一个模式最快，想要别的顺序就改 stride。

size 和 cosize

drills54–58 行：size 是坐标的个数（shape 的乘积），cosize 是最大偏移 + 1。(4,2):(1,4) 的 size 和 cosize 都是 8（紧凑）；(4,2):(2,8) 的 size 是 8，cosize 是 3·2 + 1·8 + 1 = 15，中间有空隙。

两种紧凑布局

drillsmake_layout_right / make_layout_left（64–86 行）：make_layout_right((4,2)) = (4,2):(2,1)（行主序，最后一个模式 stride 为 1）；make_layout_left((4,2)) = (4,2):(1,4)（列主序，stride 是 shape 的前缀积）。

③ coalesce：同一个函数，最少的模式

drillscoalesce（89–104 行）

规则只有两条：去掉大小为 1 的模式；相邻两个模式 s0:d0 和 s1:d1，如果 d1 == s0 × d0（第二个正好接着第一个走完的地方），合并成 (s0·s1):d0。

coalesce 不改变函数（每个下标映射到的偏移不变），只是把描述化简。在 composition 里先对 A 做 coalesce，可以减少要处理的模式数。

④ composition：用 B 去选 A 的元素

drillscomposition（107–135 行）

R = A ∘ B，R(i) = A(B(i))。B 先把 i 变成一个“A 的下标”，A 再把它变成偏移。

例子：A = (4,2):(2,1)（行主序 4×2），B = 4:2（取第 0、2、4、6 个）。

参考解给出的 R = (2,2):(4,1)：下标 0..3 → 0, 4, 1, 5 ✓。

注意结果不是简单的“stride 乘 2”：B 跨的步子跨过了 A 的第一个模式的边界（A 的第一维只有 4 个），所以结果被拆成了两个模式。参考解的算法就是在做这件事：对 B 的一个模式 s:d，沿 A 的模式往下走，先“跳过前 d 个”（除），再“取 s 个”（取模），每跨过 A 的一个模式边界，就在结果里多出一个模式。

另一个例子：A = (4,4):(4,1)（行主序 4×4），B = (4,4):(4,1)。A∘B = (4,4):(1,4)：用一个行主序的选取方式去读一个行主序的矩阵，结果相当于转置。composition 就是“重新排列数据”的统一表达。

⑤ complement：A 的空隙

drillscomplement（138–152 行）

complement(A, M) 是 A 在 [0, M) 里“没覆盖到的部分”的布局。A 和它的 complement 拼起来，恰好不重不漏地铺满 [0, M)。

第二行：A 每隔一个取一个，取 4 个，覆盖了 0–7 里的偶数位置。空隙有两种：“奇数位”（偏移 +1）和“后半段”（偏移 +8）。A(i) + complement(j) 遍历所有组合，正好是 0..15 每个恰好一次。

算法（注释 141–142 行）：把 A 的模式按 stride 排序，每两个模式之间的“空档”变成 complement 的一个模式，最后再补一个模式一直延伸到 M。

⑥ logical_divide：切成 tile × 重复

drillslogical_divide（155–159 行）

logical_divide(A, B) = A ∘ (B, complement(B, size(A)))

结果有两个模式：模式 0 是“一个 tile 里的元素”（B 选出来的），模式 1 是“tile 的重复”（complement 选出来的）。

第二行展开（模式 0 是 tile 内的第几个，模式 1 是第几个 tile），每个 tile 拿到的偏移：

tiler 4:2 决定了每个 tile 的形状（隔一个取一个），complement (2,2):(1,8) 决定了 tile 之间怎么排（先 +1，再 +8）。16 个偏移每个恰好出现一次。

这为什么有用：GEMM 里“把 C 矩阵切成 128×128 的 tile，第 (m, n) 个 block 拿第 (m, n) 个 tile”，在 CuTe 里就是 local_tile = 对 C 的 layout 做 divide 再取模式 1 的某个坐标。“把一个 tile 分给 128 个线程，每个线程拿 4 个”也是 divide。C4 里你手算的 tid % 16、tid / 16 和 4×4 小块，在 CuTe 里都可以写成 layout 运算。

⑦ 对照 CuTe 源码

• C++（CuTe 本体）：cutlasscoalesce · cutlasscomposition · cutlasscomplement · cutlasslogical_divide
• pycute（官方的纯 Python 参考实现，和这题最接近）：cutlasscoalesce · cutlasscomposition · cutlasscomplement · cutlasslogical_divide · cutlasszipped_divide
• 官方文档：cutlass01_layout.md · cutlassCoalesce · cutlassComposition · cutlassComplement · cutlassDivision
• CuTe DSL notebook：cutlasscute_layout_algebra.ipynb
• 实际 kernel：cutlassCuTe sgemm_1（C4 的同一个算法，用 layout 运算写成）· swizzle（C3 讲过的异或）cutlassSwizzle

⑧ 常见坑与自测

• 按行字典序拆下标：CuTe 是第一个模式最快。(4,2) 的下标 1 是 (1, 0)，不是 (0, 1)。
• coalesce 时比较顺序弄反：条件是 d1 == s0 × d0，用的是前一个模式的 shape 和 stride。
• 以为 composition 的结果模式数和 B 一样：B 跨过 A 的模式边界时，结果会多出模式（第 ④ 节的例子）。
• complement 用了未排序的模式：要先按 stride 排序。

源码与资料索引

本仓库：drillssolutions/g2_cute_layout.py · drillsmine/g2_cute_layout.py · drillstests/test_g2_cute_layout.py · drillscontent/g2.md · 前置：G1 讲解

CuTe：cutlasslayout algebra 文档 · cutlasspycute · cutlasslayout.hpp · cutlassCuTe DSL notebook