cute（11）hopper之TMA_copy

前置阅读

• reed博士：(4 封私信 / 2 条消息) cute 之 Hopper TMA - 知乎

• 竹佬：(4 封私信 / 2 条消息) 写给大家看的 CuTe 教程：TMA Copy - 知乎

11 TMA Copy

对应代码：11_tma_copy.cu 需要 GPU（SM90 / Hopper）。

核心概念

TMA（Tensor Memory Accelerator） 是 Hopper 架构的硬件加速搬运引擎：

• 单线程即可发起整个 tile 的 Global → Shared Memory 搬运
• 硬件自动处理地址计算、Swizzle、多维索引
• 搬运完成后自动通过 mbarrier 通知
• 与 cp.async（09 课）的根本区别：cp.async 每线程搬一小块，TMA 一个线程搬整个 tile

---

代码实现

/**
 * 第 11 课：TMA Copy —— Tensor Memory Accelerator
 *
 * 学习目标：
 *   - 用 TMA 搬运一个 2D tile 从 Global → Shared Memory
 *   - Host 端：make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout) 构造 TMA descriptor
 *   - arithTuple 坐标系：get_tma_tensor → local_tile → partition_S / partition_D
 *   - Kernel 端：单线程发起 TMA copy，mbarrier 等待完成
 *   - 与 cp.async（09 课）的对比
 *
 * TMA 的核心优势：
 *   - 单线程即可发起整个 tile 的搬运（而不是每线程搬一部分）
 *   - 硬件自动处理地址计算和 swizzle
 *   - 搬运完成后自动 arrive 到 mbarrier
 *
 * 关键 API：
 *   Host: auto tma_load = make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout);
 *   Device: copy(tma_load.with(mbar), thr.partition_S(coord_tile), thr.partition_D(smem_tensor));
 */

#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>  // GMMA::Layout_MN_SW128_Atom
#include <cute/arch/copy_sm90_desc.hpp>
#include <cstdio>
#include <cstdlib>

using namespace cute;

