cute（6）tiled_mma

鱿鱼圈 Lv4
  2026-06-13 22:13:32 2026-06-13 22:13:32 创建   2026-06-14 23:33:01 2026-06-14 23:33:01 更新      

前置阅读

6 TiledMMA

对应代码：06_tiled_mma.cu 纯 host 代码，不需要 GPU。

核心概念

MMA（Matrix Multiply-Accumulate）是 GPU 的矩阵乘硬件指令。CuTe 提供了三层抽象：

1
MMA_Atom（一条指令）→ EU Repeat（多 warp 覆盖更大面积）→ P Tile（每线程多算几轮）

第 1 节：MMA Atom 基础信息

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// ============================================================
 // 1. MMA Atom 基础信息
 // ============================================================
 {
   printf("=== 1. MMA Atom: SM80_16x8x16_F16F16F16F16_TN ===\n");

   using mma_op = SM80_16x8x16_F16F16F16F16_TN;
   using mma_traits = MMA_Traits<mma_op>;

   printf("  Shape_MNK = "); print(typename mma_traits::Shape_MNK{});
   printf("\n");  // (16, 8, 16)
   printf("  ThrID = "); print(typename mma_traits::ThrID{});
   printf("\n");  // 32 (一个 warp)
   printf("  每个 atom 需要 %d 个线程\n", int(size(typename mma_traits::ThrID{})));
   printf("\n");
 }

img

1
2
using mma_op = SM80_16x8x16_F16F16F16F16_TN;
using mma_traits = MMA_Traits<mma_op>;

命名解读 SM80_16x8x16_F16F16F16F16_TN：

字段 含义
SM80 Ampere 架构（A100 等）
16x8x16 Shape_MNK = (M=16, N=8, K=16)
F16F16F16F16 A=half, B=half, C=half, D=half
TN A 转置，B 不转置（数据排列方式）

关键属性：

1
2
Shape_MNK = (16, 8, 16)    → 一条指令计算 A(16×16) × B(8×16)^T = C(16×8)
ThrID = 32                  → 需要 32 个线程（1 个 warp）协作

一个 atom 是不可再分的最小单元，对应一条 PTX mma.sync 指令。

