前置阅读

官方文档：CuTe — NVIDIA CUTLASS Documentation
reed博士：(4 封私信 / 2 条消息) cute 之 Swizzle - 知乎

8 Shared Memory GEMM

对应代码：08_smem_gemm.cu 需要 GPU。

核心概念

在 07 的基础上加入 Shared Memory，实现两级数据搬运：

G2S：Global → Shared（cp.async 128bit 异步搬运）
S2R：Shared → Register（ldmatrix 指令）
Swizzle 消除 shared memory bank conflict
make_tiled_copy_A/B 与 TiledMMA 的线程布局自动匹配
retile_D：在 MMA 和 Copy 两种寄存器视角间切换

代码实现

template <typename Config>
__global__ void smem_gemm_kernel(half_t* Dptr, const half_t* Aptr,
                                 const half_t* Bptr, int M, int N, int K) {
  using T = half_t;
  using SmemLayoutA = typename Config::SmemLayoutA;
  using SmemLayoutB = typename Config::SmemLayoutB;
  using TiledMMA = typename Config::MMA;
  using S2RCopyAtomA = typename Config::S2RCopyAtomA;
  using S2RCopyAtomB = typename Config::S2RCopyAtomB;
  using G2SCopyA = typename Config::G2SCopyA;
  using G2SCopyB = typename Config::G2SCopyB;

  constexpr int kTileM = Config::kTileM;
  constexpr int kTileN = Config::kTileN;
  constexpr int kTileK = Config::kTileK;

  extern __shared__ T shm_data[];
  T* Ashm = shm_data;
  T* Bshm = shm_data + cosize(SmemLayoutA{});

  int idx = threadIdx.x;
  int bx = blockIdx.x;
  int by = blockIdx.y;

  // Global Tensor
  auto A = make_tensor(make_gmem_ptr(Aptr), make_shape(M, K), make_stride(K, Int<1>{}));
  auto B = make_tensor(make_gmem_ptr(Bptr), make_shape(N, K), make_stride(K, Int<1>{}));
  auto D = make_tensor(make_gmem_ptr(Dptr), make_shape(M, N), make_stride(N, Int<1>{}));

  auto gA = local_tile(A, make_tile(Int<kTileM>{}, Int<kTileK>{}), make_coord(by, _));
  auto gB = local_tile(B, make_tile(Int<kTileN>{}, Int<kTileK>{}), make_coord(bx, _));
  auto gD = local_tile(D, make_tile(Int<kTileM>{}, Int<kTileN>{}), make_coord(by, bx));

  // Shared Memory
  auto sA = make_tensor(make_smem_ptr(Ashm), SmemLayoutA{});
  auto sB = make_tensor(make_smem_ptr(Bshm), SmemLayoutB{});

  // MMA
  TiledMMA tiled_mma;
  auto thr_mma = tiled_mma.get_slice(idx);
  auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));
  auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0));
  auto tCrD = thr_mma.partition_fragment_C(gD);
  clear(tCrD);

  // S2R Copy
  auto s2r_tiled_copy_a = make_tiled_copy_A(S2RCopyAtomA{}, tiled_mma);
  auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(idx);
  auto tAsA = s2r_thr_copy_a.partition_S(sA);
  auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA);

  auto s2r_tiled_copy_b = make_tiled_copy_B(S2RCopyAtomB{}, tiled_mma);
  auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(idx);
  auto tBsB = s2r_thr_copy_b.partition_S(sB);
  auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB);

  // G2S Copy
  G2SCopyA g2s_tiled_copy_a;
  auto g2s_thr_copy_a = g2s_tiled_copy_a.get_slice(idx);
  auto tAgA_copy = g2s_thr_copy_a.partition_S(gA);
  auto tAsA_copy = g2s_thr_copy_a.partition_D(sA);

  G2SCopyB g2s_tiled_copy_b;
  auto g2s_thr_copy_b = g2s_tiled_copy_b.get_slice(idx);
  auto tBgB_copy = g2s_thr_copy_b.partition_S(gB);
  auto tBsB_copy = g2s_thr_copy_b.partition_D(sB);

  // 打印调试信息
  if (idx == 0 && bx == 0 && by == 0) {
    printf("\n  === G2S Copy ===\n");
    printf("  tAgA_copy shape: "); print(shape(tAgA_copy)); printf("\n");
    printf("  tAgA_copy layout: "); print(tAgA_copy.layout()); printf("\n");
    printf("  tAsA_copy shape: "); print(shape(tAsA_copy)); printf("\n");
    printf("  tAsA_copy layout: "); print(tAsA_copy.layout()); printf("\n");

    printf("\n  === MMA fragment ===\n");
    printf("  tCrA shape: "); print(shape(tCrA)); printf("\n");
    printf("  tCrA layout: "); print(tCrA.layout()); printf("\n");
    printf("  tCrB shape: "); print(shape(tCrB)); printf("\n");
    printf("  tCrD shape: "); print(shape(tCrD)); printf("\n");

    printf("\n  === S2R Copy (make_tiled_copy_A) ===\n");
    printf("  tAsA shape: "); print(shape(tAsA)); printf("\n");
    printf("  tAsA layout: "); print(tAsA.layout()); printf("\n");
    printf("  tCrA_view shape: "); print(shape(tCrA_view)); printf("\n");
    printf("  tCrA_view layout: "); print(tCrA_view.layout()); printf("\n");

    printf("\n  === S2R Copy (make_tiled_copy_B) ===\n");
    printf("  tBsB shape: "); print(shape(tBsB)); printf("\n");
    printf("  tBsB layout: "); print(tBsB.layout()); printf("\n");
    printf("  tCrB_view shape: "); print(shape(tCrB_view)); printf("\n");
    printf("  tCrB_view layout: "); print(tCrB_view.layout()); printf("\n");

    printf("\n  === 循环结构 ===\n");
    printf("  K 外层: K/kTileK = %d/%d = %d 轮\n", K, kTileK, K/kTileK);
    printf("  K 内层: nk = size<2>(tCrA) = %d 轮\n", int(size<2>(tCrA)));
  }

  // 主循环：沿 K 方向
  int num_tile_k = K / kTileK;
  for (int itile = 0; itile < num_tile_k; ++itile) {
    // Global -> Shared
    cute::copy(g2s_tiled_copy_a, tAgA_copy(_, _, _, itile), tAsA_copy(_, _, _));
    cute::copy(g2s_tiled_copy_b, tBgB_copy(_, _, _, itile), tBsB_copy(_, _, _));
    cp_async_fence();
    cp_async_wait<0>();
    __syncthreads();

    // K 内循环
    int nk = size<2>(tCrA);
    for (int ik = 0; ik < nk; ++ik) {
      // Shared -> Register
      cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik), tCrA_view(_, _, ik));
      cute::copy(s2r_tiled_copy_b, tBsB(_, _, ik), tCrB_view(_, _, ik));

      // MMA
      cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
    }
    __syncthreads();
  }

  // 写回（简化版：直接从寄存器写 global）
  auto tDgD = thr_mma.partition_C(gD);
  cute::copy(tCrD, tDgD);
}

