cute（12）async_pipeline

前置阅读

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

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

12 Async Pipeline

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

核心概念

PipelineTmaAsync 是 CUTLASS 对 Hopper 多阶段异步流水线的封装：

• 每个 stage 有一对 mbarrier：full_barrier（数据就绪）+ empty_barrier（buffer 空闲）
• PipelineState 管理 (index_, phase_, count_) 三元组，自动翻转 phase
• Producer/Consumer 角色分工：Producer 发 TMA，Consumer 消费数据
• 用 make_producer_start_state() 初始化 producer 的 phase=1（因为 buffer 初始为空）

---

代码实现

/**
 * 第 12 课：Async Pipeline —— 多阶段异步流水线
 *
 * 学习目标：
 *   - PipelineTmaAsync：full_barrier + empty_barrier 每 stage 一对
 *   - PipelineState：index_、phase_、count_ 的状态转移
 *   - Producer / Consumer 分工：
 *     Producer：producer_acquire → TMA copy → ++pipe_write
 *     Consumer：consumer_wait → 使用数据 → consumer_release → ++pipe_read
 *   - make_producer_start_state()：producer 的 phase 初始化为 1 的原因
 *   - 实验：3-stage TMA 搬运流水线，producer 搬多个 tile，consumer 验证
 *
 * 核心概念：
 *   full_barrier[stage]  — producer 搬运完成后翻转，consumer 等待
 *   empty_barrier[stage] — consumer 用完数据后翻转，producer 等待
 *
 *   producer phase 从 1 开始：因为 buffer 一开始是空的（empty），
 *   empty_barrier 初始 phase=0，producer 等 phase=0 翻转
 *   → 但 buffer 本来就是空的 → empty_barrier 等效已经是 "可写" 状态
 *   → 所以 producer 用 phase=1 来跳过第一轮等待
 */

#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>  // GMMA::Layout_K_SW128_Atom
#include <cute/arch/copy_sm90_desc.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/cutlass.h>
#include <cstdio>
#include <cstdlib>

using namespace cute;

// ============================================================
// Pipeline 配置
// ============================================================
constexpr int kTileM = 128;
constexpr int kTileN = 64;
constexpr int kStage = 3;
constexpr int kNumTiles = 8;   // K 方向的 tile 数
constexpr int kBlockThreads = 128;  // 1 warp group = 4 warps

// Smem Layout（带 Swizzle，TMA 兼容）
// 数据 row-major (M, N) stride (N, 1) → stride-1 沿第二维 → K-major atom
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<half_t>;
using SmemLayout = decltype(tile_to_shape(
    SmemLayoutAtom{},
    make_shape(Int<kTileM>{}, Int<kTileN>{}, Int<kStage>{})));

// ============================================================
// Async Pipeline Kernel
// ============================================================
// Warp-specialized：
//   warp 0 (threads 0-31)   = Producer（发起 TMA）
//   warp 1-3 (threads 32-127) = Consumer（验证数据）
// Shared memory structure (pipeline barriers first for alignment)
struct SharedStorage {
  using Pipeline = cutlass::PipelineTmaAsync<kStage>;
  typename Pipeline::SharedStorage pipeline;
  alignas(128) half_t data[cosize(SmemLayout{})];
};