第 2 节：只用一个 Atom（不扩展）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
// ============================================================
  // 2. 只用一个 atom（不扩展）
  // ============================================================
  {
    printf("=== 2. TiledMMA: 1 atom (no repeat) ===\n");

    using mma_op = SM80_16x8x16_F16F16F16F16_TN;
    using mma_atom = MMA_Atom<MMA_Traits<mma_op>>;

    // 不扩展：(1, 1, 1) repeat, 只覆盖 (16, 8, 16)
    auto tiled_mma = make_tiled_mma(mma_atom{});

    printf("  TiledMMA: "); print(tiled_mma); printf("\n");
    printf("  线程数: %d\n", int(size(tiled_mma)));
    printf("\n");

    // 看 partition 的 shape（需要传 tensor，不能传 layout）
    half_t dummy_a[16*16] = {};
    half_t dummy_b[8*16] = {};
    half_t dummy_c[16*8] = {};
    auto test_A = make_tensor(&dummy_a[0], make_shape(Int<16>{}, Int<16>{}));   // (M, K)
    auto test_B = make_tensor(&dummy_b[0], make_shape(Int<8>{},  Int<16>{}));   // (N, K)
    auto test_C = make_tensor(&dummy_c[0], make_shape(Int<16>{}, Int<8>{}));    // (M, N)

    auto thr = tiled_mma.get_slice(0);  // thread 0
    auto fragA = thr.partition_fragment_A(test_A);
    auto fragB = thr.partition_fragment_B(test_B);
    auto fragC = thr.partition_fragment_C(test_C);

    printf("  Thread 0 fragment shapes:\n");
    printf("    fragA: "); print(shape(fragA)); printf("\n");
    printf("    fragB: "); print(shape(fragB)); printf("\n");
    printf("    fragC: "); print(shape(fragC)); printf("\n");
    printf("\n");

    // ---- 打印 Layout_TV 的 2D 可视化 ----
    // Layout_TV 是 (Thread, Value) -> 矩阵位置 的映射
    using mma_traits = MMA_Traits<mma_op>;
    printf("  LayoutC_TV (Thread x Value -> 16x8 矩阵位置):\n");
    auto layoutC_TV = typename mma_traits::CLayout{};
    printf("    "); print(layoutC_TV); printf("\n");
    print_layout(layoutC_TV);
    printf("\n");

    printf("  LayoutA_TV (Thread x Value -> 16x16 矩阵位置):\n");
    auto layoutA_TV = typename mma_traits::ALayout{};
    printf("    "); print(layoutA_TV); printf("\n");
    print_layout(layoutA_TV);
    printf("\n");

    printf("  LayoutB_TV (Thread x Value -> 8x16 矩阵位置):\n");
    auto layoutB_TV = typename mma_traits::BLayout{};
    printf("    "); print(layoutB_TV); printf("\n");
    print_layout(layoutB_TV);
    printf("\n");

    // ---- 打印每个线程的 partition 视图（用 partition_A/B/C 而非 fragment） ----
    // partition_A/B/C 返回对原始 tensor 的分片视图，保留指针关系
    auto pA = thr.partition_A(test_A);
    auto pB = thr.partition_B(test_B);
    auto pC = thr.partition_C(test_C);
    printf("  Thread 0 partition shapes (with pointer to original data):\n");
    printf("    partA: "); print(shape(pA)); printf("  layout: "); print(pA.layout()); printf("\n");
    printf("    partB: "); print(shape(pB)); printf("  layout: "); print(pB.layout()); printf("\n");
    printf("    partC: "); print(shape(pC)); printf("  layout: "); print(pC.layout()); printf("\n");
    printf("\n");

    // ---- 打印 fragC 的值（初始化为 0） ----
    printf("  fragC values (thread 0, %d elements):\n    ", int(size(fragC)));
    for (int i = 0; i < size(fragC); ++i) {
      printf("%.1f ", float(fragC(i)));
    }
    printf("\n\n");

    // ---- 打印多个线程的 partition_C 来观察不同线程负责哪些位置 ----
    printf("  各线程在 C(16x8) 矩阵中负责的位置 (offset):\n");
    for (int t = 0; t < 4; ++t) {  // 只打印前 4 个线程
      auto thr_t = tiled_mma.get_slice(t);
      auto pC_t = thr_t.partition_C(test_C);
      printf("    thread %d: ", t);
      for (int v = 0; v < size<0>(pC_t); ++v) {
        // 计算该元素相对于 test_C 起始地址的 offset
        int offset = int(&pC_t(v, 0, 0) - &test_C(0, 0));
        printf("%3d ", offset);
      }
      printf("\n");
    }
    printf("\n");
  }

img

img

img

img

img

1
auto tiled_mma = make_tiled_mma(mma_atom{});

不做任何扩展，32 个线程覆盖 (16, 8, 16)。

TiledMMA 打印信息解读

1
2
ThrLayoutVMNK:  (_32,_1,_1,_1):(_1,_0,_0,_0)
PermutationMNK: (_,_,_)

ThrLayoutVMNK 的四个维度是 (V, M, N, K)：

维度 值 含义
V=32 一个 atom 需要 32 个线程
M=1 M 方向没有复制
N=1 N 方向没有复制
K=1 K 方向没有复制

stride 中 _0 表示该维度不影响线程编号。

PermutationMNK 三个 _ 表示没有第二层扩展（P Tile）。

Layout_TV：寄存器分配地图

TV = Thread × Value。它回答：矩阵中的每个元素，由哪个线程的第几个寄存器持有？

1
Layout_TV: (thread_id, value_id) → 矩阵中的逻辑位置 (offset)
维度 含义
T (Thread) 线程编号 0~31
V (Value) 该线程持有的第几个元素（寄存器编号）

输出是矩阵的逻辑位置（offset），不是寄存器的物理位置。

LayoutC_TV

1
2
LayoutC_TV: ((_4,_8), (_2,_2)):((_32,_1), (_16,_8))
              Thread     Value

C 矩阵 16×8 = 128 个元素，32 个线程 × 4 个 value = 128，刚好分完。

Thread 0 那一行：