// ============================================================
// TMA Copy Kernel
// ============================================================
// __grid_constant__ 让 TMA descriptor 以 kernel 参数形式传入（而不是通过 smem 或 gmem）
// 注意：Copy_Traits 包含 TmaDescriptor（128字节对齐），必须用 __grid_constant__
template <class TmaLoad, class SmemLayout, class TileShape>
__global__ void tma_copy_kernel(
    __grid_constant__ const TmaLoad tma_load,
    half_t const* gmem_ptr,
    half_t*       result_ptr,
    int M, int N) {

  using T = half_t;
  constexpr int kTileM = get<0>(TileShape{});
  constexpr int kTileN = get<1>(TileShape{});

  // Step 1: 在 smem 中分配 buffer 和 mbarrier
  extern __shared__ char smem_buf[];
  T* smem_data = reinterpret_cast<T*>(smem_buf);
  // mbarrier 需要 8 字节对齐，放在 smem_data 后面
  constexpr int smem_data_size = cosize(SmemLayout{});
  uint64_t* mbar_ptr = reinterpret_cast<uint64_t*>(
      smem_buf + sizeof(T) * smem_data_size);

  auto sA = make_tensor(make_smem_ptr(smem_data), SmemLayout{});  // (kTileM, kTileN)

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

  // Step 2: 构造 TMA 坐标 tensor
  //   get_tma_tensor 返回一个特殊的 "坐标tensor"
  //   它的 shape 和 gmem_tensor 相同，但元素是坐标（arithTuple）
  //   用 local_tile 切出当前 block 需要的 tile
  auto coord_tensor = tma_load.get_tma_tensor(make_shape(M, N));

  // local_tile：从坐标 tensor 中切出 (kTileM, kTileN) 大小的 tile
  auto coord_tile = local_tile(
      coord_tensor,
      make_tile(Int<kTileM>{}, Int<kTileN>{}),
      make_coord(by, bx));  // (tileM, tileN)

  // Step 3: 获取线程的 partition
  //   TMA 只需 1 个线程发起搬运，所以 get_slice(0)
  auto tma_thr = tma_load.get_slice(0);
  auto tSgS = tma_thr.partition_S(coord_tile);  // source partition (坐标)
  auto tSsS = tma_thr.partition_D(sA);          // destination partition (smem)

  // Step 4: 初始化 mbarrier + 设置 transaction bytes + 发起 TMA copy
  if (tid == 0) {
    cute::initialize_barrier(*mbar_ptr, 1);  // 1 个到达者

    // TMA 搬运的字节数 = tile 大小 × 元素字节数
    uint32_t tx_bytes = kTileM * kTileN * sizeof(T);
    cute::set_barrier_transaction_bytes(*mbar_ptr, tx_bytes);

    // 发起 TMA copy！
    //   tma_load.with(*mbar_ptr) 绑定 mbarrier
    //   copy 只需 thread 0 调用，TMA 引擎会自动搬运整个 tile
    copy(tma_load.with(*mbar_ptr), tSgS, tSsS);

    if (by == 0 && bx == 0) {
      printf("[TMA Copy] Block(%d,%d) 线程 0 发起 TMA 搬运\n", bx, by);
      printf("  coord_tile shape: ");
      print(shape(coord_tile));
      printf("\n");
      printf("  tSgS shape: ");
      print(shape(tSgS));
      printf("\n");
      printf("  tSsS shape: ");
      print(shape(tSsS));
      printf("\n");
      printf("  tx_bytes = %u\n", tx_bytes);
    }
  }

  // Step 5: 所有线程等待 TMA 搬运完成
  __syncthreads();  // 确保 mbarrier 初始化完成
  if (tid == 0) {
    cute::wait_barrier(*mbar_ptr, 0);  // 等待 phase 0 翻转
  }
  __syncthreads();  // 让所有线程看到搬运完成的数据

  // Step 6: 验证 — 将 smem 数据写回 global（用于 host 端校验）
  // 每个线程搬运一部分
  for (int i = tid; i < kTileM * kTileN; i += blockDim.x) {
    int row = i / kTileN;
    int col = i % kTileN;
    int global_row = by * kTileM + row;
    int global_col = bx * kTileN + col;
    // smem 可能有 swizzle，通过 sA 访问以自动处理
    result_ptr[global_row * N + global_col] = sA(row, col);
  }
}