template <class TmaLoad>
__global__ void async_pipeline_kernel(
    __grid_constant__ const TmaLoad tma_load,
    half_t const* gmem_ptr,
    int* error_count,
    int M, int N) {

  using T = half_t;
  using Pipeline = cutlass::PipelineTmaAsync<kStage>;
  using PipelineState = typename Pipeline::PipelineState;

  // ============================================================
  // Shared Memory 布局
  // ============================================================
  extern __shared__ char smem_buf[];
  auto& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);

  auto sA = make_tensor(make_smem_ptr(shared_storage.data), SmemLayout{});  // (kTileM, kTileN, kStage)

  int tid = threadIdx.x;
  int warp_idx = tid / 32;
  int lane_idx = tid % 32;

  // ============================================================
  // 确定角色：Producer or Consumer
  // ============================================================
  bool is_producer = (warp_idx == 0);
  auto role = is_producer
      ? Pipeline::ThreadCategory::Producer
      : Pipeline::ThreadCategory::Consumer;

  int num_consumers = kBlockThreads - 32;  // 96 个 consumer 线程

  typename Pipeline::Params pipeline_params;
  pipeline_params.role = role;
  pipeline_params.is_leader = is_producer && (lane_idx == 0);  // warp 0 的 lane 0
  pipeline_params.num_consumers = num_consumers;
  pipeline_params.num_producers = 1;  // 只有 leader (1个线程) 做 arrive_and_expect_tx
  pipeline_params.transaction_bytes = kTileM * kTileN * sizeof(T);

  // ============================================================
  // 初始化 Pipeline
  // ============================================================
  Pipeline pipeline(shared_storage.pipeline, pipeline_params, Shape<_1, _1>{});
  __syncthreads();

  // ============================================================
  // TMA 坐标准备
  // ============================================================
  auto coord_tensor = tma_load.get_tma_tensor(make_shape(M, N));
  auto coord_col = local_tile(
      coord_tensor,
      make_tile(Int<kTileM>{}, Int<kTileN>{}),
      make_coord(0, 0));  // 固定在 (0,0) block

  auto tma_thr = tma_load.get_slice(0);

  // ============================================================
  // Producer 逻辑
  // ============================================================
  if (is_producer) {
    // Producer 的初始状态：phase=1
    PipelineState pipe_write = cutlass::make_producer_start_state<Pipeline>();

    if (lane_idx == 0 && blockIdx.x == 0) {
      printf("[Pipeline] Producer 初始状态: index=%d, phase=%d, count=%d\n",
             pipe_write.index(), pipe_write.phase(), pipe_write.count());
      printf("  ★ phase=1 的原因：buffer 初始为空, producer 可以直接写入\n");
    }

    for (int tile = 0; tile < kNumTiles; ++tile) {
      // producer_acquire：等 empty_barrier → buffer 可写
      pipeline.producer_acquire(pipe_write);

      // 发起 TMA copy
      if (lane_idx == 0) {
        // 模拟 K 方向不同 tile 的坐标（这里简化为搬同一个 tile）
        auto tSgS = tma_thr.partition_S(coord_col);
        auto tSsS = tma_thr.partition_D(sA(_, _, pipe_write.index()));
        copy(tma_load.with(
            *pipeline.producer_get_barrier(pipe_write)), tSgS, tSsS);

        if (blockIdx.x == 0 && tile < 4) {
          printf("  Producer: tile %d → stage %d (phase=%d)\n",
                 tile, pipe_write.index(), pipe_write.phase());
        }
      }

      // 推进 write 状态
      ++pipe_write;
    }
  }

  // ============================================================
  // Consumer 逻辑
  // ============================================================
  if (!is_producer) {
    PipelineState pipe_read;

    if (warp_idx == 1 && lane_idx == 0 && blockIdx.x == 0) {
      printf("[Pipeline] Consumer 初始状态: index=%d, phase=%d, count=%d\n",
             pipe_read.index(), pipe_read.phase(), pipe_read.count());
    }

    for (int tile = 0; tile < kNumTiles; ++tile) {
      // consumer_wait：等 full_barrier → 数据已就绪
      pipeline.consumer_wait(pipe_read);

      // 使用数据：验证 smem 中的值
      int stage = pipe_read.index();
      int local_errors = 0;

      // 每个 consumer 线程验证一部分数据
      int consumer_tid = tid - 32;  // consumer 从 0 开始编号
      for (int i = consumer_tid; i < kTileM * kTileN; i += num_consumers) {
        int row = i / kTileN;
        int col = i % kTileN;
        T expected = T(float((row * N + col) % 100) * 0.01f);
        T got = sA(row, col, stage);
        if (float(expected) != float(got)) {
          local_errors++;
        }
      }

      if (local_errors > 0) {
        atomicAdd(error_count, local_errors);
      }

      if (warp_idx == 1 && lane_idx == 0 && blockIdx.x == 0 && tile < 4) {
        printf("  Consumer: tile %d from stage %d (phase=%d) — verified\n",
               tile, stage, pipe_read.phase());
      }

      // consumer_release：通知 producer 这个 stage 的 buffer 可以覆盖了
      pipeline.consumer_release(pipe_read);

      // 推进 read 状态
      ++pipe_read;
    }
  }
}

int main() {
  printf("=== 第 12 课：Async Pipeline ===\n\n");

  using T = half_t;

  // 虚拟矩阵（只搬运第一个 tile 来演示 pipeline 机制）
  int M = kTileM;
  int N = kTileN;

  printf("Tile: %d x %d, Stages: %d, Tiles to process: %d\n",
         kTileM, kTileN, kStage, kNumTiles);
  printf("Threads: %d (warp0=producer, warp1-3=consumer)\n\n", kBlockThreads);

  // 分配数据
  T* h_data = (T*)malloc(M * N * sizeof(T));
  T* d_data;
  int* d_errors;
  int h_errors = 0;

  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_errors, sizeof(int));
  cudaMemcpy(d_data, h_data, M * N * sizeof(T), cudaMemcpyHostToDevice);
  cudaMemset(d_errors, 0, sizeof(int));

  // 构造 gmem tensor & TMA descriptor
  auto gmem_tensor = make_tensor(
      make_gmem_ptr(d_data),
      make_shape(M, N),
      make_stride(N, Int<1>{}));

  // 单 stage 的 smem layout（给 TMA descriptor 用）
  using SmemLayoutSingle = decltype(tile_to_shape(
      SmemLayoutAtom{},
      make_shape(Int<kTileM>{}, Int<kTileN>{})));

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

  // 计算 smem 大小
  int smem_size = sizeof(SharedStorage);
  smem_size = (smem_size + 127) & ~127;  // 128字节对齐
  constexpr int data_smem_size = cosize(SmemLayout{}) * sizeof(T);

  printf("Shared memory: %d bytes (data=%d, pipeline=%d)\n\n",
         smem_size, data_smem_size,
         (int)sizeof(typename cutlass::PipelineTmaAsync<kStage>::SharedStorage));

  cudaFuncSetAttribute(
      async_pipeline_kernel<decltype(tma_load)>,
      cudaFuncAttributeMaxDynamicSharedMemorySize,
      smem_size);

  // ============================================================
  // PipelineState 状态转移说明
  // ============================================================
  printf("========================================\n");
  printf("PipelineState 状态转移（3-stage 示例）\n");
  printf("========================================\n");
  printf("  操作   │ index │ phase │ 说明\n");
  printf("  ───────┼───────┼───────┼────────────────\n");
  printf("  初始   │   0   │   0   │ consumer 从 stage 0 开始\n");
  printf("  ++     │   1   │   0   │ 移到 stage 1\n");
  printf("  ++     │   2   │   0   │ 移到 stage 2\n");
  printf("  ++     │   0   │   1   │ 回到 stage 0, phase 翻转!\n");
  printf("  ++     │   1   │   1   │\n");
  printf("  ++     │   2   │   1   │\n");
  printf("  ++     │   0   │   0   │ 回到 stage 0, phase 再翻转!\n");
  printf("\n  Producer 初始 phase=1（因为 buffer 空 → 可直接写）\n");
  printf("  Consumer 初始 phase=0（等待 producer 填充数据）\n\n");

  // 启动 kernel
  printf("========================================\n");
  printf("运行 Pipeline Kernel\n");
  printf("========================================\n");

  async_pipeline_kernel<decltype(tma_load)>
      <<<1, kBlockThreads, smem_size>>>(tma_load, d_data, d_errors, M, N);
  cudaDeviceSynchronize();

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

  cudaMemcpy(&h_errors, d_errors, sizeof(int), cudaMemcpyDeviceToHost);
  printf("\n验证: %s (%d errors)\n", h_errors == 0 ? "PASS" : "FAIL", h_errors);

  // ============================================================
  // Pipeline 架构图
  // ============================================================
  printf("\n========================================\n");
  printf("Pipeline 流程图\n");
  printf("========================================\n");
  printf("\n");
  printf("  Producer (warp 0)           Consumer (warp 1-3)\n");
  printf("  ─────────────────           ──────────────────\n");
  printf("  producer_acquire(S0)        \n");
  printf("    ↓ TMA copy → S0          \n");
  printf("  producer_acquire(S1)        consumer_wait(S0)\n");
  printf("    ↓ TMA copy → S1            ↓ use data in S0\n");
  printf("  producer_acquire(S2)        consumer_release(S0)\n");
  printf("    ↓ TMA copy → S2          consumer_wait(S1)\n");
  printf("  producer_acquire(S0) ←wait    ↓ use data in S1\n");
  printf("    ↓ TMA copy → S0          consumer_release(S1)\n");
  printf("  ...                         ...\n");
  printf("\n");
  printf("  full_barrier[s]:  producer 写完 → consumer 可读\n");
  printf("  empty_barrier[s]: consumer 用完 → producer 可写\n");

  free(h_data);
  cudaFree(d_data);
  cudaFree(d_errors);

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