1
2
3
4
value 0 → offset=0   → 矩阵位置 (row=0, col=0)
value 1 → offset=16  → 矩阵位置 (row=2, col=0)   ← 不连续！
value 2 → offset=8   → 矩阵位置 (row=1, col=0)
value 3 → offset=24  → 矩阵位置 (row=3, col=0)

每个线程的 4 个元素分散在矩阵中的 2×2 小块，这个模式由 NVIDIA 硬件决定，不能改。

LayoutA_TV

1
LayoutA_TV: ((_4,_8), (_2,_2,_2)):((_32,_1), (_16,_8,_128))

每个线程持有 2×2×2 = 8 个 A 元素。A 矩阵 16×16 = 256，32 × 8 = 256 刚好分完。

LayoutB_TV

1
LayoutB_TV: ((_4,_8), (_2,_2)):((_16,_1), (_8,_64))

每个线程持有 2×2 = 4 个 B 元素。B 矩阵 8×16 = 128，32 × 4 = 128 刚好分完。

为什么 A 是 8 个而 B/C 是 4 个

MMA 指令的 shape 不对称：

1
2
3
A: (M, K) = (16, 16) = 256 / 32 = 8 个/线程
B: (N, K) = (8,  16) = 128 / 32 = 4 个/线程
C: (M, N) = (16,  8) = 128 / 32 = 4 个/线程

N 维只有 8（不是 16），所以 B 和 C 比 A 少一半。

Layout_TV 的作用

1. 加载数据时告诉线程去 smem 的哪个地址取

1
2
LayoutA_TV(thread=0, value=0) = 0     ← thread 0 要去 smem[0] 取
LayoutA_TV(thread=0, value=1) = 16    ← thread 0 要去 smem[16] 取

2. 存储数据时告诉线程往 smem 的哪个地址写

1
2
LayoutC_TV(thread=0, value=0) = 0     ← 寄存器 0 写到 smem[0]
LayoutC_TV(thread=0, value=1) = 16    ← 寄存器 1 写到 smem[16]

3. 坐标变换（retile_D）的基础

MMA 和 Copy 用不同的 Layout_TV，同一个矩阵位置在两种 layout 下有不同的 value 编号：

1
2
3
4
5
MMA:  LayoutC_TV(t0, value=2) = 矩阵位置 8
Copy: LayoutA_TV(t0, value=1) = 矩阵位置 8

同一个位置 8，MMA 放在寄存器 2，Copy 需要从寄存器 1 读
→ inv + compose 做坐标变换

注意：逻辑位置和 smem 物理地址之间还隔着 smem 的 layout（可能有 Swizzle），但 CuTe 的 partition 操作自动把这两层串起来，用户不需要手动换算。

Fragment Shape 的含义

1
2
3
fragA: ((_2,_2,_2), _1, _1)     = (MMA, MMA_M, MMA_K)
fragB: ((_2,_2),    _1, _1)     = (MMA, MMA_N, MMA_K)
fragC: ((_2,_2),    _1, _1)     = (MMA, MMA_M, MMA_N)
维度 含义
第一维 MMA 一次 mma 指令中该线程负责的元素数（= Layout_TV 的 Value 维大小）
第二维 MMA_M/MMA_N M 或 N 方向要重复多少次。_1 = 不需要重复
第三维 MMA_K/MMA_N K 或 N 方向要重复多少次。_1 = 不需要重复

每个线程持有的寄存器数量 = 三个维度之积：

1
2
3
fragA：8 × 1 × 1 = 8 个 half 寄存器
fragB：4 × 1 × 1 = 4 个 half 寄存器
fragC：4 × 1 × 1 = 4 个 half 累加器

partition_fragment vs partition 的区别

partition_fragment_C partition_C
数据在哪 新创建的寄存器（栈上） 指向原始 tensor 的视图
有指针关系吗 无，独立的数据 有，和原始 tensor 共享内存
用途 创建累加器/寄存器 fragment 分析每个线程负责原始矩阵的哪些位置