int main() {
  printf("=== 第 11 课：TMA Copy ===\n\n");

  using T = half_t;
  constexpr int M = 256;
  constexpr int N = 128;
  constexpr int kTileM = 128;
  constexpr int kTileN = 64;
  using TileShape = Shape<Int<kTileM>, Int<kTileN>>;

  printf("问题规模: M=%d, N=%d\n", M, N);
  printf("Tile 大小: %d x %d\n", kTileM, kTileN);

  // ============================================================
  // 1. 分配内存 & 初始化数据
  // ============================================================
  T* h_data = (T*)malloc(M * N * sizeof(T));
  T* h_result = (T*)malloc(M * N * sizeof(T));
  T* d_data;
  T* d_result;

  srand(42);
  for (int i = 0; i < M * N; ++i) {
    h_data[i] = T(float(i % 100) * 0.01f);
  }

  cudaMalloc(&d_data, M * N * sizeof(T));
  cudaMalloc(&d_result, M * N * sizeof(T));
  cudaMemcpy(d_data, h_data, M * N * sizeof(T), cudaMemcpyHostToDevice);
  cudaMemset(d_result, 0, M * N * sizeof(T));

  // ============================================================
  // 2. 构造 GMEM Tensor（Host 端）
  // ============================================================
  // Row-major: A(M, N) with stride (N, 1)
  auto gmem_tensor = make_tensor(
      make_gmem_ptr(d_data),
      make_shape(M, N),
      make_stride(N, Int<1>{}));

  printf("\ngmem_tensor:\n");
  printf("  shape:  ");  print(shape(gmem_tensor));  printf("\n");
  printf("  stride: ");  print(stride(gmem_tensor)); printf("\n");

  // ============================================================
  // 3. 构造 SMEM Layout（带 Swizzle）
  // ============================================================
  // TMA 要求 Swizzle 的 M 参数 ≥ 4（对应 16B base alignment）
  // GMMA 提供标准 layout atoms：
  //   Layout_MN_SW128_Atom<T>: stride-1 沿第一维 (M/N major)
  //   Layout_K_SW128_Atom<T>:  stride-1 沿第二维 (K major)
  //
  // 我们的 gmem 是 (M, N) stride (N, 1) → stride-1 沿 N（第二维）
  // smem tile 是 (kTileM, kTileN) → 需要 stride-1 沿 kTileN（第二维）
  // 所以使用 Layout_K_SW128_Atom（K-major，stride-1 沿第二维）
  using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<T>;
  using SmemLayout = decltype(tile_to_shape(
      SmemLayoutAtom{},
      make_shape(Int<kTileM>{}, Int<kTileN>{})));

  printf("\nSmemLayout:\n  ");
  print(SmemLayout{});
  printf("\n");

  // ============================================================
  // 4. 构造 TMA Descriptor（Host 端）
  // ============================================================
  // make_tma_copy 在 Host 端创建 TMA descriptor
  // 它编码了：gmem 地址、形状、stride、smem swizzle 模式
  auto tma_load = make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, SmemLayout{});

  printf("\n构造 TMA descriptor 完成\n");

  // ============================================================
  // 5. 启动 Kernel
  // ============================================================
  dim3 grid(N / kTileN, M / kTileM);  // (2, 2)
  dim3 block(128);  // 128 线程（虽然只有 thread 0 发起 TMA，其余线程用于写回验证）

  // smem 大小：tile 数据 + mbarrier（8 字节对齐）
  int smem_size = sizeof(T) * cosize(SmemLayout{}) + sizeof(uint64_t) + 8;

  printf("Grid: (%d, %d), Block: %d, ShmSize: %d bytes\n",
         grid.x, grid.y, block.x, smem_size);

  tma_copy_kernel<decltype(tma_load), SmemLayout, TileShape>
      <<<grid, block, smem_size>>>(tma_load, d_data, d_result, M, N);
  cudaDeviceSynchronize();

  auto err = cudaGetLastError();
  if (err != cudaSuccess) {
    printf("CUDA error: %s\n", cudaGetErrorString(err));
    return 1;
  }

  // ============================================================
  // 6. 验证结果
  // ============================================================
  cudaMemcpy(h_result, d_result, M * N * sizeof(T), cudaMemcpyDeviceToHost);

  int errors = 0;
  for (int i = 0; i < M * N; ++i) {
    if (float(h_data[i]) != float(h_result[i])) {
      if (errors < 5) {
        printf("  MISMATCH at [%d]: expected=%.4f, got=%.4f\n",
               i, float(h_data[i]), float(h_result[i]));
      }
      errors++;
    }
  }
  printf("\n验证: %s (%d errors)\n", errors == 0 ? "PASS" : "FAIL", errors);

  // ============================================================
  // 7. TMA vs cp.async 对比总结
  // ============================================================
  printf("\n========================================\n");
  printf("TMA vs cp.async 对比\n");
  printf("========================================\n");
  printf("  ┌───────────────────┬─────────────────┬──────────────────┐\n");
  printf("  │     特性          │ cp.async (SM80) │ TMA (SM90)       │\n");
  printf("  ├───────────────────┼─────────────────┼──────────────────┤\n");
  printf("  │ 发起线程数        │ 所有线程        │ 单线程           │\n");
  printf("  │ 地址计算          │ 每线程自己算    │ 硬件自动         │\n");
  printf("  │ Swizzle           │ 需要手动 Layout │ Descriptor 编码  │\n");
  printf("  │ 同步机制          │ cp_async_wait   │ mbarrier         │\n");
  printf("  │ Descriptor        │ 不需要          │ Host 端构造      │\n");
  printf("  │ 支持多维          │ 只能 1D         │ 最多 5D          │\n");
  printf("  └───────────────────┴─────────────────┴──────────────────┘\n");

  printf("\n关键流程：\n");
  printf("  Host:   make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout)\n");
  printf("  Kernel: tma_load.get_tma_tensor(shape) → local_tile → partition\n");
  printf("          thread 0: init_barrier → set_tx_bytes → copy(tma.with(mbar), src, dst)\n");
  printf("          all threads: wait_barrier\n");

  free(h_data);
  free(h_result);
  cudaFree(d_data);
  cudaFree(d_result);

  printf("\n=== 练习完成 ===\n");
  return 0;
}