// Config
struct GemmConfig {
  using T = half_t;
  static constexpr int kTileM = 128;
  static constexpr int kTileN = 128;
  static constexpr int kTileK = 32;

  // Swizzle SmemLayout
  using SmemLayoutAtom = decltype(composition(
      Swizzle<3, 3, 3>{},
      make_layout(make_shape(Int<8>{}, Int<kTileK>{}),
                  make_stride(Int<kTileK>{}, Int<1>{}))));
  using SmemLayoutA = decltype(
      tile_to_shape(SmemLayoutAtom{}, make_shape(Int<kTileM>{}, Int<kTileK>{})));
  using SmemLayoutB = decltype(
      tile_to_shape(SmemLayoutAtom{}, make_shape(Int<kTileN>{}, Int<kTileK>{})));

  // MMA
  using mma_op = SM80_16x8x16_F16F16F16F16_TN;
  using mma_atom = MMA_Atom<MMA_Traits<mma_op>>;
  using MMA = decltype(make_tiled_mma(
      mma_atom{},
      make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})),
      Tile<Int<32>, Int<32>, Int<16>>{}));

  // G2S Copy: cp.async 128bit
  using g2s_copy_op = SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>;
  using g2s_copy_atom = Copy_Atom<Copy_Traits<g2s_copy_op>, T>;
  using G2SCopyA = decltype(make_tiled_copy(
      g2s_copy_atom{},
      make_layout(make_shape(Int<32>{}, Int<4>{}), make_stride(Int<4>{}, Int<1>{})),
      make_layout(make_shape(Int<1>{}, Int<8>{}))));
  using G2SCopyB = G2SCopyA;

  // S2R Copy: ldmatrix
  using s2r_copy_atom = Copy_Atom<Copy_Traits<SM75_U32x4_LDSM_N>, T>;
  using S2RCopyAtomA = s2r_copy_atom;
  using S2RCopyAtomB = s2r_copy_atom;

  static constexpr int kThreadNum = size(MMA{});
  static constexpr int kShmSize =
      (cosize(SmemLayoutA{}) + cosize(SmemLayoutB{})) * sizeof(T);
};

1. 整体数据流

07: Global ──直接──> Register ──mma──> Register ──直接──> Global
    (慢：global 延迟 ~400 cycles)

08: Global ──cp.async──> Shared ──ldmatrix──> Register ──mma──> Register → Global
    (快：smem 延迟 ~20 cycles)

2. GemmConfig 逐项解析

2.1 基本参数

using T = half_t;                    // fp16
static constexpr int kTileM = 128;   // 每个 block 处理 M=128 行
static constexpr int kTileN = 128;   // 每个 block 处理 N=128 列
static constexpr int kTileK = 32;    // 每次从 global 搬 K=32 列到 smem