第 3 节：第一层扩展 —— EU Repeat (2, 2, 1)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
// ============================================================
  // 3. 第一层扩展：EU Repeat (2, 2, 1)
  // ============================================================
  {
    printf("=== 3. EU Repeat (2, 2, 1) ===\n");

    using mma_op = SM80_16x8x16_F16F16F16F16_TN;
    using mma_atom = MMA_Atom<MMA_Traits<mma_op>>;

    // M 方向 2 个 warp, N 方向 2 个 warp => 4 个 warp = 128 线程
    auto tiled_mma = make_tiled_mma(
        mma_atom{},
        make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})));

    printf("  线程数: %d (= 32 x 2 x 2)\n", int(size(tiled_mma)));

    // 覆盖范围：(16*2, 8*2, 16) = (32, 16, 16)
    half_t da[32*16]={}, db[16*16]={}, dc[32*16]={};
    auto test_A = make_tensor(&da[0], make_shape(Int<32>{}, Int<16>{}));
    auto test_B = make_tensor(&db[0], make_shape(Int<16>{}, Int<16>{}));
    auto test_C = make_tensor(&dc[0], make_shape(Int<32>{}, Int<16>{}));

    auto thr = tiled_mma.get_slice(0);
    auto fragA3 = thr.partition_fragment_A(test_A);
    auto fragB3 = thr.partition_fragment_B(test_B);
    auto fragC3 = thr.partition_fragment_C(test_C);
    printf("  Thread 0 fragment shapes:\n");
    printf("    fragA: "); print(shape(fragA3)); printf("\n");
    printf("    fragB: "); print(shape(fragB3)); printf("\n");
    printf("    fragC: "); print(shape(fragC3)); printf("\n");
    printf("\n");

    // 打印 partition 的 layout 细节
    auto pA3 = thr.partition_A(test_A);
    auto pB3 = thr.partition_B(test_B);
    auto pC3 = thr.partition_C(test_C);
    printf("  Thread 0 partition layouts:\n");
    printf("    partA: shape="); print(shape(pA3)); printf("  layout="); print(pA3.layout()); printf("\n");
    printf("    partB: shape="); print(shape(pB3)); printf("  layout="); print(pB3.layout()); printf("\n");
    printf("    partC: shape="); print(shape(pC3)); printf("  layout="); print(pC3.layout()); printf("\n");
    printf("\n");

    // ---- 打印 C(32x16) 矩阵的线程分配全景图 ----
    // 每个格子填写负责它的线程编号
    printf("  C(32x16) 矩阵线程分配图 (每格=线程号):\n");
    {
      // 建立 offset -> thread_id 的反向映射
      int thr_map[32*16];
      for (int i = 0; i < 32*16; ++i) thr_map[i] = -1;
      for (int t = 0; t < 128; ++t) {
        auto thr_t = tiled_mma.get_slice(t);
        auto pC_t = thr_t.partition_C(test_C);
        for (int v = 0; v < size<0>(pC_t); ++v) {
          int off = int(&pC_t(v, 0, 0) - &test_C(0, 0));
          if (off >= 0 && off < 32*16) thr_map[off] = t;
        }
      }
      // 打印：行=M，列=N
      printf("     ");
      for (int n = 0; n < 16; ++n) printf(" %3d", n);
      printf("\n");
      for (int m = 0; m < 32; ++m) {
        printf("  %2d:", m);
        for (int n = 0; n < 16; ++n) {
          int off = m * 16 + n;  // row-major in test_C (col-major default)
          // test_C 是 make_shape(32, 16) 默认 col-major: offset = m + 32*n
          int off_cm = m + 32 * n;
          printf(" %3d", thr_map[off_cm]);
        }
        printf("\n");
      }
    }
    printf("\n");
  }

img

1
2
3
auto tiled_mma = make_tiled_mma(
    mma_atom{},
    make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})));

EU Repeat = 用更多 warp 覆盖更大面积，增加线程数。

1
2
3
4
AtomLayoutMNK = (2, 2, 1)
→ M 方向 2 个 warp，N 方向 2 个 warp
→ 4 个 warp = 32 × 4 = 128 个线程
→ 覆盖范围：(16×2, 8×2, 16) = (32, 16, 16)

Fragment Shape

1
2
3
fragA: ((_2,_2,_2), _1, _1)    每线程还是 8 个元素
fragB: ((_2,_2),    _1, _1)    每线程还是 4 个元素
fragC: ((_2,_2),    _1, _1)    每线程还是 4 个元素