打印信息如下

[yuyanqi@set-hldy-mlp-codelab-pc75 build]$ ./11_tma_copy 
=== 第 11 课：TMA Copy ===

问题规模: M=256, N=128
Tile 大小: 128 x 64

gmem_tensor:
  shape:  (256,128)
  stride: (128,_1)

SmemLayout:
  Sw<3,4,3> o smem_ptr[16b](unset) o ((_8,_16),(_64,_1)):((_64,_512),(_1,_0))

构造 TMA descriptor 完成
Grid: (2, 2), Block: 128, ShmSize: 16400 bytes
[TMA Copy] Block(0,0) 线程 0 发起 TMA 搬运
  coord_tile shape: (_128,_64)
  tSgS shape: (((_64,_128),_1),_1,_1)
  tSsS shape: (((_64,_128),_1),_1,_1)
  tx_bytes = 16384

验证: PASS (0 errors)

========================================
TMA vs cp.async 对比
========================================
  ┌───────────────────┬─────────────────┬──────────────────┐
  │     特性          │ cp.async (SM80) │ TMA (SM90)       │
  ├───────────────────┼─────────────────┼──────────────────┤
  │ 发起线程数        │ 所有线程        │ 单线程           │
  │ 地址计算          │ 每线程自己算    │ 硬件自动         │
  │ Swizzle           │ 需要手动 Layout │ Descriptor 编码  │
  │ 同步机制          │ cp_async_wait   │ mbarrier         │
  │ Descriptor        │ 不需要          │ Host 端构造      │
  │ 支持多维          │ 只能 1D         │ 最多 5D          │
  └───────────────────┴─────────────────┴──────────────────┘

关键流程：
  Host:   make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout)
  Kernel: tma_load.get_tma_tensor(shape) → local_tile → partition
          thread 0: init_barrier → set_tx_bytes → copy(tma.with(mbar), src, dst)
          all threads: wait_barrier

=== 练习完成 ===

1. TMA 的工作流程

Host 端（一次性准备）:
  make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout)
       ↓
  TMA Descriptor（128 字节，编码 gmem 信息 + smem swizzle 模式）
       ↓
  通过 kernel 参数传入 (__grid_constant__)

Kernel 端（每个 block 执行）:
  1. get_tma_tensor(shape) → 坐标 tensor（arithTuple）
  2. local_tile → 切出当前 block 的 tile 坐标
  3. get_slice(0) → TMA 只需 1 个线程
  4. partition_S / partition_D → source（坐标）和 dest（smem）
  5. thread 0: init_mbarrier → set_tx_bytes → copy(tma.with(mbar), src, dst)
  6. 所有线程: wait_barrier → 数据就绪

---

2. Host 端：构造 TMA Descriptor

2.1 GMEM Tensor

auto gmem_tensor = make_tensor(
    make_gmem_ptr(d_data),
    make_shape(M, N),         // (256, 128)
    make_stride(N, Int<1>{})); // Row-major: stride (128, 1)

2.2 SMEM Layout（带 Swizzle）

using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<half_t>;
using SmemLayout = decltype(tile_to_shape(
    SmemLayoutAtom{}, make_shape(Int<128>{}, Int<64>{})));

为什么用 Layout_K_SW128_Atom 而不是 Layout_MN_SW128_Atom？

属性	Layout_K	Layout_MN
stride-1 方向	第二维 (K/N)	第一维 (M/N)
适用场景	gmem row-major: stride-1 沿 N	gmem col-major: stride-1 沿 M

我们的 gmem 是 (M, N) stride (N, 1) → stride-1 沿 N（第二维）→ smem 也应 stride-1 沿第二维 → Layout_K。

