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

重點總覽

TMA 概觀：兩種模式與完成機制

TMA（在 PTX 中稱 bulk / bulk-tensor 非同步複製）的 source 與 destination 可在 shared 或 global，支援：global → shared、shared → global，以及 cluster 內 block 之間的 shared → distributed shared。在 cluster 中還可指定 multicast，把 global 的資料一次廣播到多個 block 的 shared memory（為 sm_90a 最佳化，其他 target 效能可能大幅下降）。

                       TMA 完成機制
  ┌─────────────────────────────────────────────────────────┐
  │ global → shared    : shared memory barrier (transaction)  │  讀：block 內任一 thread 可等
  │ global → s::cluster: shared memory barrier (multicast)    │
  │ shared → global    : bulk async-group (commit/wait)       │  寫：只有發起 thread 可等
  │ s::cta → s::cluster: shared memory barrier                │
  └─────────────────────────────────────────────────────────┘

• 讀（資料進 shared）：硬體完成讀取後更新 barrier 的 transaction count，barrier flip 後資料才安全可讀。
• 寫（資料出 shared）：用 thread-local 的 bulk async-group；多個寫可批次成一個 group，再等該 group「讀完 shared」或「寫完 global」。

4.11.2.1 用 TMA 搬 1D 陣列

1D 連續陣列不需要 tensor map，只用指標與大小即可在 device 端發起。對應 API：cuda::memcpy_async、cuda::device::memcpy_async_tx、cuda::ptx::cp_async_bulk。

1D read-modify-write 的七個步驟（global → shared → +1 → global）：

 1. init barrier(blockDim.x)                  ── 完成機制
 2. cp_async_bulk: global → shared            ── 發起讀（單一 thread）
 3. bar.arrive() ; bar.wait(token)            ── 等資料到齊
 4. smem_data[i] += 1                          ── 計算（所有 thread）
 5. fence_proxy_async(space_shared)+syncthreads── 讓 generic 寫對 async proxy 可見
 6. cp_async_bulk: shared → global            ── 發起寫（單一 thread）
 7. commit_group ; wait_group_read(0)          ── 等寫端讀完 shared

// 步驟 2/6 的三種寫法（擇一）
cuda::memcpy_async(smem_data, data + offset,
                   cuda::aligned_size_t<16>(sizeof(smem_data)), bar);   // 自動設 tx count
// 或 memcpy_async_tx / cp_async_bulk，需自行：
// cuda::ptx::mbarrier_expect_tx(...) 告知預期 transaction（以 byte 計）
ptx::cp_async_bulk(ptx::space_global, ptx::space_shared,
                   data + offset, smem_data, sizeof(smem_data));        // shared → global
ptx::cp_async_bulk_commit_group();
ptx::cp_async_bulk_wait_group_readn32_t<0>();

重點：memcpy_async 會自動更新 barrier 的預期 transaction count；用 memcpy_async_tx / cp_async_bulk 則必須手動 mbarrier_expect_tx。barrier 要等「所有 thread 已 arrive」且「所有 byte 已到」兩條件都成立才 flip。

步驟 5 的 fence_proxy_async 不可省：計算階段是透過 generic proxy 寫 shared，而 TMA 透過 async proxy 讀，必須用 proxy fence 把 generic 寫排在 async 讀之前，再以 __syncthreads() 讓所有 thread 的 fence 對發起 thread 0 生效。

1D 對齊需求：global 位址 16B、shared 位址 16B、barrier 位址 8B（cuda::barrier 保證）、傳輸大小須為 16B 倍數。

4.11.2.1.1 Prefetching（多階段預取）

在「複製-計算」迭代中，提前預取未來迭代的資料，可用當前迭代的計算掩蓋未來資料的搬移延遲，提高 bytes-in-flight。範例用 cuda::device::memcpy_async_tx 做 TMA 複製，並用 explicit phase tracking 的 shared memory barrier 同步：

• 初始化階段：每個 stage 一個 barrier（init(&bar[i], 1)），先載入前 num_stages 個 batch 到不同 shared 區段。
• 主迴圈：(a) mbarrier_try_wait_parity() 等當前 batch 到齊 → (b) 計算 → (c) 預取領先 num_stages 的下一筆 → (d) 用 rotating buffer 輪替 stage 並追蹤 barrier parity。

4.11.2.2 用 TMA 搬多維陣列（tensor map）

多維（最多 5D）與 1D 的關鍵差異：必須在 host 端建立 tensor map 並傳給 kernel。device code 用 cuda::ptx::cp_async_bulk_tensor 發起，完成機制與 1D 相同（讀用 barrier、寫用 bulk async-group）。

用 Driver API 建立 tensor map