每个线程的工作量没变（还是一个 atom 的量），只是用了更多线程来覆盖更大的矩阵。

C(32×16) 线程分配全景图

从程序输出可以清晰看到 4 个 warp 的分工：

1
2
3
4
5
6
7
8
9
10
11
              N 方向
         col 0-7         col 8-15
        +---------------+---------------+
row 0-7 | warp 0 (t0-31)| warp 2 (t64-95)|   M 方向
        +---------------+---------------+   上半
row 8-15| warp 0 重复    | warp 2 重复    |
        +---------------+---------------+
row16-23| warp 1 (t32-63)| warp 3 (t96-127)|  M 方向
        +---------------+---------------+   下半
row24-31| warp 1 重复    | warp 3 重复    |
        +---------------+---------------+

每个线程负责 4 个元素（2×2 小块），每个格子出现两次是因为 MMA 指令的寄存器分配模式在 M 方向有 8 行的周期。

可视化代码原理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// 1. 建反向映射表：offset → thread_id
int thr_map[32*16];
for (int i = 0; i < 32*16; ++i) thr_map[i] = -1;

// 2. 遍历所有 128 个线程
for (int t = 0; t < 128; ++t) {
    auto thr_t = tiled_mma.get_slice(t);           // 取线程 t
    auto pC_t = thr_t.partition_C(test_C);          // 该线程负责的元素
    for (int v = 0; v < size<0>(pC_t); ++v) {       // 遍历 4 个 value
        int off = int(&pC_t(v, 0, 0) - &test_C(0, 0));  // 指针减法算 offset
        thr_map[off] = t;                            // 填表
    }
}

// 3. 按矩阵形状打印（注意 test_C 默认 col-major）
for (int m = 0; m < 32; ++m)
    for (int n = 0; n < 16; ++n)
        printf(" %3d", thr_map[m + 32 * n]);  // col-major: offset = row + num_rows × col

核心技巧是指针减法：partition_C 返回的 tensor 和原始 test_C 共享内存，&pC_t(v,0,0) - &test_C(0,0) 得到该元素在矩阵中的 offset。