TMA 对 Swizzle 的要求：

Swizzle<B, M, S>  中  M 必须 ≥ 4

Layout_K_SW128_Atom 的 swizzle: Swizzle<3, 4, 3>  → M=4 ✓
Swizzle<3, 3, 3>（09 课用的）→ M=3 ✗ → TMA 不兼容！

TMA 的 swizzle base 需要 16B 对齐（M=4）、32B（M=5）或 64B（M=6），对应：

• Layout_K_INTER_Atom: Swizzle<0,4,3> — 无 swizzle
• Layout_K_SW32_Atom: Swizzle<1,4,3> — 32B swizzle
• Layout_K_SW64_Atom: Swizzle<2,4,3> — 64B swizzle
• Layout_K_SW128_Atom: Swizzle<3,4,3> — 128B swizzle

2.3 make_tma_copy

auto tma_load = make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, SmemLayout{});

这一步在 Host 端执行，创建一个包含 TmaDescriptor（128 字节对齐结构体）的 Copy_Traits 对象。Descriptor 编码了：

• gmem 的基地址、shape、stride
• smem 的 swizzle 模式
• 搬运的元素类型和大小

---

3. Kernel 端：发起 TMA Copy

3.1 grid_constant

__global__ void tma_copy_kernel(
    __grid_constant__ const TmaLoad tma_load, ...)

__grid_constant__ 是 Hopper 引入的修饰符：

• 让 kernel 参数存在于 constant memory 中（而非被拷贝到每个线程的 local memory）
• TMA descriptor 是 128 字节的大结构体，必须用此修饰
• 所有线程共享同一份，减少 register 压力

3.2 坐标 Tensor 和 Partition

// 1. 获取坐标 tensor（arithTuple 格式）
auto coord_tensor = tma_load.get_tma_tensor(make_shape(M, N));
// coord_tensor 的 shape 和 gmem 一样，但元素是坐标而非数据

// 2. 切出当前 block 的 tile
auto coord_tile = local_tile(coord_tensor,
    make_tile(Int<128>{}, Int<64>{}),
    make_coord(by, bx));

// 3. TMA 只需 1 个线程发起搬运
auto tma_thr = tma_load.get_slice(0);  // slice 0（只有 1 个）
auto tSgS = tma_thr.partition_S(coord_tile);  // 坐标 partition
auto tSsS = tma_thr.partition_D(sA);          // smem partition

为什么 TMA 用"坐标 tensor"而不是数据 tensor？

cp.async 需要每线程自己算地址→搬一小块。TMA 不同：

• 线程不直接访问 gmem 地址
• 而是告诉 TMA 引擎"请搬坐标 (row, col) 处的 tile"
• 地址计算由硬件完成

3.3 发起搬运

if (tid == 0) {
    cute::initialize_barrier(*mbar_ptr, 1);
    cute::set_barrier_transaction_bytes(*mbar_ptr, kTileM * kTileN * sizeof(T));
    copy(tma_load.with(*mbar_ptr), tSgS, tSsS);
}

关键步骤：

1. initialize_barrier: 设 arrive_count=1（set_barrier_transaction_bytes 内部会 arrive）
2. set_barrier_transaction_bytes: 设置期望搬运量 + 隐式 arrive
3. copy(tma.with(mbar)): .with(mbar) 绑定 mbarrier，TMA 搬运完成后自动 complete_tx
4. 所有线程 wait_barrier(mbar, 0) → 等待搬运完成

时间线：
  thread 0: init_mbar → set_tx(16KB) → copy(TMA) → TMA 引擎异步搬运
                                                          ↓
  所有线程:                                   wait_barrier ←── TMA 完成, phase 翻转
                                                          ↓
                                              数据在 smem 中就绪！

3.4 同步代码详解：两个 __syncthreads 和 tid==0 wait

__syncthreads();           // ① 确保 mbarrier 初始化完成
if (tid == 0) {
    cute::wait_barrier(*mbar_ptr, 0);  // ② 等待 phase 0 翻转
}
__syncthreads();           // ③ 让所有线程看到搬运完成的数据