打印信息如下

[yuyanqi@set-hldy-mlp-codelab-pc75 build]$ ./12_async_pipeline 
=== 第 12 课：Async Pipeline ===

Tile: 128 x 64, Stages: 3, Tiles to process: 8
Threads: 128 (warp0=producer, warp1-3=consumer)

Shared memory: 49280 bytes (data=49152, pipeline=48)

========================================
PipelineState 状态转移（3-stage 示例）
========================================
  操作   │ index │ phase │ 说明
  ───────┼───────┼───────┼────────────────
  初始   │   0   │   0   │ consumer 从 stage 0 开始
  ++     │   1   │   0   │ 移到 stage 1
  ++     │   2   │   0   │ 移到 stage 2
  ++     │   0   │   1   │ 回到 stage 0, phase 翻转!
  ++     │   1   │   1   │
  ++     │   2   │   1   │
  ++     │   0   │   0   │ 回到 stage 0, phase 再翻转!

  Producer 初始 phase=1（因为 buffer 空 → 可直接写）
  Consumer 初始 phase=0（等待 producer 填充数据）

========================================
运行 Pipeline Kernel
========================================
[Pipeline] Consumer 初始状态: index=0, phase=0, count=0
[Pipeline] Producer 初始状态: index=0, phase=1, count=0
  ★ phase=1 的原因：buffer 初始为空, producer 可以直接写入
  Producer: tile 0 → stage 0 (phase=1)
  Consumer: tile 0 from stage 0 (phase=0) — verified
  Producer: tile 1 → stage 1 (phase=1)
  Consumer: tile 1 from stage 1 (phase=0) — verified
  Producer: tile 2 → stage 2 (phase=1)
  Consumer: tile 2 from stage 2 (phase=0) — verified
  Producer: tile 3 → stage 0 (phase=0)
  Consumer: tile 3 from stage 0 (phase=1) — verified

验证: PASS (0 errors)

========================================
Pipeline 流程图
========================================

  Producer (warp 0)           Consumer (warp 1-3)
  ─────────────────           ──────────────────
  producer_acquire(S0)        
    ↓ TMA copy → S0          
  producer_acquire(S1)        consumer_wait(S0)
    ↓ TMA copy → S1            ↓ use data in S0
  producer_acquire(S2)        consumer_release(S0)
    ↓ TMA copy → S2          consumer_wait(S1)
  producer_acquire(S0) ←wait    ↓ use data in S1
    ↓ TMA copy → S0          consumer_release(S1)
  ...                         ...

  full_barrier[s]:  producer 写完 → consumer 可读
  empty_barrier[s]: consumer 用完 → producer 可写

---

1. Pipeline 的双 Barrier 机制

1.1 为什么需要两个 Barrier

问题：Producer 写 buffer 的同时 Consumer 在读同一 buffer → 数据竞争

解决：每个 stage 两个 barrier，形成"请求-释放"协议

  full_barrier[s]:
    Producer 写完数据 → arrive
    Consumer wait → 确认数据就绪，可以读

  empty_barrier[s]:
    Consumer 读完数据 → arrive
    Producer wait → 确认 buffer 空闲，可以写

1.2 3-Stage Pipeline 流程

时间 →
                  Stage 0      Stage 1      Stage 2
Producer:
  tile 0 →        写入                                   pipe_write: S0→S1
  tile 1 →                     写入                      pipe_write: S1→S2
  tile 2 →                                  写入         pipe_write: S2→S0 (环绕)
  tile 3 →        写入 (等S0空)                          pipe_write: S0→S1
  ...