cuTensorMapEncodeTiled 是 driver API，需直接連 -lcuda 或經 cudaGetDriverEntryPointByVersion 取得函式指標（見 04-CUDA-Features/25-Driver-Entry-Point-Access）：

PFN_cuTensorMapEncodeTiled_v12000 get_cuTensorMapEncodeTiled() {
  void* p = nullptr; cudaDriverEntryPointQueryResult st;
  cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &p, 12000,
                                   cudaEnableDefault, &st);   // version 12000
  return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(p);
}

建立需要：base 指標、各維 size（元素數）、stride（byte，且須 16B 倍數）、shared buffer 的 box_size（元素數）、elem_stride。注意最快變動維排第一。其他欄位包含資料型別、interleave、swizzle、L2 promotion（見 04-CUDA-Features/17-L2-Cache-Control）、OOB fill 等：

constexpr uint32_t rank = 2;
uint64_t size[rank]       = {GMEM_WIDTH, GMEM_HEIGHT};   // 最快維在前
uint64_t stride[rank-1]   = {GMEM_WIDTH * sizeof(int)};  // 須 16B 倍數
uint32_t box_size[rank]   = {SMEM_WIDTH, SMEM_HEIGHT};
uint32_t elem_stride[rank]= {1, 1};                      // 2 可只取複數的實部
cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_INT32, rank,
    tensor_ptr, size, stride, box_size, elem_stride,
    CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE,
    CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

Host-to-device：三種傳遞方式

使用 tensor map（2D tile 搬移）

device 端只給 tile 左上角座標 {x, y}，TMA 自動算位址：

__shared__ alignas(128) int smem_buffer[SMEM_HEIGHT][SMEM_WIDTH];  // 多維需 128B 對齊
int32_t coords[2] = { x, y };
ptx::cp_async_bulk_tensor(ptx::space_shared, ptx::space_global,
    &smem_buffer, &tensor_map, coords, cuda::device::barrier_native_handle(bar));
token = cuda::device::barrier_arrive_tx(bar, 1, sizeof(smem_buffer)); // 設 tx count

多維對齊需求：global 位址 16B、global stride 16B 倍數、shared 位址 128B、barrier 8B、傳輸大小 16B 倍數；global size 不必是 16B 倍數但須 ≥ 1。

4.11.2.2.1 / 4.11.2.2.4 在 device 上編碼 tensor map（template 模式）

當「典型的 __grid_constant__ 傳遞」不理想時（例如一個 kernel launch 要處理一批大小各異的 tensor），可在 device 上編碼 tiled-type tensor map。推薦流程：

host:   make_tensormap_template()  ── 用 cuTensorMapEncodeTiled 建最小 template
          │  (傳給 device：參數 / 指標 / __constant__ 皆可)
device: encode_tensor_map<<<1,32>>>(template, params, global_tensor_map)
          │  copy template → shared，tensormap_replace* 改欄位，cp_fenceproxy 寫回 global
device: consume_tensor_map<<<1,1>>>(global_tensor_map)
          │  fence acquire → cp_async_bulk_tensor 使用

Template via Driver API（4.11.2.2.4）：用 cuTensorMapEncodeTiled 建立一個 rank=1、UINT8 的最小可用 tensor map 當 template，之後在 device 修改其欄位即可。

4.11.2.2.2 device 端編碼與修改

步驟：① 把 template 傳進 kernel（任意方式）→ ② copy-initialize 到 shared 的 CUtensorMap（128B 對齊）→ ③ 用 cuda::ptx::tensormap_replace* 改欄位 → ④ 用 tensormap_cp_fenceproxy 從 shared 寫回 global 並 fence。

__launch_bounds__(32) __global__
void encode_tensor_map(const __grid_constant__ CUtensorMap template_map,
                       tensormap_params p, CUtensorMap* out) {
  __shared__ alignas(128) CUtensorMap smem_tmap;
  if (threadIdx.x == 0) {
    smem_tmap = template_map;                                  // copy template
    ptx::tensormap_replace_global_addressspace_shared, &smem_tmap, p.global_address;
    ptx::tensormap_replace_rankspace_shared, &smem_tmap, p.rank - 1; // 零基！
    // tensormap_replace_box_dim / global_dim / global_stride / element_size ...
  }
  __syncwarp();                                                // warp 內同步修改
  ptx::n32_t<128> bytes_128;
  ptx::tensormap_cp_fenceproxysem_release, ptx::scope_gpu, out, &smem_tmap, bytes_128;
}

要點：.rank 欄位採零基編號，要填「想要的 rank 減一」。tensormap_cp_fenceproxy 由 warp 內 threads 集體把 tensor map 複製到 global，同時以 release 語意讓更新對其他 thread（之後的 cp.async.bulk.tensor）可見。