为什么是 tid==0 做 wait_barrier，而不是所有线程？

wait_barrier 内部是 spin-wait 忙等循环：

LAB_WAIT:
  mbarrier.try_wait.parity.shared::cta.b64 P1, [mbar], phase, ticks;
  @P1 bra DONE;
  bra LAB_WAIT;    // 不断轮询直到 phase 翻转
DONE:

128 个线程全部 spin-wait → 128 个线程同时轮询同一个 smem 地址 → 浪费算力、增加 smem bank 争用。只让 1 个线程探测就够了，然后用 __syncthreads 广播"数据已就绪"。

三个同步点各自的作用：

tid==0:  init_barrier → set_tx_bytes → copy(TMA)
其他线程: 不知道 mbarrier 是否已初始化

① __syncthreads
   作用：确保 tid==0 对 smem 的写入（mbarrier 初始化）对所有线程可见
   虽然只有 tid==0 做后续的 wait，但这是 smem 内存可见性的保证

tid==0:  wait_barrier(phase=0) — 独自 spin-wait 等 TMA 搬运完成
其他线程: 在 ③ 的 __syncthreads 处等待

② TMA 完成 → full_barrier phase 翻转 → tid==0 的 wait 通过
   此时 smem 中的数据对 tid==0 可见

③ __syncthreads
   作用：tid==0 已知数据就绪，通过 __syncthreads 让所有线程看到 TMA 写入的数据
   如果没有 ③：其他线程可能读到 TMA 还没写完的旧数据

能不能让所有线程都 wait_barrier？

可以，但没必要。wait_barrier 是纯轮询，N 个线程同时轮询不会比 1 个更快通过——翻转时机由 TMA 完成决定，与轮询线程数无关。1 个线程 wait + 1 个 __syncthreads 是更高效的做法。

对比 12 课 Pipeline： 12 课中 pipeline.consumer_wait() 是所有 consumer 线程都调用的，因为 PipelineTmaAsync 内部做了优化（分散到不同 warp、支持 cluster），不存在简单 spin-wait 的问题。

---

4. 验证方式

搬运完成后，每个线程从 smem 读数据写回另一块 gmem，然后 host 端对比原始数据和写回数据：

// 通过 smem tensor sA 访问（自动处理 swizzle 解码）
result_ptr[global_row * N + global_col] = sA(row, col);

注意：直接用裸指针读 smem 数据会因为 swizzle 而得到错误的位置。必须通过 sA(row, col) 访问，CuTe 会自动应用 swizzle 映射。

---

5. TMA vs cp.async 对比

特性	cp.async (SM80/Ampere)	TMA (SM90/Hopper)
发起线程数	所有线程各搬一部分	单线程发起整个 tile
地址计算	每线程自己算	硬件自动（通过 Descriptor）
Swizzle	需要手动 Layout	Descriptor 编码，硬件处理
同步机制	cp_async_fence / wait	mbarrier (phase 翻转)
Descriptor	不需要	Host 端构造，Kernel 端传入
多维支持	只能 1D（逐行搬）	最多 5D tensor 直接搬
Warp 专用化	困难（所有线程都搬运）	自然支持（单线程搬运）

TMA 的关键优势：搬运只需 1 个线程 → 其余线程可以做计算 → Warp Specialization 的基础。

---

6. TMA 与 LSU 的代理可见性（Proxy Visibility）

6.1 两条独立的硬件通路

Hopper 上共享内存有两条独立的访问通路：

LSU (Load/Store Unit):  普通线程的 ld.shared / st.shared 指令
TMA 引擎:               cp.async.bulk（TMA Load / TMA Store）

它们是不同的硬件代理（proxy），互相看不到对方的 in-flight 写入。

6.2 TMA 写 → LSU 读：无需额外操作

本课使用的场景：

TMA 引擎写入 smem → mbarrier wait → LSU 读取 smem

这是安全的，因为：

• mbarrier.wait 保证 TMA 搬运已完成（transaction bytes 归零）
• TMA 写入完成后，数据对 LSU 可见（TMA → LSU 方向有硬件保证）

