昨天的 TTIR 還看得出這個 kernel 在算 (64×32) @ (32×64)，但 tensor element 到底怎麼分給 GPU thread，這一層還看不到。今天繼續看同一次編譯的 TTGIR，看 SM121 target、warp layout、async copy、shared memory 與 dot lowering。

以下片段來自實際實驗的 TTGIR，為了方便閱讀我省略 loc(...)、重複 type 與不影響本段結論的 SSA value，保留 layout、memory descriptor 和 operation 主線。

想看完整版本可以直接開啟 matmul_kernel.ttgir。文中使用的 num_warps、shared memory 與 target 等編譯結果，則記錄在 matmul_kernel.json。

本篇大綱

• 對照 TTIR 與 TTGIR，找出 layout、shared memory 與 pipeline 從哪些 pass 出現。
• 確認 TTGIR 的 target 與 launch mapping。
• 讀 blocked、shared 與 NVIDIA MMA encoding。
• 追 A、B tile 從 global memory 到 shared memory。
• 用 IR 與 metadata 交叉計算 shared-memory 用量。
• 判斷 GB10 上這次實際選到哪條 Tensor Core 路徑。

.ttir → .ttgir

NVIDIA backend 在 CUDABackend.make_ttgir() 建立 pass manager。Triton 3.6.0 的完整 make_ttgir() 實作位於 compiler.py 第 247～318 行。其中第 255～267 行建立初始 GPU layout，SM121 會進入的 capability // 10 >= 10 分支在第 278～293 行，最後的 TMA、layout cleanup、fence 與 MMA lowering 則在第 296～314 行。

這次 target 是 capability 121，因此 capability // 10 >= 10 的 Blackwell 分支會進入 pipeline。可以把實際 pass 分成四組讀

表格依教學目的分組，實際執行順序請直接對照上方連結的 make_ttgir() 原始程式碼。其中 add_lower_mma 在第 311 行加入 pass manager，不是明天 make_llir() 才呼叫。

這個分支還會呼叫 TMEM 相關 pass，例如 add_hoist_tmem_alloc、add_promote_lhs_to_tmem 與 add_remove_tmem_tokens。但 pass 被執行不代表每支 kernel 都必定改寫。本次 TTGIR 是 NVIDIA MMA v2，metadata 的 tmem_size=0，表示這些 TMEM pass 沒有將本 kernel 轉成 TMEM/tcgen05 路徑。

add_tma_lowering 也會因 SM121 滿足 capability // 10 >= 9 而加入 pass manager。但這支 kernel 沒有 tensor descriptor，metadata 的 tensordesc_meta=[]，因此不能只因 pass 名稱就宣稱 artifact 已走 TMA 資料路徑。

三組前後對照

1. 純 shape 變成 GPU 分工

TTIR 只有

tensor<64x64xf32>

add_convert_to_ttgpuir 先加入 target 與基本 blocked encoding，matmul/layout passes 再把 accumulator 改成

tensor<64x64xf32, #ttg.nvidia_mma<{
  versionMajor = 2, warpsPerCTA = [2, 2], instrShape = [16, 8]
}>>

Shape 沒變，但 type 現在多了「哪些 warp/thread 持有哪些 element」的資訊。

2. tt.load 變成 global → shared → fragment

TTIR loop 內原本是

%a = tt.load %a_ptrs, %a_mask, %zero
%b = tt.load %b_ptrs, %b_mask, %zero
%next = tt.dot %a, %b, %acc

經 matmul acceleration、scheduling 與 pipeline passes 後，TTGIR 出現

%a_smem = ttg.local_alloc : () -> !ttg.memdesc<1x64x32xf16, ...>
%token = ttg.async_copy_global_to_local %a_ptrs, %slot mask %a_mask
%ready = ttg.async_wait %token {num = 0 : i32}
%a_fragment = ttg.local_load %slot token %ready
%next = tt.dot %a_fragment, %b_fragment, %acc

這是本篇最重要的變化：tt.dot 還沒有變成 machine instruction，但它的輸入已改為從 swizzled shared memory 讀出的 MMA fragment。

3. 一般 tensor 寫回前多了 layout conversion

TTIR 的 FP16 tensor 可直接交給 tt.store。TTGIR 中，accumulator 由 #mma layout 持有，global store 則期待 #blocked，因此必須插入

%c_blocked = ttg.convert_layout %c
    : tensor<64x64xf16, #mma> -> tensor<64x64xf16, #blocked>