跟 07 对比，kTileK 从 16 变成了 32。因为有了 smem，一次可以搬更大的 K 块，减少搬运次数。

2.2 SmemLayout（Swizzle）

using SmemLayoutAtom = decltype(composition(
    Swizzle<3, 3, 3>{},
    make_layout(make_shape(Int<8>{}, Int<kTileK>{}),         // (8, 32)
                make_stride(Int<kTileK>{}, Int<1>{}))));      // row-major

using SmemLayoutA = decltype(
    tile_to_shape(SmemLayoutAtom{}, make_shape(Int<kTileM>{}, Int<kTileK>{})));
using SmemLayoutB = decltype(
    tile_to_shape(SmemLayoutAtom{}, make_shape(Int<kTileN>{}, Int<kTileK>{})));

详见第 3 节 Swizzle 详解。

smem 总大小：

1
2
3

SmemLayoutA: 128 × 32 = 4096 个 half_t = 8192 bytes
SmemLayoutB: 128 × 32 = 4096 个 half_t = 8192 bytes
总计: 16384 bytes = 16 KB

2.3 MMA（跟 07 完全一样）

using mma_op = SM80_16x8x16_F16F16F16F16_TN;
using mma_atom = MMA_Atom<MMA_Traits<mma_op>>;
using MMA = decltype(make_tiled_mma(
    mma_atom{},                                                    // 16×8×16 atom
    make_layout(make_shape(Int<2>{}, Int<2>{}, Int<1>{})),         // EU Repeat
    Tile<Int<32>, Int<32>, Int<16>>{}));                           // P Tile

Atom:       16×8×16,   32 线程
EU (2,2,1): 32×16×16, 128 线程
P <32,32,16>: 不变（P Tile 只做寄存器级 A 复用）

TiledMMA 覆盖: 32(M) × 32(N) × 16(K)
实际 mma.sync 一次: 32×16（P Tile 不增大硬件覆盖面积）

2.4 G2S Copy（Global → Shared Memory）

using g2s_copy_op = SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>;   // cp.async 128bit
using g2s_copy_atom = Copy_Atom<Copy_Traits<g2s_copy_op>, T>;     // 元素类型 half_t

using G2SCopyA = decltype(make_tiled_copy(
    g2s_copy_atom{},
    make_layout(make_shape(Int<32>{}, Int<4>{}),                   // 线程布局: 32行×4列
                make_stride(Int<4>{}, Int<1>{})),                   // row-major
    make_layout(make_shape(Int<1>{}, Int<8>{}))));                  // 每线程每次搬 1×8
using G2SCopyB = G2SCopyA;

三个参数拆解：

参数	值	含义
CopyAtom	`cp.async 128bit`	一条指令搬 128bit = 8 个 half_t
线程布局	`(32, 4):(4, 1)` row-major	128 线程排成 32 行 × 4 列
值布局	`(1, 8)`	每线程每次搬 1 行 × 8 列

一次覆盖多大：

1	线程布局 × 值布局 = (32×1, 4×8) = (32, 32) = 1024 个 half_t

tile 多大：(128, 32) = 4096

需要几轮：4096 / 1024 = 4 轮

关键：cp.async 是异步的，不经过寄存器，直接 global → smem。

2.5 S2R Copy（Shared Memory → Register）

1
2
3

using s2r_copy_atom = Copy_Atom<Copy_Traits<SM75_U32x4_LDSM_N>, T>;
using S2RCopyAtomA = s2r_copy_atom;
using S2RCopyAtomB = s2r_copy_atom;

SM75_U32x4_LDSM_N = ldmatrix 指令：

ldmatrix.sync.aligned.x4.m8n8.shared.b16

一条指令：
  - warp 内 32 个线程协作
  - 从 smem 加载 4 个 8×8 矩阵（每个元素 16bit）
  - 加载到寄存器，直接满足 mma.sync 的寄存器布局

ldmatrix 的价值：

1 2	不用 ldmatrix：smem → 寄存器 → 还需要 shuffle 才能满足 mma.sync 布局用 ldmatrix： smem → 寄存器（直接就是 mma.sync 需要的布局），省掉 shuffle

2.6 完整配置总结

┌─────────────────────────────────────────────────────────┐
│                    GemmConfig 总览                        │
├─────────────────────────────────────────────────────────┤
│ Tile:      128(M) × 128(N) × 32(K)                     │
│ 线程数:    128                                           │
│ Smem:      16 KB (A: 8KB + B: 8KB, with Swizzle)       │
│                                                         │
│ G2S:  cp.async 128bit, 128 线程, 4 轮搬完一个 128×32   │
│ S2R:  ldmatrix x4, 与 MMA 布局自动匹配                   │
│ MMA:  SM80_16x8x16, EU(2,2,1), P<32,32,16>            │
│                                                         │
│ K 循环:                                                  │
│   外层: K/kTileK = 128/32 = 4 轮 (每轮搬 G→S)          │
│   内层: kTileK/atom_K = 32/16 = 2 轮 (每轮搬 S→R + MMA)│
└─────────────────────────────────────────────────────────┘