6.3 LSU 写 → TMA 读：需要 fence！

如果反过来，线程用普通指令写 smem，然后让 TMA 引擎读 smem（如 TMA_STORE：Smem → Global）：

// 危险！LSU 写入可能对 TMA 不可见
sdata[idx] = value;
__syncthreads();         // 只保证线程间同步，不保证跨代理可见性
if (elected) {
    TMA::store(sdata, ...);  // TMA 可能读到旧数据！
}

问题：__syncthreads 保证所有线程执行到同一点，但 LSU 的 st.shared 可能还在流水线上，数据尚未真正写入 smem。如果是 LSU 自己读（走同一条流水线），顺序性保证数据可见；但 TMA 引擎走的是另一条通路，看不到 LSU 流水线中尚未提交的写入。

修复：在 __syncthreads 前加 fence.proxy.async.shared::cta：

sdata[idx] = value;
asm volatile("fence.proxy.async.shared::cta;\n");  // 确保 LSU 写入对 TMA 可见
__syncthreads();         // 线程同步
if (elected) {
    TMA::store(sdata, ...);  // 现在 TMA 能看到正确数据 ✓
}

fence.proxy.async.shared::cta 的作用：

• 等待 fence 之前所有 LSU 对 smem 的写入完成
• 使这些写入对异步代理（TMA）可见
• TMA 后续的读取会等待 fence 完成后才开始

6.4 CUTLASS 中的封装

CUTLASS 在 cute/arch/copy_sm90_tma.hpp 中提供了封装函数：

// cute::tma_store_fence() — 在 LSU 写 smem 和 TMA Store 之间调用
CUTE_HOST_DEVICE static void tma_store_fence() {
    asm volatile ("fence.proxy.async.shared::cta;");
}

典型的 TMA Store 流程：

// 1. 线程通过 LSU 写入 smem（例如 MMA 累加器 → smem）
sA(row, col) = accumulator_value;

// 2. fence：确保 LSU 写入对 TMA 可见
cute::tma_store_fence();

// 3. 线程同步
__syncthreads();

// 4. 单线程发起 TMA Store（smem → global）
if (tid == 0) {
    copy(tma_store, smem_src, gmem_dst);
    cute::tma_store_arrive();    // cp.async.bulk.commit_group
}
cute::tma_store_wait<0>();       // 等待 TMA Store 完成

6.5 总结：四种组合

写入方	读取方	是否需要 fence	原因
TMA	LSU	不需要	mbarrier.wait 已保证可见性
LSU	LSU	不需要	同一流水线，顺序保证
LSU	TMA	需要 fence.proxy.async.shared::cta	跨代理，LSU 写入对 TMA 不可见
TMA	TMA	不需要	同一代理，顺序保证

本课只用了 TMA_LOAD（TMA 写 → LSU 读），所以不需要 fence。但在实际 GEMM kernel 的 Epilogue 中（累加器 → smem → TMA Store → global），必须加 fence。

---

7. API 总结

Host 端

API	作用
make_tma_copy(SM90_TMA_LOAD{}, gmem, smem_layout)	构造 TMA descriptor
SM90_TMA_LOAD{}	TMA Load 操作类型（Global→Smem）

Kernel 端

API	作用
tma_load.get_tma_tensor(shape)	获取坐标 tensor（arithTuple）
local_tile(coord, tile_shape, coord)	切出当前 block 的 tile 坐标
tma_load.get_slice(0)	获取线程 0 的搬运 partition
thr.partition_S(coord_tile)	Source partition（坐标）
thr.partition_D(smem_tensor)	Destination partition（smem）
tma_load.with(mbar)	绑定 mbarrier，搬运完成自动 arrive
copy(tma.with(mbar), src, dst)	发起 TMA 搬运

mbarrier（第 10 课 API）

API	作用
initialize_barrier(mbar, 1)	初始化（1 个 arrive）
set_barrier_transaction_bytes(mbar, bytes)	设置期望搬运字节 + 隐式 arrive
wait_barrier(mbar, 0)	等待 phase 翻转（搬运完成）