昨天的 LLVM IR 已經出現 cp.async、ldmatrix 與 mma.sync。今天沿用同一份 artifact，看 LLVM NVPTX backend 產生的 PTX，再真的呼叫一次 CUDA 13.0 的 ptxas。一路追到這裡，終於可以拿到 physical register 與 spill 數字了。

完整版可以參考 matmul_kernel.ptx 與 Triton 產生的 matmul_kernel.cubin 已放在 GitHub，本文講解部分片段的 PTX kernel entry、資料搬移與 MMA 。

本篇大綱

• 對照 LLIR 與 PTX，分清 LLVM backend 與 ptxas 各自負責的轉換
• 讀 PTX header、kernel entry 與 thread 數
• 從 PTX 找出資料搬移與 Tensor Core instruction
• 分清 PTX virtual register 與 physical register
• 用 ptxas -v 重新組譯並讀 resource log
• 說明手動產生的 cubin 為何不保證與 Triton cache 逐 byte 相同

.llir → .ptx

這裡已經不是 MLIR pass pipeline，make_llir() 完成後，CUDABackend.make_ptx() 取得 target triple、SM processor 與 feature，然後呼叫 llvm.translate_to_asm()。Triton 3.6.0 的完整 make_ptx() 實作在 NVIDIA backend compiler.py 第 435～459 行，實際呼叫 LLVM NVPTX backend 的程式碼在第 438～442 行。

llvm.translate_to_asm(
    src,
    "nvptx64-nvidia-cuda",
    "sm_121a",
    features,
    ["nvptx-mad-wide-opt"],
    ...,
)

這裡 LLVM NVPTX backend 將 LLVM module 產生 PTX assembly text，Triton 之後再用少量文字後處理修正 .version、.target，並在不需要時移除會阻擋 ptxas 最佳化的 debug flag。

LLIR 和 PTX 到底差在哪裡

這也解釋了為什麼 LLIR 與 PTX 都看得到 cp.async 或 mma.sync，Triton 在 LLIR 已經藉由 NVVM intrinsic 或 inline asm 選定了部分 NVIDIA-specific operation，LLVM backend 再完成其餘 instruction selection、register 虛擬化與 PTX 文字產生。

.ptx → .cubin

Triton 不會用 LLVM 直接寫出本次的 cubin，Triton 3.6.0 的完整 make_cubin() 實作在第 461～535 行。這個函式會把 PTX 寫入暫存檔，組出 ptxas command，執行後再讀回 binary bytes。ptxas_cmd 的參數排列在第 469～492 行：

PTX text
  → temporary .ptx
  → ptxas -lineinfo -v --gpu-name=sm_121a ... -o ...
  → cubin bytes

enable_fp_fusion=false 時會多加 --fmad=false；關閉 ptxas 最佳化時會加 --opt-level 0；ptx_options 也會被附加到 command。文中列出的 command 只對應本次設定。其他 Triton 編譯可能會加入不同 option。

PTX 的閱讀順序

PTX 已經比 LLVM IR 接近 GPU，但仍是 virtual ISA。第一次閱讀可以照以下順序

header / target
→ kernel parameter 與 thread 數
→ memory instruction
→ synchronization
→ matrix instruction
→ store
→ ptxas resource report

先用 instruction mnemonic 建立主線，再追 %r、%rd register。如果一開始就追所有 register 編號，很容易把 address calculation 和主要計算混在一起。

先看 PTX header

這次實際輸出的開頭是

.version 8.8
.target sm_121a
.address_size 64