第 4 节：第二层扩展 —— Permutation Tile

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
// ============================================================
  // 4. 第二层扩展：Permutation Tile
  // ============================================================
  {
    printf("=== 4. EU Repeat (2,2,1) + P Tile (32,32,16) ===\n");

    using mma_op = SM80_16x8x16_F16F16F16F16_TN;
    using mma_atom = MMA_Atom<MMA_Traits<mma_op>>;

    // EU Repeat 后覆盖 (32, 16, 16)
    // P Tile 要求 (32, 32, 16) => N 方向每线程多算一倍
    auto tiled_mma = make_tiled_mma(
        mma_atom{},
        make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})),  // EU Repeat
        Tile<Int<32>, Int<32>, Int<16>>{});                      // P Tile

    printf("  线程数: %d (还是 128，P Tile 不增加线程)\n", int(size(tiled_mma)));

    // 覆盖范围已是 (32, 32, 16)
    // 对 128x128 的 tile 做 partition
    half_t da4[128*32]={}, db4[128*32]={}, dc4[128*128]={};
    auto test_A = make_tensor(&da4[0], make_shape(Int<128>{}, Int<32>{}));
    auto test_B = make_tensor(&db4[0], make_shape(Int<128>{}, Int<32>{}));
    auto test_C = make_tensor(&dc4[0], make_shape(Int<128>{}, Int<128>{}));

    auto thr = tiled_mma.get_slice(0);
    auto fragA = thr.partition_fragment_A(test_A);
    auto fragB = thr.partition_fragment_B(test_B);
    auto fragC = thr.partition_fragment_C(test_C);

    printf("  Thread 0 fragment shapes (for 128x128x32 tile):\n");
    printf("    fragA (MMA, MMA_M, MMA_K): "); print(shape(fragA)); printf("\n");
    printf("    fragB (MMA, MMA_N, MMA_K): "); print(shape(fragB)); printf("\n");
    printf("    fragC (MMA, MMA_M, MMA_N): "); print(shape(fragC)); printf("\n");

    // 解读 fragC shape: (MMA, MMA_M, MMA_N)
    //   MMA    = 一次 mma 指令中该线程负责的元素数
    //   MMA_M  = M 方向重复次数 = 128 / 32 = 4
    //   MMA_N  = N 方向重复次数 = 128 / 32 = 4
    printf("\n");
    printf("  MMA_M = 128/32 = %d\n", int(size<1>(fragC)));
    printf("  MMA_N = 128/32 = %d\n", int(size<2>(fragC)));
    printf("  MMA_K = 32/16  = %d\n", int(size<2>(fragA)));
    printf("  每个线程持有的累加器元素数: %d\n", int(size(fragC)));
    printf("\n");

    // 打印 partition 的 layout 细节
    auto pA4 = thr.partition_A(test_A);
    auto pB4 = thr.partition_B(test_B);
    auto pC4 = thr.partition_C(test_C);
    printf("  Thread 0 partition layouts:\n");
    printf("    partA: shape="); print(shape(pA4)); printf("\n");
    printf("           layout="); print(pA4.layout()); printf("\n");
    printf("    partB: shape="); print(shape(pB4)); printf("\n");
    printf("           layout="); print(pB4.layout()); printf("\n");
    printf("    partC: shape="); print(shape(pC4)); printf("\n");
    printf("           layout="); print(pC4.layout()); printf("\n");
    printf("\n");

    // 打印几个线程在 C(128x128) 中第一个 MMA_M, MMA_N 块的位置
    printf("  各线程在 C(128x128) 中 MMA_M=0,MMA_N=0 时负责的位置:\n");
    for (int t = 0; t < 4; ++t) {
      auto thr_t = tiled_mma.get_slice(t);
      auto pC_t = thr_t.partition_C(test_C);
      printf("    thread %3d: ", t);
      for (int v = 0; v < size<0>(pC_t); ++v) {
        int offset = int(&pC_t(v, 0, 0) - &test_C(0, 0));
        printf("%5d ", offset);
      }
      printf("\n");
    }
    printf("    ...\n");
    // warp 2 的前两个线程 (thread 64-65)
    for (int t = 64; t < 66; ++t) {
      auto thr_t = tiled_mma.get_slice(t);
      auto pC_t = thr_t.partition_C(test_C);
      printf("    thread %3d: ", t);
      for (int v = 0; v < size<0>(pC_t); ++v) {
        int offset = int(&pC_t(v, 0, 0) - &test_C(0, 0));
        printf("%5d ", offset);
      }
      printf("\n");
    }
    printf("\n");

    // 打印 thread 0 在不同 MMA_M, MMA_N 块负责的首元素位置
    printf("  Thread 0 在各 (MMA_M, MMA_N) 块的首元素 offset:\n");
    for (int mm = 0; mm < size<1>(pC4); ++mm) {
      for (int mn = 0; mn < size<2>(pC4); ++mn) {
        int offset = int(&pC4(0, mm, mn) - &test_C(0, 0));
        printf("    (MMA_M=%d, MMA_N=%d): offset=%d\n", mm, mn, offset);
      }
    }
    printf("\n");

    // ---- 打印 C(128x128) 左上角 32x32 子块的线程分配图 ----
    printf("  C(128x128) 左上角 32x32 子块的线程分配图 (MMA_M=0, MMA_N=0):\n");
    printf("  (每格=线程号, 只显示 MMA_M=0,MMA_N=0 对应的第一个 atom 块)\n");
    {
      const int show_M = 32, show_N = 32;
      int thr_map4[show_M * show_N];
      for (int i = 0; i < show_M * show_N; ++i) thr_map4[i] = -1;
      for (int t = 0; t < 128; ++t) {
        auto thr_t = tiled_mma.get_slice(t);
        auto pC_t = thr_t.partition_C(test_C);
        // 只取 MMA_M=0, MMA_N=0 的 value
        for (int v = 0; v < size<0>(pC_t); ++v) {
          int off = int(&pC_t(v, 0, 0) - &test_C(0, 0));
          // test_C 是 (128, 128) col-major: offset = m + 128*n
          int m = off % 128;
          int n = off / 128;
          if (m < show_M && n < show_N) {
            thr_map4[m * show_N + n] = t;
          }
        }
      }
      printf("     ");
      for (int n = 0; n < show_N; ++n) printf(" %3d", n);
      printf("\n");
      for (int m = 0; m < show_M; ++m) {
        printf("  %2d:", m);
        for (int n = 0; n < show_N; ++n) {
          int val = thr_map4[m * show_N + n];
          if (val >= 0) printf(" %3d", val);
          else printf("   .");
        }
        printf("\n");
      }
    }
    printf("\n");
  }