Consumer:
  tile 0 ←        读取                                   pipe_read: S0→S1
  tile 1 ←                     读取                      pipe_read: S1→S2
  tile 2 ←                                  读取         pipe_read: S2→S0 (环绕)
  tile 3 ←        读取                                   pipe_read: S0→S1
  ...

Producer 领先 Consumer 最多 kStage 步。当 Producer 写到第 4 个 tile 时，需要等 Consumer 释放 Stage 0。

---

2. PipelineState 状态转移

2.1 三元组 (index, phase, count)

struct PipelineState {
    int index_;  // 当前 stage 编号 (0 ~ kStage-1)
    int phase_;  // 当前 phase (0 或 1)
    int count_;  // 已处理的总步数
};

2.2 operator++ 的行为

初始:    index=0, phase=0, count=0
++:      index=1, phase=0, count=1
++:      index=2, phase=0, count=2
++:      index=0, phase=1, count=3  ← 回到 stage 0, phase 翻转！
++:      index=1, phase=1, count=4
++:      index=2, phase=1, count=5
++:      index=0, phase=0, count=6  ← 回到 stage 0, phase 再翻转！

Phase 翻转时机：当 index 从 kStage-1 回绕到 0 时，phase 翻转。这与 mbarrier 的 phase 自动翻转一致。

2.3 make_producer_start_state

PipelineState pipe_write = cutlass::make_producer_start_state<Pipeline>();
// 返回: index=0, phase=1, count=0

为什么 Producer phase 从 1 开始？

初始化时：
  empty_barrier 每个 stage 的初始 phase = 0
  buffer 一开始是空的 → 逻辑上已经"释放"了

如果 producer phase=0:
  producer_acquire 会 wait empty_barrier(phase=0)
  但 phase 确实是 0 → 需要等 consumer 先 arrive → 死锁！

如果 producer phase=1:
  producer_acquire 会 wait empty_barrier(phase=1)
  初始 phase=0 ≠ 1 → 条件已满足 → 不等待 → 直接写入 ✓

一句话：buffer 初始为空 = empty_barrier 已经"翻转过"了 = producer 直接用 phase=1 跳过等待。

---

3. Producer / Consumer API

3.1 Producer 流程

PipelineState pipe_write = cutlass::make_producer_start_state<Pipeline>();

for (int tile = 0; tile < num_tiles; ++tile) {
    // 1. 等待 buffer 空闲（empty_barrier）
    pipeline.producer_acquire(pipe_write);

    // 2. 发起 TMA copy（只有 leader 线程）
    if (lane_idx == 0) {
        auto barrier = pipeline.producer_get_barrier(pipe_write);
        copy(tma_load.with(*barrier), src, dst);
    }

    // 3. 推进状态
    ++pipe_write;
}

// 4. 收尾（通知 consumer 不再有更多数据）
pipeline.producer_tail(pipe_write);

3.2 Consumer 流程

PipelineState pipe_read;  // index=0, phase=0, count=0

for (int tile = 0; tile < num_tiles; ++tile) {
    // 1. 等待数据就绪（full_barrier）
    pipeline.consumer_wait(pipe_read);

    // 2. 使用数据
    int stage = pipe_read.index();
    // ... 读 sA(_, _, stage) ...

    // 3. 释放 buffer（通知 producer 可以覆盖）
    pipeline.consumer_release(pipe_read);

    // 4. 推进状态
    ++pipe_read;
}

3.3 API 含义总结

API	操作的 Barrier	含义
producer_acquire	wait empty_barrier	等 buffer 空闲
producer_get_barrier	返回 full_barrier	给 TMA 绑定，搬完自动 arrive
producer_tail	—	通知 consumer：数据全部发完
consumer_wait	wait full_barrier	等数据就绪
consumer_release	arrive empty_barrier	通知 producer：buffer 可覆盖

---

4. Pipeline 参数配置

4.1 Pipeline::Params

typename Pipeline::Params pipeline_params;
pipeline_params.role = role;              // Producer 或 Consumer
pipeline_params.is_leader = is_producer && (lane_idx == 0);
pipeline_params.num_consumers = 96;       // consumer 线程数
pipeline_params.num_producers = 1;        // 只有 leader 做 arrive_and_expect_tx
pipeline_params.transaction_bytes = kTileM * kTileN * sizeof(T);

4.2 关键参数说明

num_producers = 1

pipeline 构造函数内部用 num_producers 初始化 full_barrier：
  full_barrier.init(num_producers)  → 期望 num_producers 个 arrive

如果设 num_producers=32（整个 warp）：
  full_barrier 期望 32 个 arrive
  但只有 leader（1 个线程）做 arrive_and_expect_tx
  → 永远收不到 32 个 arrive → 死锁！

正确做法：num_producers=1，只有 leader arrive

is_leader

is_leader = is_producer && (lane_idx == 0)

只有 leader 做：
  - arrive_and_expect_tx（设置期望搬运字节数）
  - 发起 TMA copy

其他 producer 线程只做 arrive（不设 tx bytes）

cluster_shape = Shape<_1, _1>{}

Pipeline pipeline(shared_storage.pipeline, pipeline_params, Shape<_1, _1>{});

PipelineTmaAsync 构造函数访问 size<0> 和 size<1>，必须是至少 2D。即使单 block cluster，也要写 Shape<_1, _1> 而非 Shape<_1>。

---

5. SharedStorage 结构

5.1 为什么用 struct

struct SharedStorage {
  using Pipeline = cutlass::PipelineTmaAsync<kStage>;
  typename Pipeline::SharedStorage pipeline;   // mbarrier 数组（必须先声明！）
  alignas(128) half_t data[cosize(SmemLayout{})];  // 数据 buffer
};

Pipeline barriers 必须放在 struct 最前面：