3. Swizzle 详解

3.1 什么是 Bank Conflict

Shared memory 有 32 个 bank，每个 bank 宽 4 bytes。

地址到 bank 的映射：

bank 编号 = (地址 / 4) % 32

地址 0~3   → bank 0
地址 4~7   → bank 1
地址 8~11  → bank 2
...
地址 124~127 → bank 31
地址 128~131 → bank 0    ← 回绕

同一个 warp 的 32 个线程，如果两个线程访问同一个 bank 的不同地址 → conflict → 串行。

3.2 不加 Swizzle 时的问题

(8, 32) : (32, 1) 的 half_t (2 bytes) smem，一行 32 个 half = 64 bytes：

half_t 是 2 bytes，一个 bank = 4 bytes = 2 个 half_t
一行 32 个 half_t = 64 bytes = 16 bank

row 0:  col 0~1 在 bank 0,  col 2~3 在 bank 1, ... col 30~31 在 bank 15
row 1:  col 0~1 在 bank 16, col 2~3 在 bank 17, ... 到 bank 31, 然后回绕 bank 0...
...
row 4:  col 0~1 又回到 bank 0 ← 跟 row 0 同 bank！

row 0 和 row 4 的 col 0 都在 bank 0。如果 ldmatrix 需要同时读不同行的同一列，就会 conflict。

3.3 Swizzle 的核心思想

用行号来打乱列号，让不同行的同一逻辑列映射到不同物理列，从而落到不同 bank。

操作就是：新列号 = 原列号 XOR 行号

3.4 最简单的例子：4×4 矩阵

原始 (不 Swizzle):
         col0  col1  col2  col3
row 0:   (0,0) (0,1) (0,2) (0,3)
row 1:   (1,0) (1,1) (1,2) (1,3)
row 2:   (2,0) (2,1) (2,2) (2,3)
row 3:   (3,0) (3,1) (3,2) (3,3)

读 col 0 的四行: 都在同一物理列 → 同一 bank → conflict！

做 新列号 = 原列号 XOR 行号：

row 0: 新列号 = col XOR 0 = col       → col0 col1 col2 col3  (不变)
row 1: 新列号 = col XOR 1             → col1 col0 col3 col2  (0↔1, 2↔3)
row 2: 新列号 = col XOR 2             → col2 col3 col0 col1  (0↔2, 1↔3)
row 3: 新列号 = col XOR 3             → col3 col2 col1 col0  (全反转)

Swizzle 后的存储位置：

         物理col0  物理col1  物理col2  物理col3
row 0:   (0,0)     (0,1)     (0,2)     (0,3)       ← 原样
row 1:   (1,1)     (1,0)     (1,3)     (1,2)       ← 0↔1 交换
row 2:   (2,2)     (2,3)     (2,0)     (2,1)       ← 0↔2 交换
row 3:   (3,3)     (3,2)     (3,1)     (3,0)       ← 全反转

现在读逻辑 col 0 的四行：

row 0 的逻辑 col 0 在 → 物理 col 0  → bank 0
row 1 的逻辑 col 0 在 → 物理 col 1  → bank 1
row 2 的逻辑 col 0 在 → 物理 col 2  → bank 2
row 3 的逻辑 col 0 在 → 物理 col 3  → bank 3

四个不同 bank → 无 conflict！

3.5 Swizzle<B, M, S> 三个参数的含义

上面的简单 XOR 是"逐列"做的。但实际 smem 中：

一个 bank = 4 bytes，而 half_t = 2 bytes，连续 2 个 half_t 共享一个 bank
ldmatrix 需要连续的一组元素保持连续，不能逐个打乱

所以 Swizzle 不是逐个元素 XOR，而是按组 XOR。

Swizzle<B, M, S> 定义了怎么分组：

把 offset 的二进制位分成三段：

offset bit:  [...] [B 段] [S 段] [M 段]
                    高位    中间    低位

M 段: 最低 M 位 → 一组内的偏移 → 不动（保持连续，给 ldmatrix 用）
S 段: 中间 S 位 → 组号         → 被 XOR（打乱）
B 段: 高位 B 位 → 行号         → XOR 源

操作: S 段 = S 段 XOR B 段

3.6 Swizzle<3, 3, 3> 对应到 (8, 32) half_t

一行 32 个 half_t，列号 0~31：

M = 3: 低 3 位 = col % 8
       8 个 half_t 一组 = 16 bytes
       组内不动，保持连续（ldmatrix 需要）

S = 3: 中间位 = col / 8
       一行 32/8 = 4 组（组号 0~3）
       被行号 XOR

B = 3: 高 3 位 = 行号（0~7）
       作为 XOR 源

