非同步資料複製：STAS (Async Data Copies with STAS)

#cuda #async-execution #async-copy #concept

重點總覽

項目	重點
用途	在 thread block cluster 內，於不同 thread block 之間搬移小型資料元素
硬體需求	STAS 指令需 CC 9.0+（Hopper 起，cluster 機制）
方向	唯一方向：register → distributed shared memory（無反向、無 global 端）
支援大小	一次複製 4、8 或 16 bytes
對齊	目的地指標需依大小對齊到 4、8 或 16 bytes
API	僅透過低階 `cuda::ptx::st_async`（libcu++）暴露，無高階包裝
非同步模型	視為 async thread operation（Async Thread / Async Proxy）
完成機制	只能用 shared memory barrier（mbarrier）標示完成
典型場景	cluster 內 producer-consumer，例如 8 blocks 環狀通訊

STAS 與兄弟機制的定位

STAS 是三種非同步資料複製機制中唯一「register → distributed shared memory」者，專為 cluster 內跨 block 的小資料交換而生；LDGSTS 與 TMA 都是 global ↔ shared 之間的搬運。三者並列見 04-CUDA-Features/13-Async-Copies-LDGSTS 與 04-CUDA-Features/14-Async-Copies-TMA。

STAS 是什麼

使用 thread block cluster 的應用，常需要在 cluster 內的各 thread block 之間搬移小型資料元素。STAS 指令（CC 9.0+，詳見 PTX 文件）讓資料能直接從 register 非同步複製到 distributed shared memory。

STAS 只透過較低階的 cuda::ptx::st_async API 暴露，該 API 位於 libcu++ 函式庫。
沒有 cuda::memcpy_async 這類高階包裝，因此使用上需要直接操作 PTX 層級的 helper。
distributed shared memory 是 cluster 機制提供的能力：cluster 內任一 block 可用 cluster.map_shared_rank 映射並存取「另一個 block」的 shared memory，STAS 正是把資料寫進這種遠端 shared 位址。

為什麼叫 STAS

STAS 對應 PTX 的 st.async（async store）。它不是「載入」也不是「bulk 拷貝」，而是把暫存器中的小值，以非同步 store 的形式投遞到遠端 block 的 shared memory。

維度、來源與目的地

Dimensions（維度）：STAS 支援複製 4、8 或 16 bytes。
Source and destination（來源與目的地）：唯一支援的非同步複製方向是 register → distributed shared memory。目的地指標需依資料大小對齊到 4、8 或 16 bytes。

方向與大小受限

STAS 不能反向（shared → register）、不能接觸 global memory，也不做大區塊搬運；單次只搬 4/8/16 bytes 且對齊要求嚴格。需要 global ↔ shared 的大量資料時請改用 LDGSTS 或 TMA。

非同步性與完成機制

STAS 的資料傳輸是非同步的，被建模為 async thread operation（見 Async Thread 與 Async Proxy）。
這讓發起的 thread 可以在硬體於背景非同步複製資料時，繼續做運算。
是否真的非同步發生取決於硬體實作，未來可能改變（屬於語意保證，而非時序保證）。
STAS 操作用來標示「已完成」的完成機制是 shared memory barrier（即 mbarrier）。

發起 thread:  st_async(...)  ──┐ (投遞後立即返回，可繼續運算)
                              │  硬體背景非同步複製
遠端 shared:                   └─▶ 寫入完成 ──▶ 透過 mbarrier 通知

完成不靠回傳值或 fence，而是靠目的端的 mbarrier 計數（透過 mbarrier_arrive_expect_tx 宣告預期位元組數、mbarrier_try_wait_parity 等待），與 04-CUDA-Features/11-Asynchronous-Barriers-Deep-Dive 的 transaction barrier 機制一致。

Producer-Consumer 範例

原文以一個 cluster 內 producer-consumer 範例示範 STAS：8 個 thread block 排成一個環，每個 block 同時：

為序列中的下一個 block 生產資料（producer）。
從序列中的上一個 block 消費資料（consumer）。

每個 block 需要 2 個 shared memory barrier：filled（通知 consumer 資料已複製進 shared buffer）與 ready（通知 producer consumer 端 buffer 已可被填寫）。

       st_async →        st_async →
  ┌────────┐        ┌────────┐        ┌────────┐
  │ block0 │ ─────▶ │ block1 │ ─────▶ │ block2 │ ─ ... ─▶ block7 ─┐
  └────────┘        └────────┘        └────────┘                  │
       ▲                                                          │
       └──────────────────────────────────────────────────────────┘
  每個 block：對右鄰是 producer，對左鄰是 consumer（環狀）

初始化與位址映射：

#include <cooperative_groups.h>
#include <cuda/barrier>
#include <cuda/ptx>