img

img

1
2
3
4
auto tiled_mma = make_tiled_mma(
    mma_atom{},
    make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})),  // EU Repeat
    Tile<Int<32>, Int<32>, Int<16>>{});                      // P Tile

P Tile = 让每个线程多执行几次 MMA，不增加线程数。

EU Repeat 后覆盖 (32, 16, 16)，P Tile 要求 (32, 32, 16)：

1
2
3
M: 32 = 32  → 刚好
N: 32 / 16 = 2  → 每个线程在 N 方向多算一倍
K: 16 = 16  → 刚好

线程数不变（还是 128），但每个线程做更多工作。

两层扩展的区别

EU Repeat P Tile
做什么 用更多 warp 覆盖更大面积 让每个线程多执行几次 MMA
是否增加线程数 是 (32→128) 否 (还是 128)
覆盖范围怎么变 (16,8,16)→(32,16,16) (32,16,16)→(32,32,16)
代价 需要更多线程 每个线程做更多工作

Fragment Shape（对 128×128×32 tile）

1
2
3
fragA: ((_2,_2,_2), _4, _2)    = (MMA=8, MMA_M=4, MMA_K=2)
fragB: ((_2,_2),    _8, _2)    = (MMA=4, MMA_N=8, MMA_K=2)
fragC: ((_2,_2),    _4, _8)    = (MMA=4, MMA_M=4, MMA_N=8)

维度解读：

1
2
3
4
5
MMA_M = 128 / 32 = 4    → M 方向要重复 4 次
MMA_N = 128 / 32 = 8    → N 方向要重复 8 次（注意不是 128/16=8，而是 tile 被 P Tile 扩展到 32 后，128/16=8）
MMA_K = 32 / 16 = 2     → K 方向要重复 2 次

每线程累加器数 = 4 × 4 × 8 = 128 个 half 寄存器

为什么 MMA_N=8 而不是 4？

P Tile 把 N 的"一次覆盖量"从 16 扩展到 32，但 P Tile 内部 N=32 中每个线程实际覆盖的是 16（分两轮算）。对 128 列来说需要 128/16 = 8 次重复。

Thread 0 在各 (MMA_M, MMA_N) 块的首元素

1
2
3
(MMA_M=0, MMA_N=0): offset=0        ← 左上角
(MMA_M=0, MMA_N=1): offset=2048     ← 128×16=2048，跳 16 列
(MMA_M=1, MMA_N=0): offset=32       ← 往下跳 32 行（但是 col-major 所以 offset 差 32）

offset 差 2048 = 128×16，说明 MMA_N 每次覆盖 16 列。 offset 差 32，说明 MMA_M 每次覆盖 32 行（col-major 下 32 行 = offset 差 32）。

但这只是第一个 atom 块（MMA_M=0, MMA_N=0）的分配。

C(128×128) 左上角 32×32 的线程分配图

输出中可以看到：

  • 左半 16 列 (col 0-15)：有线程号，pattern 和第 3 节完全一样
  • 右半 16 列 (col 16-31)：全是 .（空）

这是因为只打印了 MMA_M=0, MMA_N=0 的分配。MMA_N=0 的一个 atom 块只覆盖前 16 列，col 16-31 属于 MMA_N=1 的下一轮。

可视化代码原理

1
2
3
4
5
6
7
8
9
10
11
// 和第 3 节相同思路，但需要处理 128×128 的 col-major tensor
for (int t = 0; t < 128; ++t) {
    auto pC_t = thr_t.partition_C(test_C);
    for (int v = 0; v < size<0>(pC_t); ++v) {
        int off = int(&pC_t(v, 0, 0) - &test_C(0, 0));  // 只取 MMA_M=0, MMA_N=0
        int m = off % 128;    // col-major 反算 row
        int n = off / 128;    // col-major 反算 col
        if (m < 32 && n < 32) // 只取左上角 32×32
            thr_map4[m * 32 + n] = t;
    }
}