就是第 3.4 节的简单 XOR，但单位从"1 个元素"变成了"8 个元素一组"。

用组号来画：

一行分 4 组: 组0(col0~7) 组1(col8~15) 组2(col16~23) 组3(col24~31)

不加 Swizzle:
row 0: 组0 组1 组2 组3
row 1: 组0 组1 组2 组3     ← 同列同 bank，conflict
row 2: 组0 组1 组2 组3
row 3: 组0 组1 组2 组3

加 Swizzle (组号 XOR 行号):
row 0: 0^0=组0  1^0=组1  2^0=组2  3^0=组3   (不变)
row 1: 0^1=组1  1^1=组0  2^1=组3  3^1=组2   (0↔1, 2↔3 交换)
row 2: 0^2=组2  1^2=组3  2^2=组0  3^2=组1   (0↔2, 1↔3 交换)
row 3: 0^3=组3  1^3=组2  2^3=组1  3^3=组0   (全反转)
row 4: 回绕，同 row 0
...

3.7 验证 bank conflict 消除

一组 = 8 个 half_t = 16 bytes = 4 个 bank

组0 占 bank 0~3
组1 占 bank 4~7
组2 占 bank 8~11
组3 占 bank 12~15

Swizzle 后，读逻辑 col 0 的 row 0~3：
  row 0 逻辑 col 0 在组 0 → bank 0
  row 1 逻辑 col 0 在组 1 → bank 4    ← 不同 bank
  row 2 逻辑 col 0 在组 2 → bank 8    ← 不同 bank
  row 3 逻辑 col 0 在组 3 → bank 12   ← 不同 bank

四行访问四个不同的 bank 组 → 无 conflict ✓

3.8 M 段为什么不动

M = 3 → 低 3 位不参与 XOR → 一组内的 8 个元素保持物理连续

原因：ldmatrix 一次读 128 bit = 8 个 half_t
     这 8 个必须在 smem 中连续存放，否则 ldmatrix 一条指令读不完
     如果把这 8 个也打乱 → ldmatrix 不能用 → 性能下降

3.9 tile_to_shape 扩展

1 2	using SmemLayoutA = decltype( tile_to_shape(SmemLayoutAtom{}, make_shape(Int<kTileM>{}, Int<kTileK>{})));

SmemLayoutAtom 是 (8, 32) with Swizzle，目标 (128, 32)：

在 M 方向重复 128/8 = 16 次：

  行  0~7:    atom 0 (with Swizzle)
  行  8~15:   atom 1 (with Swizzle, 同样的 XOR 模式)
  行 16~23:   atom 2
  ...
  行 120~127: atom 15

每个 8 行块内部都有 Swizzle 消除 bank conflict

3.10 Swizzle 总结

Swizzle 就是：按组对列号做 XOR

  原始:  每行的列排列相同 → 同列同 bank → conflict
  XOR:   不同行的列排列不同 → 同列不同 bank → 无 conflict

Swizzle<B, M, S>:
  M = 一组多大（低 M 位不动，保持连续，给 ldmatrix 用）
  S = 一行多少组（中间 S 位，被 XOR 打乱）
  B = 多少行参与（高 B 位，作为 XOR 源）

Swizzle<3, 3, 3>:
  M=3 → 8 个 half_t 一组（16 bytes），组内不动
  S=3 → 一行最多 8 组（实际只有 4 组，因为 32/8=4）
  B=3 → 8 行为一个 XOR 周期

核心一句话：不同行的列组做 XOR 交换，让同列不同行散布到不同 bank。

4. cp.async 的内部循环

4.1 问题：cp.async 是自己内部循环 4 轮吗？

是的。 cute::copy 内部自动遍历 CPY_M 维度。

从打印结果看 G2S copy 的 shape：

1 2	tAgA_copy shape: ((_8,_1), _4, _1, 4) CPY CPY_M CPY_K num_tile_k

维度	值	含义
CPY = `(8,1)` = 8	每线程每次搬 8 个 half_t (128bit)	cp.async 一条指令
CPY_M = 4	M 方向循环 4 次	128 行 / 32 行线程 = 4 轮
CPY_K = 1	K 方向 1 次	32 列 / (4 列线程 × 8 值) = 1
num_tile_k = 4	K 外层	128/32 = 4

一次 cute::copy(g2s_tiled_copy_a, src, dst) 的内部：

// cute::copy 内部等价于：
for (int cpy_m = 0; cpy_m < 4; ++cpy_m) {      // 4 轮
    for (int cpy_k = 0; cpy_k < 1; ++cpy_k) {   // 1 轮
        // 128 个线程同时发出 cp.async，每个搬 128bit
        cp.async(smem_addr, global_addr);
    }
}
// 一个线程发了 4 条 cp.async 指令

画出来：

128×32 的 tile：