• mbarrier 要求 8 字节对齐
• Pipeline::SharedStorage 内部是 mbarrier 数组（每 stage 2 个 barrier = kStage × 2 × 8 bytes）
• 放在 struct 最前面自然满足对齐
• 如果用手动 offset 计算，很容易对齐出错导致死锁

5.2 alignas(128)

数据数组 alignas(128) 确保 TMA 搬运的目标地址满足 128 字节对齐要求。

---

6. 源码解析：数据结构与方法

6.1 PipelineState — 环形游标（不是 barrier）

源码位置：cutlass/pipeline/sm90_pipeline.hpp:170-250

template<uint32_t Stages_>   // Stages_ = 3
struct PipelineState {
    static constexpr uint32_t Stages = Stages_;
    int index_ = 0;           // 当前 stage 编号 (0 ~ Stages-1)
    uint32_t phase_ = 0;      // 当前 phase (0 或 1)
    uint32_t count_ = 0;      // 已推进的总步数（仅统计用）
};

PipelineState 本身不是 barrier，它只是一个游标，告诉 producer/consumer 当前该操作第几个 stage、用哪个 phase 去 wait。

operator++ 源码：

void operator++() {
    if constexpr (Stages > 0) {
        ++index_;
        ++count_;
        if (index_ == Stages) {   // 走完一圈
            index_ = 0;           // 回绕到 stage 0
            phase_ ^= 1;          // phase XOR 翻转
        }
    }
}

为什么每圈翻转 phase？ mbarrier 的 phase 也是每轮所有 arrive 完成后自动翻转。PipelineState 的 phase 必须与 mbarrier 保持同步，这样 wait(phase) 才能正确工作。

advance 源码 — 批量前进：

PipelineState& advance(uint32_t num_iterations) {
    // 情况 1：前进不到一圈，但跨过了 stage 边界 → 翻转一次
    if ((num_iterations < Stages) && (index_ + num_iterations) >= Stages)
        phase_ ^= 1;
    // 情况 2：前进超过一圈，跨越奇数次边界 → 翻转一次
    if ((num_iterations >= Stages) && (((index_ + num_iterations) / Stages) % 2) == 1)
        phase_ ^= 1;
    index_ = (index_ + num_iterations) % Stages;
    count_ += num_iterations;
}

operator++ 的批量版本。producer_tail 中用到它（等待 Stages 轮）。

make_producer_start_state 源码：

template<class Pipeline>
PipelineState<Pipeline::Stages> make_producer_start_state() {
    // Producer starts with an opposite phase as the buffers are initially empty
    constexpr int InitialProducerStage = 0;
    constexpr uint32_t InitialProducerPhase = 1;   // ← 关键！
    constexpr uint32_t InitialProducerCount = 0;
    return {InitialProducerStage, InitialProducerPhase, InitialProducerCount};
}

注释说得很清楚：“opposite phase because buffers are initially empty”。

6.2 ClusterBarrier — empty_barrier 的类型

源码位置：cutlass/arch/barrier.h:341-539

struct ClusterBarrier {
    uint64_t barrier_;   // 64-bit mbarrier，必须在 smem 中
};

用于 empty_barrier——追踪 consumer 是否释放了 buffer。只基于 arrive count 同步，不涉及 transaction bytes。

init — PTX: mbarrier.init

static void init(ValueType const* smem_ptr, uint32_t arrive_count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;"
        : : "r"(arrive_count), "r"(smem_addr));
}

初始化 barrier：phase=0，arrive_count=参数值。对于 empty_barrier，arrive_count = num_consumers = 96。

wait — PTX: mbarrier.try_wait.parity 忙等循环

static void wait(ValueType const* smem_ptr, uint32_t phase) {
    uint32_t ticks = 0x989680;    // ~10M 次超时重试
    asm volatile(
        "LAB_WAIT: \n\t"
        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
        "@P1 bra DONE; \n\t"
        "bra     LAB_WAIT; \n\t"
        "DONE: \n\t"
        : : "r"(smem_addr), "r"(phase), "r"(ticks));
}

语义：wait(phase) = “等到 barrier 当前 phase ≠ phase”。注意不是"等 phase 变成某个值"，而是"等 phase 不再是传入的值"。

arrive — PTX: mbarrier.arrive

// 本地 arrive（consumer_release 在 cluster_size=1 时使用）
static void arrive(ValueType const* smem_ptr) {
    asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(smem_addr));
}

// 带 cta_id 的 arrive（支持 cluster 跨 block）
static void arrive(ValueType const* smem_ptr, uint32_t cta_id, uint32_t pred) {
    asm volatile(
        "mapa.shared::cluster.u32  remAddr32, %0, %1;\n\t"
        "mbarrier.arrive.shared::cluster.b64  _, [remAddr32];\n\t"
        : : "r"(smem_addr), "r"(cta_id));
}

arrive_count 减 1。当减到 0 时，phase 自动翻转，barrier 自动重置 arrive_count。

6.3 ClusterTransactionBarrier — full_barrier 的类型

源码位置：cutlass/arch/barrier.h:545-663

struct ClusterTransactionBarrier : public ClusterBarrier {
    // 继承了 ClusterBarrier 的 init/wait/arrive
    // 新增 transaction bytes 相关方法
};

用于 full_barrier——追踪 TMA 搬运是否完成。翻转条件比 ClusterBarrier 更严格：arrive_count == 0 && transaction_count == 0。

arrive_and_expect_tx — PTX: mbarrier.arrive.expect_tx

static void arrive_and_expect_tx(ValueType const* smem_ptr, uint32_t transaction_bytes) {
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0;"
        : : "r"(transaction_bytes), "r"(smem_addr));
}