与第 3 节的区别：

  1. pC_t(v, 0, 0) — 固定 MMA_M=0, MMA_N=0，只看第一个 atom 块
  2. off % 128 / off / 128 — test_C 是 (128, 128) col-major，需要用 128 而不是 32 来反算坐标
  3. 只显示左上角 32×32，因为 128×128 太大

完整扩展链路

1
2
3
4
5
6
7
8
9
10
11
atom:               (16,  8, 16)  ×  32 线程
                         ↓ EU Repeat (2, 2, 1)
第一层:             (32, 16, 16)  × 128 线程    ← 线程数增加
                         ↓ P Tile (32, 32, 16)
第二层:             (32, 32, 16)  × 128 线程    ← 线程数不变，每线程多算

对 128×128×32 tile:
  MMA_M = 128/32 = 4 次
  MMA_N = 128/16 = 8 次    (P Tile 的 N=32 中每线程实际覆盖 16)
  MMA_K = 32/16  = 2 次
  每线程累加器 = 4 × 4 × 8 = 128 个

make_tiled_mma 参数详解

1
make_tiled_mma(MMA_Atom, AtomLayoutMNK, PermutationMNK)
参数 含义 示例
MMA_Atom 基础 MMA 指令 MMA_Atom<MMA_Traits<SM80_16x8x16_F16F16F16F16_TN>>
AtomLayoutMNK EU Repeat：在 M/N/K 方向各复制几份 Layout<Shape<_2, _2, _1>>
PermutationMNK P Tile：最终期望的 tile 大小 Tile<Int<32>, Int<32>, Int<16>>

API 速查

API 作用
make_tiled_mma(atom, repeat, tile) 创建 TiledMMA
tiled_mma.get_slice(threadIdx.x) 取当前线程的 MMA 切片
thr.partition_fragment_A/B/C(tensor) 在寄存器上创建 fragment（独立数据）
thr.partition_A/B/C(tensor) 对原始 tensor 分片（共享内存，保留指针）
size(tiled_mma) 总线程数
size<0>(fragC) 每线程每次 MMA 的元素数
size<1>(fragC) MMA_M 重复次数
size<2>(fragC) MMA_N 重复次数

易错点总结

问题 原因 解决方法
partition_fragment_A 传 layout 报错 需要传 tensor（有数据指针），不能传 layout 创建 dummy 数组包装成 tensor
fragment shape 的 MMA_N 比预期大 P Tile 扩展了覆盖范围 MMA_N = tile_N / (atom_N × 每线程覆盖的 N)
可视化时 offset 换算错误 tensor 默认 col-major，不是 row-major row = off % num_rows, col = off / num_rows
  • 标题: cute（6）tiled_mma
  • 作者: 鱿鱼圈
  • 创建于 : 2026-06-13 22:13:32
  • 更新于 : 2026-06-14 23:33:01
  • 链接: https://yuyanqi.com/2026/06/13/cute(6)tiled_mma/
  • 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。
评论
目录
cute（6）tiled_mma
  1. 前置阅读
  2. 6 TiledMMA
    1. 核心概念
    2. 第 1 节：MMA Atom 基础信息
    3. 第 2 节：只用一个 Atom（不扩展）
      1. TiledMMA 打印信息解读
      2. Layout_TV：寄存器分配地图
      3. Layout_TV 的作用
      4. Fragment Shape 的含义
      5. partition_fragment vs partition 的区别
    4. 第 3 节：第一层扩展 —— EU Repeat (2, 2, 1)
      1. Fragment Shape
      2. C(32×16) 线程分配全景图
      3. 可视化代码原理
    5. 第 4 节：第二层扩展 —— Permutation Tile
      1. 两层扩展的区别
      2. Fragment Shape（对 128×128×32 tile）
      3. Thread 0 在各 (MMA_M, MMA_N) 块的首元素
      4. C(128×128) 左上角 32×32 的线程分配图
      5. 可视化代码原理
    6. 完整扩展链路
    7. make_tiled_mma 参数详解
    8. API 速查
    9. 易错点总结