add_remove_layout_conversions 會嘗試移除或合併多餘 conversion，最終還留著這一個，表示輸出 ownership 的轉換不能在這層消除。明天可以繼續查它的實作成本。

先看 GPU target

Module attribute 寫得很直接

module attributes {
  "ttg.num-ctas" = 1 : i32,
  "ttg.num-warps" = 4 : i32,
  ttg.target = "cuda:121"
}

cuda:121 對應 GB10 的 compute capability 12.1。num-warps=4 代表一個 CTA 使用 128 threads。

Layout encoding 是資料分工表

TTGIR type 不只保留 shape，還多了 encoding。這次 artifact 可看到

#blocked = #ttg.blocked<{
  sizePerThread = [1, 8],
  threadsPerWarp = [4, 8],
  warpsPerCTA = [4, 1],
  order = [1, 0]
}>
#mma = #ttg.nvidia_mma<{
  versionMajor = 2,
  versionMinor = 0,
  warpsPerCTA = [2, 2],
  instrShape = [16, 8]
}>
#shared_a = #ttg.swizzled_shared<{
  vec = 8, perPhase = 2, maxPhase = 4, order = [1, 0]
}>
#shared_b = #ttg.swizzled_shared<{
  vec = 8, perPhase = 1, maxPhase = 8, order = [1, 0]
}>

先看 #blocked

• sizePerThread=[1,8]：這個 layout 中，每個 thread 持有的基本資料塊沿兩個 tensor 軸排成 1×8。
• threadsPerWarp=[4,8]：32 個 lane 以 4×8 排列。
• warpsPerCTA=[4,1]：四個 warp 在 CTA 內沿第一個 tensor 軸排列。
• order=[1,0]：axis 1 是較快變動的軸。

這些欄位要合起來看。sizePerThread 不能單獨理解成一個 thread 最後只算 8 個元素，因為 operation 可能在同一個 thread 中持有多個 fragment。

#mma 描述 accumulator 如何由 Tensor Core fragment 與四個 warp 共同持有。warpsPerCTA=[2,2] 表示四個 warp 以 2×2 方式合作處理 M、N 方向。#shared_a 與 #shared_b 則描述 A/B tile 放進 shared memory 後的 swizzle，兩者 shape 與存取方向不同，因此參數也不同。

TTGIR 把資料 ownership 寫進 tensor type 的 encoding。兩個 tensor 即使 shape 相同，只要 encoding 不同，就可能需要 ttg.convert_layout。這項轉換可能使用 register shuffle，也可能需要 shared memory 與 synchronization。

A、B tile 真的進入 shared memory

這次可以找到兩個 local allocation

%a_smem = ttg.local_alloc
  : () -> !ttg.memdesc<
      1x64x32xf16, #shared_a, #smem, mutable>
%b_smem = ttg.local_alloc
  : () -> !ttg.memdesc<
      1x32x64xf16, #shared_b, #smem, mutable>

!ttg.memdesc 表示一塊有 layout 的 memory region；#smem 表示 shared memory；mutable 表示這塊區域會被 async copy 寫入。前面的 1 是 pipeline buffer 軸。就 element 數量計算

A tile：64 × 32 × 2 bytes = 4096 bytes
B tile：32 × 64 × 2 bytes = 4096 bytes
合計：                         8192 bytes

"shared": 8192

這和 metadata 的一致，這個數字有 IR allocation、dtype 大小與 compiler metadata 三份證據可以交叉驗證。

Global → shared → register fragment

先看 A tile 的 pipeline prologue

%a_slot = ttg.memdesc_index %a_smem[%c0]
%a_token = ttg.async_copy_global_to_local
    %a_ptrs, %a_slot mask %a_mask other %zero
    {contiguity = 8 : i32}
%a_group = ttg.async_commit_group tokens %a_token

這三行分別表示

1. 從 A 的 shared-memory allocation 取出目前 buffer slot。
2. 把 global-memory pointer tensor 指向的資料搬進該 slot。
3. 把這次 async copy 提交成一個 group。

B 也有相同結構，只是使用 32×64 與另一套 shared layout。進入 K loop 後，consumer 端會等待資料，再載入 dot fragment