一条 PTX 指令做两件事：

1. arrive_count 减 1
2. transaction_count 加 transaction_bytes

TMA 搬运完成后硬件自动调用 complete_tx 减 transaction_count。当 arrive_count 和 transaction_count 都归零 → phase 翻转。

complete_transaction — PTX: mbarrier.complete_tx

static void complete_transaction(ValueType const* smem_ptr, uint32_t dst_cta_id,
                                  uint32_t transaction_bytes, uint32_t pred) {
    asm volatile("@p mbarrier.complete_tx.shared::cluster.relaxed.cluster.b64 [%1], %0;"
        : : "r"(transaction_bytes), "r"(smem_addr), "r"(pred));
}

正常 TMA 场景不需要手动调用——copy(tma_load.with(*barrier), ...) 绑定 barrier 后，TMA 硬件搬运完自动 complete_tx。

6.4 Pipeline::SharedStorage — 每 stage 一对 barrier

// PipelineTmaAsync 内部定义
struct SharedStorage {
    FullBarrier  full_barrier_[Stages];    // ClusterTransactionBarrier[3]
    EmptyBarrier empty_barrier_[Stages];   // ClusterBarrier[3]
};

smem 中的布局：

偏移 0:   [full_barrier[0]]  [full_barrier[1]]  [full_barrier[2]]    ← 3 × 8 bytes
偏移 24:  [empty_barrier[0]] [empty_barrier[1]] [empty_barrier[2]]   ← 3 × 8 bytes
共 48 bytes barrier 数据

初始化（init_barriers 内部）：

for (int i = 0; i < Stages; i++) {
    full_barrier[i].init(num_producers);     // init(1) → arrive_count=1, phase=0
    empty_barrier[i].init(num_consumers);    // init(96) → arrive_count=96, phase=0
}

6.5 PipelineTmaAsync — Pipeline 类

class PipelineTmaAsync {
    FullBarrier  *full_barrier_ptr_;    // 指向 SharedStorage 中的 full_barrier 数组
    EmptyBarrier *empty_barrier_ptr_;   // 指向 SharedStorage 中的 empty_barrier 数组
    Params params_;                     // 角色、leader、线程数、tx_bytes
    uint32_t dst_blockid_;             // cluster 中的目标 block id
    uint32_t is_signaling_thread_;     // 是否负责 cluster arrive
};

Producer 和 Consumer 共享同一对 barrier 数组。 构造函数传入同一个 shared_storage 引用：

// 12 课代码：所有线程（producer + consumer）构造同一个 pipeline
Pipeline pipeline(shared_storage.pipeline, pipeline_params, Shape<_1, _1>{});

内部：

PipelineTmaAsync(SharedStorage& storage, Params params, ...)
    : full_barrier_ptr_(&storage.full_barrier_[0])      // 共享
    , empty_barrier_ptr_(&storage.empty_barrier_[0]) {  // 共享
    init_barriers(storage, params_, cluster_shape);      // 初始化所有 barrier
}

6.6 producer_acquire 源码

void producer_acquire(uint32_t stage, uint32_t phase) {
    // ① 等 empty_barrier：consumer 是否释放了这个 stage
    empty_barrier_ptr_[stage].wait(phase);

    // ② leader arrive full_barrier + 设置 tx_bytes
    if (params_.is_leader) {
        full_barrier_ptr_[stage].arrive_and_expect_tx(params_.transaction_bytes);
    }
}

做了两件事，顺序很关键：

1. 先确认 buffer 空闲（wait empty_barrier）
2. 然后告诉 full_barrier “我要开始搬了，预期搬 N 字节”

为什么 arrive 在 TMA copy 之前？因为 arrive_and_expect_tx 同时做了 arrive（arrive_count 1→0）和 expect_tx（设 transaction_count=16KB）。此时 arrive_count=0 但 transaction_count=16KB → 不翻转。等 TMA 搬运完成后硬件 complete_tx → transaction_count 归零 → 两个都是 0 → phase 翻转 → consumer 的 wait 通过。

6.7 consumer_wait 源码

void consumer_wait(uint32_t stage, uint32_t phase) {
    full_barrier_ptr_[stage].wait(phase);
}

就一行：等 full_barrier[stage] 的 phase 翻转。翻转 = arrive_count 和 transaction_count 都归零 = TMA 搬运完成。

6.8 consumer_release 源码

void consumer_release(uint32_t stage, uint32_t skip = false) {
    empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signaling_thread_ & (!skip));
}

所有 consumer 线程都调用，但内部通过 is_signaling_thread_ 控制实际只有部分线程执行 PTX arrive 指令。当 96 个 consumer 全部 arrive → empty_barrier 的 arrive_count 归零 → phase 翻转 → producer 的 empty_barrier.wait 通过。

6.9 producer_tail 源码

void producer_tail(PipelineState state) {
    for (int count = 0; count < Stages; ++count) {
        empty_barrier_ptr_[state.index()].wait(state.phase());
        ++state;
    }
}

Producer 搬完所有 tile 后调用。等待所有 Stages 个 stage 的 empty_barrier 翻转，确保 consumer 消费完所有数据。如果不等，producer 提前退出，consumer 可能还在消费最后几个 stage。

6.10 producer_get_barrier 源码

ProducerBarrierType* producer_get_barrier(uint32_t stage) {
    return reinterpret_cast<ProducerBarrierType*>(&full_barrier_ptr_[stage]);
}

返回 full_barrier[stage] 的指针。在 copy(tma_load.with(*barrier), ...) 中绑定给 TMA，搬运完成后 TMA 硬件自动对这个 barrier 执行 complete_tx。