4.11.2.2.3 使用被修改過的 tensor map（release-acquire）

與 __grid_constant__ 參數不同，使用 global memory 中的 tensor map 必須在「修改方」與「使用方」之間明確建立 tensor map proxy 的 release-acquire 關係：

 修改 thread ── tensormap_cp_fenceproxy(sem_release, ...) ──► global tensor map
                                                                 │
 使用 thread ── fence_proxy_tensormap_generic(sem_acquire,...) ◄─┘
                  │  (同 device 用 .gpu scope；跨 device 用 .sys)
                  ▼  __syncthreads() 後同 block 其他 thread 才可用
              cp_async_bulk_tensor(...)

__global__ void consume_tensor_map(CUtensorMap* tensor_map) {
  ptx::n32_t<128> size_bytes;
  ptx::fence_proxy_tensormap_generic(ptx::sem_acquire, ptx::scope_sys,
                                     tensor_map, size_bytes);   // acquire
  // 之後安全使用 tensor_map 發 cp_async_bulk_tensor
}

4.11.2.2.5 Shared-Memory Bank Swizzling

shared memory 有 32 個 bank，連續的 32-bit word 映到連續 bank，每 bank 每 cycle 32-bit。若同一 transaction 內重複用到同一 bank 即發生 bank conflict，頻寬下降。TMA 預設「shared 佈局 = global 佈局」，但這對某些存取模式會造成衝突。tensor map 可編入 swizzle mode，讓 TMA 存入 shared 時 'swizzle'、寫回 global 時 'unswizzle'。

矩陣轉置範例：8×8 的 int4（每元素 16B）row-major 載入，每組 8 條 thread 讀一列、寫到轉置 buffer 的一行 → 八路 bank conflict。改用 CU_TENSOR_MAP_SWIZZLE_128B（128B 正好對應一列長度）後，列存取與行存取都不再撞同一組 bank，衝突消除。轉置時用 swizzle 過的索引：

__shared__ alignas(1024) int4 smem_buffer[8][8];     // 128B swizzle 須 1024B 對齊
__shared__ alignas(1024) int4 smem_buffer_tr[8][8];
// ... cp_async_bulk_tensor 載入（與無 swizzle 寫法相同）...
for (int sidx_j = threadIdx.x; sidx_j < 8; sidx_j += blockDim.x)
  for (int sidx_i = 0; sidx_i < 8; ++sidx_i) {
    const int swiz_j     = (sidx_i % 8) ^ sidx_j;     // 異或得 swizzle 索引
    const int swiz_i_tr  = (sidx_j % 8) ^ sidx_i;
    smem_buffer_tr[sidx_j][swiz_i_tr] = smem_buffer[sidx_i][swiz_j];
  }
// host 端 encode 時改 swizzle 欄位為 CU_TENSOR_MAP_SWIZZLE_128B

swizzle 規則（CC 9）：

考量：swizzle 映射粒度固定 16B；shared box 的內維須 ≤ swizzle span；global 須 128B 對齊。若 shared buffer 未按「重複週期」對齊（但仍 128B 對齊），swizzle pattern 與 shared 間會有 row offset，需用 offset 公式校正：

// CU_TENSOR_MAP_SWIZZLE_128B 的偏移與索引關係
int offset = (reinterpret_cast<uintptr_t>(smem_ptr) / 128) % 8;
smem[y][((y + offset) % 8) ^ x] = ...;   // 64B 用 %4、32B 用 %2

要點：offset 反映 shared buffer 相對 swizzle pattern 被「位移」幾次，加在 row index y 上，再與 column index x 做 XOR 得到真正的 swizzled 索引。

考試/測驗重點

Related Notes

• 04-CUDA-Features/13-Async-Copies-LDGSTS
• 04-CUDA-Features/15-Async-Copies-STAS
• 04-CUDA-Features/11-Asynchronous-Barriers-Deep-Dive
• 04-CUDA-Features/12-Pipelines-Deep-Dive
• 04-CUDA-Features/07-Cooperative-Groups-Deep-Dive
• 04-CUDA-Features/17-L2-Cache-Control
• 04-CUDA-Features/25-Driver-Entry-Point-Access
• 03-Advanced-CUDA/07-Async-Data-Copies-and-L1-Config
• 03-Advanced-CUDA/06-Asynchronous-Barriers-and-Pipelines
• 03-Advanced-CUDA/08-CUDA-Driver-API
• 03-Advanced-CUDA/10-Tour-of-CUDA-Features
• 04-CUDA-Features/Practice-CUDA-Features
• 00-Dashboard/Exam-Traps