%ready = ttg.async_wait %a_group, %b_group {num = 0 : i32}
%a_view = ttg.memdesc_index %a_smem[%slot]
%a_fragment = ttg.local_load %a_view token %ready
%b_view = ttg.memdesc_index %b_smem[%slot]
%b_fragment = ttg.local_load %b_view token %ready
%next_acc = tt.dot %a_fragment, %b_fragment, %acc

現在可以把資料路徑完整列出來

global pointer tensor
  → async copy
swizzled shared-memory tile
  → async wait
dot operand fragment
  → tt.dot
MMA-layout FP32 accumulator

Shared encoding 還帶有 swizzle/layout 資訊，用來配合 matrix load 並減少 shared-memory bank conflict。只看 Python 的 tl.load 不會看到這些安排。

num_stages=2 在哪裡

Python launch 指定

num_stages=2

TTGIR 中已經出現 async copy、buffer indexing、等待與同步相關 operation。Loop 在計算目前 tile 後，也準備下一輪 pointer 與 copy

%next_a_ptrs = tt.addptr %a_ptrs_iter, %a_step
%next_b_ptrs = tt.addptr %b_ptrs_iter, %b_step
%next_a = ttg.async_copy_global_to_local
    %next_a_ptrs, %a_slot mask %next_a_mask other %zero
%next_b = ttg.async_copy_global_to_local
    %next_b_ptrs, %b_slot mask %next_b_mask other %zero

這是 software pipeline 的主要形狀，consumer 計算目前 tile，producer 準備下一個 tile。不過 stage 增加會占用更多 shared memory 或 register，不能只靠越深越好判斷。

本例的 K loop 有八次 iteration。到 LLIR/PTX 後，pipeline 會變成更低階的 cp.async、commit、wait 與 barrier。

GB10 為什麼不會使用 tcgen05

Blackwell 底下有不同的 instruction families。tcgen05 與 Tensor Memory（TMEM）屬於支援這組能力的資料中心 Blackwell target；PTX ISA 的 Target ISA Notes 沒有把它列給 SM120 / SM121。GB10 的 target 是 SM121，因此不能把資料中心 Blackwell 的 tcgen05 kernel 直接拿來使用。

這次 TTGIR 的實際 encoding 是

#ttg.nvidia_mma<{
  versionMajor = 2,
  versionMinor = 0,
  ...
}>

Metadata 同時顯示

"tmem_size": 0

後面的 PTX 是 mma.sync.m16n8k16，SASS 則是 HMMA.16816.F32。這四項證據彼此一致

SM121 target
  → TTGIR：NVIDIA MMA v2
  → metadata：tmem_size=0
  → PTX：mma.sync
  → SASS：HMMA

tmem_size=0 是這次 codegen 沒有配置 TMEM 的結果證據，不是 compiler 放棄 tcgen05 的原因。根本限制來自 SM121 可用的 instruction set。GB10 仍有其他 Blackwell 能力，但不支援 tcgen05／TMEM。

寫回前為什麼要 convert_layout

Dot accumulator 使用 #mma layout，但一般 global store 使用 #blocked layout。TTGIR 在輸出端明確插入

%c_f16 = arith.truncf %acc
    : tensor<64x64xf32, #mma>
      to tensor<64x64xf16, #mma>
%c_blocked = ttg.convert_layout %c_f16
    : tensor<64x64xf16, #mma>
      -> tensor<64x64xf16, #blocked>
tt.store %c_ptrs, %c_blocked, %c_mask

convert_layout 表示 accumulator 的 thread ownership 不適合直接執行目前的 global store。這類轉換是效能分析的重要訊號，它可能 lower 成 register shuffle，也可能需要 shared memory。要到 LLIR/PTX 繼續確認成本。

TTGIR 的效能檢查清單

清單要和 IR 證據一起使用

TTGIR 已經回答 tile 如何映射到 GPU，但還不是 GPU 可執行的指令。明天我們來看看 async copy、barrier、matrix load 與 MMA 如何進入 LLVM/NVVM。

參考資料

• Triton GPU dialect documentationhttps://triton-lang.org/main/dialects/TritonGPUOps.html
• Triton NVIDIA backend compilerhttps://github.com/triton-lang/triton/blob/v3.6.0/third_party/nvidia/backend/compiler.py
• CUDA C++ Programming Guide：shared memory 與 asynchronous data copieshttps://docs.nvidia.com/cuda/cuda-c-programming-guide/
• PTX ISA：target-specific features 與 tcgen05https://docs.nvidia.com/cuda/archive/13.0.2/parallel-thread-execution/