---

7. 完整 Barrier 变化时间线（3 stage, 8 tiles）

7.1 初始状态

初始化后（所有 barrier phase=0）：

full_barrier[0]:  phase=0, arrive_count=1 (等 producer arrive)
full_barrier[1]:  phase=0, arrive_count=1
full_barrier[2]:  phase=0, arrive_count=1

empty_barrier[0]: phase=0, arrive_count=96 (等 consumer arrive)
empty_barrier[1]: phase=0, arrive_count=96
empty_barrier[2]: phase=0, arrive_count=96

pipe_write = (index=0, phase=1)    ← producer 从 phase=1 开始
pipe_read  = (index=0, phase=0)    ← consumer 从 phase=0 开始

7.2 Tile 0 — pipe_write=(0, phase=1)

Producer: producer_acquire(index=0, phase=1)
  ① empty_barrier[0].wait(phase=1)
    empty_barrier[0] 当前 phase=0, 0≠1 → 条件已满足 → 不等待！
    （buffer 初始为空 = "已释放"，这就是 producer phase=1 的作用）
  ② full_barrier[0].arrive_and_expect_tx(16KB)
    full_barrier[0]: arrive_count 1→0, transaction_count 0→16KB
    (arrive_count=0 但 tx_count≠0 → 不翻转)

  TMA copy → stage 0（TMA 引擎异步搬运）
  TMA 完成后硬件自动 complete_tx → transaction_count 16KB→0
  → arrive_count=0 && transaction_count=0
  → full_barrier[0] phase 翻转: 0→1

  ++pipe_write → (1, phase=1)

Consumer: consumer_wait(index=0, phase=0)
  ③ full_barrier[0].wait(phase=0)
    等待中... TMA 完成后 phase 翻转为 1 → 1≠0 → wait 通过 ✓

  验证 stage 0 数据...

  consumer_release(index=0)
  ④ empty_barrier[0]: 96 个 consumer arrive → arrive_count 96→0
    → phase 翻转: 0→1

  ++pipe_read → (1, phase=0)

7.3 Tile 1 — pipe_write=(1, phase=1)

Producer: producer_acquire(index=1, phase=1)
  empty_barrier[1].wait(1): phase=0, 0≠1 → 不等待
  full_barrier[1].arrive_and_expect_tx(16KB)
  TMA copy → stage 1
  TMA 完成 → full_barrier[1] phase: 0→1
  ++pipe_write → (2, phase=1)

Consumer: consumer_wait(index=1, phase=0)
  full_barrier[1].wait(0) → TMA 完成 → phase=1, 1≠0 → 通过
  consumer_release(1) → empty_barrier[1] phase: 0→1
  ++pipe_read → (2, phase=0)

7.4 Tile 2 — pipe_write=(2, phase=1)

Producer: producer_acquire(index=2, phase=1)
  empty_barrier[2].wait(1): phase=0, 不等待
  full_barrier[2].arrive_and_expect_tx
  TMA copy → stage 2
  TMA 完成 → full_barrier[2] phase: 0→1
  ++pipe_write → (0, phase=0)  ← index 回绕, phase 翻转！

Consumer: consumer_wait(index=2, phase=0)
  full_barrier[2].wait(0) → TMA 完成 → 通过
  consumer_release(2) → empty_barrier[2] phase: 0→1
  ++pipe_read → (0, phase=1)  ← index 回绕, phase 翻转！

7.5 Tile 3 — pipe_write=(0, phase=0) — 第一次可能等待！

Producer: producer_acquire(index=0, phase=0)
  empty_barrier[0].wait(phase=0)
    empty_barrier[0] 当前 phase：
      tile 0 的 consumer_release 已将 phase 翻转到 1
      → 1≠0 → 不等待 ✓ （consumer 已经释放了 stage 0）

      如果 consumer 还没来得及 release → phase 仍是 0 → 等待！
      → 这就是反压机制：producer 不会覆盖 consumer 正在读的数据

  full_barrier[0].arrive_and_expect_tx
    full_barrier[0] 上一轮翻转后 phase=1
    arrive → arrive_count 1→0, transaction_count += 16KB
  TMA copy → stage 0（覆盖旧数据）
  TMA 完成 → full_barrier[0] phase: 1→0
  ++pipe_write → (1, phase=0)

Consumer: consumer_wait(index=0, phase=1)
  full_barrier[0].wait(phase=1)
    等 TMA 完成 → phase 1→0 → 0≠1 → 通过 ✓
  consumer_release(0) → empty_barrier[0] phase: 1→0
  ++pipe_read → (1, phase=1)

7.6 完整总结表

         pipe_write        empty_barrier[s]       full_barrier[s]          pipe_read
Tile  (idx,phase)    wait(phase) 是否等待   arrive+expect_tx→TMA完成   wait(phase)→通过   release→arrive   (idx,phase)
───── ──────────── ──────────────────────── ──────────────────────── ──────────────────── ──────────────── ────────────
  0   (0, ph=1)    wait(1): ph=0≠1 不等    full[0]: ph 0→1          wait(0): ph=1 通过  empty[0]: 0→1    (0, ph=0)
  1   (1, ph=1)    wait(1): ph=0≠1 不等    full[1]: ph 0→1          wait(0): ph=1 通过  empty[1]: 0→1    (1, ph=0)
  2   (2, ph=1)    wait(1): ph=0≠1 不等    full[2]: ph 0→1          wait(0): ph=1 通过  empty[2]: 0→1    (2, ph=0)