线程布局 (32,4) 每线程搬 (1,8)  →  一次覆盖 (32, 32)

  行  0~31   ← CPY_M=0  (128 线程同时搬，每个搬 8 个 half_t)
  行 32~63   ← CPY_M=1
  行 64~95   ← CPY_M=2
  行 96~127  ← CPY_M=3

stride 验证：
  tAgA_copy stride: ((_1,_0), 4096, _0, _32)
  CPY_M 的 stride = 4096 = 32行 × 128(全局K stride)  ✓
  num_tile_k 的 stride = 32 = kTileK  ✓

一次 cute::copy 调用 = 128 线程 × 4 轮 × 128bit/次 = 搬完整个 128×32 tile。

不需要手动写 for 循环，cute::copy 自动遍历 CPY_M/CPY_K 维度。

5. make_tiled_copy_A 与 retile_D

5.1 核心问题：为什么 G→S 不需要，S→R 需要

根本原因：看目的地（Destination）是谁在用。

G→S：目的地是 shared memory（公共的）

1
2
3

Global ──cp.async──> Shared Memory
                     ↑
                     公共资源，谁写都行，只要数据到位就行

smem 是 block 内所有线程共享的。线程 0 写 smem[100] 和线程 5 写 smem[100]，效果一样。所以 G2S 的线程布局可以随便选，怎么高效怎么来：

// G2S: 线程布局自己选，怎么高效怎么来
using G2SCopyA = decltype(make_tiled_copy(
    g2s_copy_atom{},
    make_layout(make_shape(Int<32>{}, Int<4>{}), ...),   // ← 自己选
    make_layout(make_shape(Int<1>{}, Int<8>{}))));       // ← 自己选

S→R：目的地是寄存器（私有的）

Shared Memory ──ldmatrix──> Register
                            ↑
                            私有！只有这个线程能用
                            而且 mma.sync 要求特定布局

寄存器是线程私有的。线程 0 的 reg[0] 只有线程 0 能用，别人写不了也读不了。

所以 S2R 有两个硬约束。

5.2 约束 1：线程必须对得上

mma.sync 硬件规定死了每个线程负责哪些矩阵元素：

mma.sync 硬件规定：
  线程 0 → 必须提供 A[0][0], A[0][1], A[8][0], A[8][1], A[0][8], A[0][9], A[8][8], A[8][9]
  线程 1 → 必须提供 A[1][0], A[1][1], A[9][0], A[9][1], A[1][8], A[1][9], A[9][8], A[9][9]
  ...
  这个映射是芯片设计时焊死的，不能改。

如果 S2R 的线程分配选错了：

                    正确的分配                    错误的分配
                    ──────────                    ──────────
smem A[0][0] ──→ 线程 0 的 reg[0]   smem A[0][0] ──→ 线程 3 的 reg[0]  ✗
smem A[1][0] ──→ 线程 1 的 reg[0]   smem A[1][0] ──→ 线程 0 的 reg[0]  ✗
smem A[2][0] ──→ 线程 2 的 reg[0]   smem A[2][0] ──→ 线程 1 的 reg[0]  ✗
...

错误分配后，mma.sync 一执行，每个线程拿到的都不是自己该用的数据
→ 计算结果全错
→ 而且无法修复（线程间不能交换寄存器）

make_tiled_copy_A 怎么解决

auto s2r_tiled_copy_a = make_tiled_copy_A(S2RCopyAtomA{}, tiled_mma);
//                                                        ↑
//                                              从 MMA 的 Layout_TV_A 查：
//                                              线程 0 需要 A 的哪些位置
//                                              线程 1 需要 A 的哪些位置
//                                              ...
//                                              然后让 Copy 的线程分配与之一致

保证了：线程 0 从 smem 读出来的，恰好就是 mma.sync 要求线程 0 提供的。

	G2S	S2R
线程布局	你随便选	必须和 MMA 一致
用什么 API	`make_tiled_copy`	`make_tiled_copy_A`
为什么	smem 是公共的	寄存器是私有的

5.3 约束 2：寄存器视角必须对得上

MMA 视角（partition_fragment_A 创建的 tCrA）：

1
2
3

tCrA shape: ((_2,_2,_2), _4, _2)
             MMA内8个值    M重复  K内循环
stride:     ((_1,_2,_4), _16, _8)

这是 mma.sync 要求的寄存器布局。8 个值按 (2,2,2) 嵌套排列。

Copy 视角（ldmatrix 需要的布局）：

1
2
3

tCrA_view shape: ((_8,_1), _4, _2)
                  一次8个    M重复  K内循环
stride:          ((_1,_0), _16, _8)

实际上物理数据排列是一样的——ldmatrix 和 mma.sync 是配套设计的。真正的问题在 CuTe 的类型系统层面：

cute::copy 函数：需要 dst 的第 0 维 shape 是 (8,1)     ← "我一次搬 8 个连续值"
cute::gemm 函数：需要 src 的第 0 维 shape 是 (2,2,2)   ← "我按三层嵌套访问"

如果 shape 不匹配 → 编译报错