.extern .shared .align 16 .b8 global_smem[];
.visible .entry matmul_kernel(
    .param .u64 .ptr .global matmul_kernel_param_0,
    .param .u64 .ptr .global matmul_kernel_param_1,
    .param .u64 .ptr .global matmul_kernel_param_2,
    ...
)
.reqntid 128
{
    .reg .pred %p<4>;
    .reg .b32  %r<196>;
    .reg .b64  %rd<32>;

幾個資訊可以直接核對前面的 artifact

• .target sm_121a：針對 SM121 的 architecture-accelerated target
• .visible .entry matmul_kernel：保留 runtime 搜尋的 kernel symbol
• .reqntid 128：四個 warp，每個 warp 32 threads
• .extern .shared ... global_smem[]：dynamic shared memory，大小由 launch configuration 傳入
• .param .u64 .ptr .global：64-bit global-memory pointer parameter
• .address_size 64：使用 64-bit device address

Metadata 常用 arch: sm121 表示 architecture family，PTX 則寫成 sm_121a。尾端的 a 表示這份 PTX 使用 architecture-accelerated feature，不能把它當成可自由前向相容的純 virtual target

PTX 中的 matmul 主線

1. Global memory 搬到 shared memory

cp.async.cg.shared.global
    [%r101 + 0], [%rd3 + 0], 0x10, %r23;
cp.async.cg.shared.global
    [%r102 + 0], [%rd4 + 0], 0x10, %r23;
cp.async.commit_group;
cp.async.wait_group ...;

這一行可以按 operand 讀

%r101：shared-memory destination address
%rd3：global-memory source address
0x10：預計搬 16 bytes
%r23：實際有效 bytes，和 boundary mask 有關

它們來自 TTGIR 的 async_copy_global_to_local，並延續 LLIR 的 inline asm / NVVM intrinsic。commit_group 與 wait_group 負責 software pipeline 的 producer/consumer ordering。

2. Shared memory 讀成 matrix fragment

ldmatrix.sync.aligned.m8n8.x4.shared.b16
    {%r69, %r70, %r71, %r72}, [%r14];
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16
    {%r73, %r74, %r89, %r90}, [%r18];

ldmatrix 讓 warp 按 matrix instruction 所需格式，從 shared memory 載入 A、B fragment。x4 表示一次載入四個 8×8 matrix fragment 對應的 register 資料，b16 表示 16-bit element。
第二條有 trans，對應 B operand 的轉置讀法。

3. 執行 Tensor Core MMA

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
    {%r163, %r164, %r165, %r166},
    {%r69, %r70, %r71, %r72},
    {%r73, %r74},
    {%r163, %r164, %r165, %r166};

這次 unrolled body 中共有 16 條同型 mma.sync。

指令名稱也交代了

• instruction tile：m16n8k16
• A 為 row-major、B 為 column-major fragment
• A/B 輸入是 FP16
• accumulator 與結果是 FP32

四組 operand 依序是 D、A、B、C，對應

D = A × B + C

第一輪某些 accumulator 會從 0 開始，後續指令再把先前 %r163...%r166 當成 C 輸入並覆寫成新的 D。完整 64×64 tile 需要多條 m16n8k16 指令與多個 warp fragment 合作完成，所以不能把一條 mma.sync 當成整個 Triton tt.dot。

最後輸出包含 vectorized store

st.global.v4.b32 ...;

一次寫入四個 32-bit word，對應 8 個 FP16 element。

.reg %r<196> 不是用了 196 個實體 register

PTX 是 virtual ISA，這幾行只宣告 virtual register namespace

.reg .pred %p<4>;
.reg .b32  %r<196>;
.reg .b64  %rd<32>;

它們有不同 bit width，lifetime 也可能不重疊。Physical register allocation、spill 與 instruction scheduling 由 ptxas 完成，不能把 %r<196> 直接拿去算 occupancy。

這份 PTX 中

• %p 保存 predicate，常用於 mask 和 branch
• %r 是 32-bit virtual register，可裝 integer、打包 FP16 或 FP32 bit pattern
• %rd 是 64-bit virtual register，主要保存 global-memory address

PTX 使用 untyped bit register 的風格，所以 %r163 出現在 FP32 MMA operand 時，應依 instruction signature 解讀為 FP32 值，不能只看名字中的 .b32 宣告。

實際跑 ptxas -v

GB10 上使用 CUDA 13.0

ptxas -v -arch=sm_121a kernel.ptx -o kernel.cubin

實測重點如下

ptxas info    : Compiling entry function 'matmul_kernel' for 'sm_121a'
ptxas info    : Function properties for matmul_kernel
    0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 80 registers, used 1 barriers
ptxas info    : Compile time = 23.414 ms

這段文字是當次實驗紀錄，目前 repository 沒有保存完整原始 ptxas -v log，所以 23.414 ms 與 used 1 barriers 無法只靠現有檔案重新驗證，80 registers 和沒有 spill 則能從 cubin resource metadata 取得證據。

這個 kernel 每個 thread 使用 80 個 physical registers，而且這次沒有 register spill。
另外 register、shared memory、warp 數量和硬體上限會一起限制 occupancy，最後還是要量 kernel。

Triton 實際怎麼呼叫 ptxas

Triton 3.6.0 的 NVIDIA backend 會建立類似以下指令

ptxas
  -lineinfo
  -v
  --gpu-name=sm_121a
  <temporary-input.ptx>
  -o <temporary-output.cubin>

實際選項還會受到 debug、CUDA toolkit 與 Triton build 影響，查 compiler 行為時，要看目前安裝版本的 NVIDIA backend 。

為何手動 cubin 不一定逐 byte 相同

這次用 -v -arch=sm_121a 手動組出的 cubin，SHA-256 與 Triton cache 中的 cubin 不同；補上 -lineinfo 後仍不同。反組譯可看到主要 instruction family 相同，但部分 scheduling/control encoding 不完全一致。

這也是當次實驗紀錄，Repository 目前只保存 Triton 產生的 cubin，沒有保存手動組譯版本、兩份 SHA-256 或 binary diff，因此不能從現有 artifact 獨立重做這項比較。

因此手動 ptxas 的用途是

• 驗證 PTX 能否針對目標架構組譯
• 取得 register、spill、stack 與 barrier 資訊
• 嘗試不同 assembler option
• 對照 SASS instruction family

它不自動構成精確重現 Triton cubin 的證明。如果要 bitwise reproduction，必須完整保存 Triton 版本、ptxas binary、所有 option、環境與輸入 bytes，不能只憑一條近似 command。

PTX 與 ptxas 各回答什麼

cp.async：global → shared
ldmatrix：shared → MMA fragment registers
mma.sync：A/B fragment + FP32 accumulator
stmatrix / st.global：重排後寫回 C
ptxas：配置 80 個 physical registers 並產生 cubin

明天會看 cubin ，還有 cuobjdump、nvdisasm 與 Nsight Systems 完成最後一哩路。

參考資料

• PTX ISA
  https://docs.nvidia.com/cuda/parallel-thread-execution/
• CUDA Binary Utilities
  https://docs.nvidia.com/cuda/cuda-binary-utilities/
• LLVM NVPTX backend guide
  https://llvm.org/docs/NVPTXUsage.html
• Triton NVIDIA backend make_ptx / make_cubin
  https://github.com/triton-lang/triton/blob/v3.6.0/third_party/nvidia/backend/compiler.py