__global__ __cluster_dims__(8, 1, 1) void producer_consumer_kernel()
{
    using namespace cooperative_groups;
    using namespace cuda::device;
    using namespace cuda::ptx;
    using barrier_t = cuda::barrier<cuda::thread_scope_block>;
    auto cluster = this_cluster();

    __shared__ int      buffer[BLOCK_SIZE];
    __shared__ barrier_t filled;   // 通知 consumer：資料已寫入
    __shared__ barrier_t ready;    // 通知 producer：buffer 可再填

    if (threadIdx.x == 0) {        // 各 block 第一個 thread 初始化 barrier
        init(&filled, 1);
        init(&ready, BLOCK_SIZE);
    }
    cluster.sync();                // 確保所有遠端 barrier 都初始化完成

    int rk      = cluster.block_rank();
    int rk_next = (rk + 1) % 8;
    int rk_prev = (rk + 7) % 8;

    // 映射右鄰 buffer 與左右鄰 barrier 的遠端位址
    auto buffer_next = cluster.map_shared_rank(buffer, rk_next);
    auto bar_next = cluster.map_shared_rank(barrier_native_handle(filled), rk_next);
    auto bar_prev = cluster.map_shared_rank(barrier_native_handle(ready),  rk_prev);
}

filled 初始化為 1（單一 arrive），ready 初始化為 block 內 thread 數；cluster.sync() 保證任何 thread 開始通訊前所有 barrier 都已就緒。map_shared_rank 把本地物件映射成指向特定 rank 的遠端 shared 位址。

每次迭代的核心（producer 送資料、consumer 等待）：

int phase = 0;
for (int it = 0; it < 1000; ++it) {
    // Producer：用 STAS 把 register 值送到右鄰的 buffer
    st_async(&buffer_next[threadIdx.x], rk, bar_next);

    if (threadIdx.x == 0) {
        // thread 0 在本地 filled barrier arrive，並宣告預期接收的 bytes
        mbarrier_arrive_expect_tx(sem_release, scope_cluster, space_shared,
                                  barrier_native_handle(filled), sizeof(buffer));
    }
    // Consumer：等本地 filled barrier，等左鄰資料抵達
    while (!mbarrier_try_wait_parity(barrier_native_handle(filled), /* ... */)) { }
    // ... 使用資料、再用 ready barrier 通知左鄰可繼續 ...
}

迭代流程依原文為：1) producer 送資料給右鄰；2) consumer 的 thread 0 在 filled arrive 並宣告預期位元組；3) consumer 等 filled；4) consumer 使用資料；5) consumer 用 ready 通知左鄰已用完；6) producer 等本地 ready 直到右鄰可接收新資料。

barrier 的 space 必須正確

對「映射後的遠端 barrier」要用 space_cluster space；對「本地 barrier」要用 space_shared space。範例中 st_async 寫向遠端 bar_next（cluster space），而 mbarrier_arrive_expect_tx／mbarrier_try_wait_parity 操作本地 filled（shared space）。用錯 space 會語意不符。

與 LDGSTS / TMA 的差異

機制	方向	搬運型態	觸發 API	完成機制
LDGSTS (cp.async)	global → shared	小元素逐元素	`cuda::memcpy_async` / `cp.async`	mbarrier 或 pipeline
TMA (cp.async.bulk)	global ↔ shared	bulk / tensor 大區塊	TMA descriptor / `cuda::ptx`	mbarrier (transaction)
STAS (st.async)	register → distributed shared (cluster)	4/8/16 bytes 小資料	`cuda::ptx::st_async`	shared memory barrier (mbarrier)

LDGSTS 與 TMA 都把資料從 global memory 拉進 shared；STAS 完全不碰 global，而是在 cluster 內 block 之間直接傳 register 值。
TMA 適合大區塊／多維 tensor；STAS 反之，專做極小（≤16 bytes）的點對點投遞。
三者都用 mbarrier 標示完成，可與 04-CUDA-Features/12-Pipelines-Deep-Dive 的 pipeline 模型搭配；STAS 特別倚賴 cluster 與 distributed shared memory。

考試/測驗重點

問題	答案
STAS 的資料流向是？	register / distributed shared memory（唯一方向，不可反向）
STAS 支援的單次複製大小？	4 / 8 / 16 bytes
目的地指標的對齊需求？	依大小對齊 4 / 8 / 16 bytes
STAS 透過哪個 API 暴露？	低階 `cuda::ptx::st_async`（libcu++），無高階包裝
最低 compute capability？	CC 9.0+（需 cluster 支援）
STAS 用什麼標示完成？	shared memory barrier（mbarrier），非回傳值或 fence
STAS 被建模為哪種操作？	async thread operation（Async Thread / Async Proxy）
範例 producer-consumer 用幾個 barrier？	每 block 2 個：filled（資料已寫入）/ ready（buffer 可再填）
遠端 barrier 與本地 barrier 各用哪個 space？	遠端用 space_cluster / 本地用 space_shared
STAS vs LDGSTS / TMA 最大差異？	STAS 不碰 global、做 cluster 內小資料交換；LDGSTS / TMA 為 global ↔ shared 搬運