如果不做 retile_D，直接传 tCrA：

// 假设不做 retile，直接传 tCrA
cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik), tCrA(_, _, ik));
//                                            ↑
//                               tCrA 的第 0 维 shape 是 (2,2,2)
//                               但 s2r_tiled_copy_a 期望第 0 维是 (8,1)
//                               shape 不匹配 → 编译报错！

retile_D 做的事

1	auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA);

retile_D 不改变数据，不分配新寄存器，只改变看数据的方式（reshaping）。

tCrA:      reg[0..7], 索引方式 (2,2,2) stride (1,2,4)
           访问 reg[5] = tCrA(1,0,1) = 1×1 + 0×2 + 1×4 = 5  ✓

tCrA_view: reg[0..7], 索引方式 (8,1) stride (1,0)
           访问 reg[5] = tCrA_view(5,0) = 5×1 + 0×0 = 5     ✓

同一个 reg[5]，两种方式都能访问到，只是坐标写法不同。

retile_D 的效果：

  tCrA       shape: ((_2,_2,_2), _4, _2)  ← gemm 用这个（MMA 视角）
                      ↕ 同一片寄存器
  tCrA_view  shape: ((_8,_1),    _4, _2)  ← copy 用这个（Copy 视角）

5.4 约束总结

约束 1（线程对齐）：
  mma.sync 硬件规定了每个线程负责哪些矩阵元素
  S→R 的 Copy 必须让每个线程从 smem 读到自己该用的那份
  → make_tiled_copy_A 从 MMA 的 Layout_TV 自动对齐

约束 2（视角对齐）：
  cute::copy 和 cute::gemm 对同一片寄存器的索引约定（shape）不同
  copy 期望 (8,1)，gemm 期望 (2,2,2)
  → retile_D 把 shape 从 MMA 视角转成 Copy 视角
  → 不搬数据，只换索引方式，编译期完成

一句话：约束 1 是"谁读什么"的问题，约束 2 是"怎么看同一片数据"的问题。

5.5 G→S 为什么不需要

G→S:  Destination 是 smem (公共) → 线程布局随便选 → make_tiled_copy
                                 → 无视角问题     → 不需要 retile

S→R:  Destination 是 reg (私有)  → 线程布局必须匹配 MMA → make_tiled_copy_A
                                 → 两种视角切换        → retile_D

核心：smem 谁写都行，寄存器只有自己能用。

5.6 类比

G→S 就像把货搬进仓库：
  仓库(smem)是公共的，谁搬都行，放到对的位置就好
  搬运工(线程)怎么分配无所谓，效率优先
  不需要关心"谁搬的"

S→R 就像把工具发给工人：
  每个工人(线程)只能用自己手里的工具(寄存器)
  必须把正确的工具发到正确的工人手里     ← make_tiled_copy_A
  而且工具要按工人习惯的方式摆好         ← retile_D

6. 两层 K 循环

6.1 为什么有两层循环

核心在于：一次搬到 smem 的 K 维度（32）比一次 MMA 能消耗的 K 维度（16）大。

kTileK = 32    ← 每次从 global 搬 32 列到 smem
MMA K  = 16    ← 一次 mma.sync 只能算 K=16

所以 smem 里的 32 列要分 32/16 = 2 次 MMA 才能用完

6.2 对比 07（只有一层循环）

07 的 kTileK = 16，刚好等于 MMA 的 K=16：

// 07: kTileK=16, MMA K=16, 一次就用完
for (int ik = 0; ik < 4; ++ik) {     // 64/16 = 4 轮
    copy(global → register);          // 搬 16 列
    gemm();                            // 算 16 列，刚好用完
}

只需要一层循环。

6.3 为什么 kTileK 要大于 MMA K

因为搬运比计算贵。cp.async 从 global 搬到 smem 有很高的开销（发起异步请求、fence、wait、syncthreads），所以希望一次多搬点，减少搬运次数：

1 2	kTileK=16: K=128 需要 128/16 = 8 次 G→S ← 8 次 fence+wait+sync kTileK=32: K=128 需要 128/32 = 4 次 G→S ← 4 次，减半

但 MMA 硬件一次只能算 K=16，所以搬进来的 32 列要分 2 次用。

6.4 两层循环的结构

int num_tile_k = K / kTileK;                          // 外层: 128/32 = 4 轮
for (int itile = 0; itile < num_tile_k; ++itile) {
    // Global -> Shared: 搬 128×32 到 smem
    cute::copy(g2s_tiled_copy_a, tAgA_copy(_, _, _, itile), tAsA_copy(_, _, _));
    cute::copy(g2s_tiled_copy_b, tBgB_copy(_, _, _, itile), tBsB_copy(_, _, _));
    cp_async_fence();
    cp_async_wait<0>();
    __syncthreads();

    int nk = size<2>(tCrA);                            // 内层: 32/16 = 2 轮
    for (int ik = 0; ik < nk; ++ik) {
        // Shared -> Register: ldmatrix 搬 16 列
        cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik), tCrA_view(_, _, ik));
        cute::copy(s2r_tiled_copy_b, tBsB(_, _, ik), tCrB_view(_, _, ik));
        // MMA: 32 次 mma.sync
        cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
    }
    __syncthreads();
}