───── 回绕 ───────────────────────────────────────────────────────────────────────────────────── 回绕 ────────────
  3   (0, ph=0)    wait(0): ph=1≠0 不等*   full[0]: ph 1→0          wait(1): ph=0 通过  empty[0]: 1→0    (0, ph=1)
  4   (1, ph=0)    wait(0): ph=1≠0 不等*   full[1]: ph 1→0          wait(1): ph=0 通过  empty[1]: 1→0    (1, ph=1)
  5   (2, ph=0)    wait(0): ph=1≠0 不等*   full[2]: ph 1→0          wait(1): ph=0 通过  empty[2]: 1→0    (2, ph=1)
───── 回绕 ───────────────────────────────────────────────────────────────────────────────────── 回绕 ────────────
  6   (0, ph=1)    wait(1): ph=0≠1 不等*   full[0]: ph 0→1          wait(0): ph=1 通过  empty[0]: 0→1    (0, ph=0)
  7   (1, ph=1)    wait(1): ph=0≠1 不等*   full[1]: ph 0→1          wait(0): ph=1 通过  empty[1]: 0→1    (1, ph=0)

* "不等" 的前提是 consumer 已经 release，否则 producer 会被阻塞（反压）

核心规律：

• Tile 0-2（第一圈）：producer 的 empty_barrier.wait(phase=1) 全不等待，因为 buffer 初始为空（phase=0 ≠ 1）
• Tile 3 开始：producer 可能需要等 consumer release，这是反压机制
• 每走完一圈 3 个 stage，pipe_write 和 pipe_read 的 phase 各翻转一次，与 barrier 的 phase 翻转保持同步

7.7 整体调用链图

Producer 线程                              Consumer 线程
    │                                          │
    ▼                                          │
producer_acquire(pipe_write)                   │
  │ empty_barrier[s].wait(phase)  ←── 等 ──── consumer_release 的 arrive
  │ full_barrier[s].arrive_and_expect_tx       │
  ▼                                            │
copy(tma.with(*full_barrier[s]), ...)          │
  │ TMA 引擎异步搬运                            │
  │ 完成后硬件 complete_tx                       │
  │ → full_barrier[s] phase 翻转 ────→  consumer_wait(pipe_read)
  ▼                                     │ full_barrier[s].wait(phase) 通过
++pipe_write                            ▼
                                   使用 stage s 数据
                                        │
                                        ▼
                                   consumer_release(pipe_read)
                                     │ empty_barrier[s].arrive
                                     │ → 所有 consumer arrive 后 phase 翻转
                                     │ → producer 下一次 empty_barrier.wait 通过
                                     ▼
                                   ++pipe_read

---

8. 实验设计

6.1 配置

Tile:    128 × 64
Stages:  3
Tiles:   8（同一个 tile 重复搬运，演示 pipeline 多轮行为）
Threads: 128（warp 0 = producer, warp 1-3 = consumer）

6.2 Producer 行为

Producer 搬运 8 个 tile 到 3-stage 环形 buffer：

tile 0 → stage 0     pipe_write: (0, phase=1)
tile 1 → stage 1     pipe_write: (1, phase=1)
tile 2 → stage 2     pipe_write: (2, phase=1)
tile 3 → stage 0     pipe_write: (0, phase=0)  ← 需等 consumer 释放 S0
tile 4 → stage 1     pipe_write: (1, phase=0)
tile 5 → stage 2     pipe_write: (2, phase=0)
tile 6 → stage 0     pipe_write: (0, phase=1)
tile 7 → stage 1     pipe_write: (1, phase=1)

6.3 Consumer 行为

Consumer 从每个 stage 读取数据并验证正确性：

tile 0 ← stage 0     pipe_read: (0, phase=0)
tile 1 ← stage 1     pipe_read: (1, phase=0)
tile 2 ← stage 2     pipe_read: (2, phase=0)
tile 3 ← stage 0     pipe_read: (0, phase=1)  ← phase 翻转
tile 4 ← stage 1     pipe_read: (1, phase=1)
...

---

7. 与 09 课 cp.async 流水线对比

	09 课 (cp.async)	12 课 (PipelineTmaAsync)
搬运方式	全线程 cp.async	单线程 TMA
同步机制	cp_async_fence / wait	mbarrier (full + empty)
流水线管理	手动 index/fence/wait	PipelineState 自动管理
Warp 专用化	无（所有线程同一代码路径）	有（producer / consumer 分离）
Phase 管理	无需	PipelineState 自动翻转
Stage 状态	ismem_read / ismem_write	pipe_read / pipe_write
Producer 初始	无需特殊处理	phase=1（跳过首轮等待）
Consumer 释放	不需要（没有 empty_barrier）	consumer_release

09 课的简洁性：所有线程都参与搬运和计算，用 cp_async_fence/wait 全局控制批次。

12 课的灵活性：Producer/Consumer 完全分离，Producer 不参与计算，Consumer 不参与搬运。这样 producer warp 可以做其他工作（如 TMA store），consumer warp 专注计算。

---

8. API 总结

API	作用	使用者
PipelineTmaAsync<kStage>	管理 kStage 对 barrier 的 pipeline	构造函数
make_producer_start_state<Pipeline>()	返回 (0, phase=1, 0)	Producer
pipeline.producer_acquire(state)	等 buffer 空闲	Producer
pipeline.producer_get_barrier(state)	获取 full_barrier 指针	Producer
pipeline.producer_tail(state)	收尾通知	Producer
pipeline.consumer_wait(state)	等数据就绪	Consumer
pipeline.consumer_release(state)	释放 buffer	Consumer
PipelineState::operator++()	推进 (index, phase, count)	两者
PipelineState::index()	当前 stage 编号	两者
PipelineState::phase()	当前 phase (0/1)	两者