时间线：

K 总共 128 列：

  [G→S 32列] [S→R 16列][MMA] [S→R 16列][MMA]
  [G→S 32列] [S→R 16列][MMA] [S→R 16列][MMA]
  [G→S 32列] [S→R 16列][MMA] [S→R 16列][MMA]
  [G→S 32列] [S→R 16列][MMA] [S→R 16列][MMA]

6.5 从打印结果验证

tCrA shape: ((_2,_2,_2), _4, _2)
                               ↑
                           nk=2  ← 这就是内层循环次数

循环结构:
  K 外层: K/kTileK = 128/32 = 4 轮
  K 内层: nk = size<2>(tCrA) = 2 轮

tCrA 的第 3 维 _2 就是内层循环的 2 轮，对应 smem 里 32 列分成两个 16 列。

6.6 kTileK 的权衡

kTileK=16: 内层=1, 外层=8, G→S 次数多，但 smem 占用小
kTileK=32: 内层=2, 外层=4, G→S 次数减半，smem 16KB
kTileK=64: 内层=4, 外层=2, G→S 更少，但 smem 32KB（可能影响 occupancy）

kTileK 的选择是搬运次数 vs smem 占用的权衡。

7. Kernel 完整流程

7.1 kernel 中的使用流程

// 第 1 步：G→S（不关心谁写，不关心布局）
cute::copy(g2s_tiled_copy_a, tAgA_copy(_, _, _, itile), tAsA_copy);
// 任意线程布局 → cp.async → smem
// make_tiled_copy 自己选线程布局，怎么快怎么来

// 第 2 步：__syncthreads()
// 确保 smem 数据完整

// 第 3 步：S→R（必须对得上）
cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik), tCrA_view(_, _, ik));
// make_tiled_copy_A → 保证线程 0 读 mma 要求线程 0 用的数据
// retile_D → 让 copy 能正确索引寄存器
// ldmatrix: smem → 寄存器 (Copy 视角写入)

// 第 4 步：MMA
cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
// mma.sync: 从寄存器读 (MMA 视角读取)
// 同一片寄存器，换个视角就能用

7.2 打印结果汇总

=== G2S Copy ===
tAgA_copy shape: ((_8,_1), _4, _1, 4)     ← 每线程8值, 4轮M, 1轮K, 4个K tile
tAsA_copy shape: ((_8,_1), _4, _1)         ← smem 侧，无 K tile 维度

=== MMA fragment ===
tCrA shape: ((_2,_2,_2), _4, _2)           ← MMA 视角
tCrB shape: ((_2,_2), (_2,_4), _2)
tCrD shape: ((_2,_2), _4, _8)

=== S2R Copy ===
tAsA shape:      ((_8,_1), _4, _2)         ← smem source
tCrA_view shape: ((_8,_1), _4, _2)         ← 寄存器 Copy 视角（与 tAsA 一一对应）

=== 循环结构 ===
K 外层: 128/32 = 4 轮
K 内层: 2 轮

8. 与 07/09 的对比

	07（Simple GEMM）	08（Smem GEMM）	09（Multi-Stage）
数据搬运	Global→Reg	Global→Smem→Reg	Global→Smem(multi)→Reg
smem	无	有，单缓冲	有，多 stage 环形缓冲
Swizzle	不需要	有	有
K 循环	1 层	2 层	2 层 + 流水线
cp.async	不用	用，但同步等待	用，异步重叠
make_tiled_copy_A	不需要	需要	需要
retile_D	不需要	需要	需要
性能	很慢	中等	最快
学习目标	partition + gemm	G2S + S2R + Swizzle	流水线重叠

9. API 总结

API	作用
`make_tiled_copy(atom, thr_layout, val_layout)`	创建 Copy，线程布局自己选（用于 G→S）
`make_tiled_copy_A(atom, tiled_mma)`	创建 Copy，线程布局自动从 MMA 提取（用于 S→R of A）
`make_tiled_copy_B(atom, tiled_mma)`	同上，用于 S→R of B
`thr_copy.retile_D(tCrA)`	把 MMA 视角的 fragment 转成 Copy 视角（zero-copy reshape）
`thr_copy.partition_S(smem)`	切分 smem Source
`thr_copy.partition_D(smem)`	切分 smem Destination
`cp_async_fence()`	标记一批 cp.async 的边界
`cp_async_wait<N>()`	等待直到最多还剩 N 批未完成
`composition(Swizzle, Layout)`	在 Layout 上叠加 Swizzle 消除 bank conflict
`tile_to_shape(atom, shape)`	把小 layout atom 重复扩展到目标 shape

cute（8）smem_gemm