昨天的 TT-Metal 範例要手動安排 reader、compute、writer、circular buffer 與 NoC synchronization。今天看 TT-MLIR 怎麼用多個 MLIR dialect 串接 tensor graph、TTNN operation、TT-Metal host program 與 device kernel。

今天用 TT-MLIR 官方 Overview 的架構圖與 IR 範例。我們會先用同一個 simple_linear 比較 TTIR 與 TTNN，再用官方 simple_matmul 追蹤 D2M、TTKernel 與 TTMetal 新增了哪些執行資訊，文中架構以目前commit 為主。

本篇大綱

• 先確認 TT-MLIR 在 Tenstorrent software stack 的位置。
• 用官方架構圖分清 TTNN 與 D2M 兩條 backend 路徑。
• 用官方範例比較每次轉換加入了哪些 layout、grid、buffer、kernel 與 host command 資訊。
• 看 ttmlir-opt 如何選擇 TTNN 或 TTMetal pipeline。
• 分清 compiler IR、FlatBuffer 與 runtime support 的狀態

TT-MLIR 放在哪裡

TT-MLIR 定義 Tenstorrent 相關的 dialect、transformation pass、binary interface 與 runtime component，也可以作為其他 compiler 專案取用的元件庫。

圖片來源：TT-MLIR Architecture & Dialect Overview

這張圖可以看到分成兩個主要方向

• TTIR → TTNN → TTNN FlatBuffer → TTNN runtime：保留高階 TTNN API operation，交給 tt-nn 執行。
• TTIR → D2M → TTKernel + TTMetal：把 generic compute、data movement、device kernel 與 host program 展開。

TTNN 和 D2M 之間的雙向箭頭，表示它們是可供 compiler 組合的不同抽象

昨天畫過的 software stack 可以再細化成

PyTorch / JAX / ONNX 等 model source
  ↓
StableHLO 或其他 front-end dialect
  ↓
TT-MLIR
  TTIR
  ├─→ TTNN dialect → TTNN FlatBuffer → TTNN runtime
  └─→ D2M → TTKernel + TTMetal → TTMetal FlatBuffer
  ↓
tt-metal / Tenstorrent hardware

根據這個版本的官方 FlatBuffers 文件，.ttb 所對應的 TTMetal backend runtime 仍標示為 unsupported，架構圖表達的是 compiler 元件與介面，不代表圖中每條 runtime 路徑都有相同的完整度。

五個 dialect 各自負責什麼

這五個 dialect 可以對回前兩天的硬體與 TT-Metal 概念

tensor 要做什麼
  → ttir / ttnn

tile 長什麼樣、放 DRAM 還是 L1、映射到哪個 grid
  → d2m

reader / compute / writer 的 circular buffer、NoC 與 tile operation
  → ttkernel

host 如何配置 buffer 並啟動 program
  → ttmetal

MLIR dialect 讓這些決策在不同階段仍有明確名稱。Compiler pass 可以先處理 tensor graph，後面才將 layout、memory space、grid、kernel synchronization 與 host dispatch 逐步寫進 IR。

每層 IR 到底改了什麼

可以把 lowering 看成用更多執行細節，換掉一部分高階名稱。例如 ttir.linear 在 TTNN 路徑還會保留為 ttnn.linear，走 D2M / TTKernel 路徑時，後面讀到的是 reduction iterator、tile matmul、CB 同步與 program dispatch。它們共同實現 matmul，但 IR 不再依賴單一個高階 matmul operation 來表達整個執行過程。

從 ttir.linear 讀高階 tensor 語意

官方 Overview 用 simple_linear 當例子。TT-MLIR TTIR 階段可以簡化成

func.func @simple_linear(
  %arg0: tensor<64x128xbf16>,
  %arg1: tensor<128x64xbf16>,
  %bias: tensor<64x64xbf16>) -> tensor<64x64xbf16> {
  %0 = ttir.empty() : tensor<64x64xbf16>
  %1 = "ttir.linear"(%arg0, %arg1, %bias, %0)
      : (...) -> tensor<64x64xbf16>
  return %1 : tensor<64x64xbf16>
}

這個階段已經知道 input、weight、bias 與 output shape，也保留 linear 的高階語意。它還沒指定實際的 DRAM / L1 layout、core grid、circular buffer 與 NoC transfer。

這樣的 IR 適合做 graph-level transformation。例如 compiler 可以在 linear 還是 named op 時做 decomposition、fusion 或 layout-related rewrite，不用從一大段低階 loop 與記憶體操作重新推回它的 tensor 語意。

Lower 到 TTNN 後多了什麼

同一個 simple_linear lower 到 TTNN dialect 後，operation 變成 ttnn.linear，tensor type 開始帶 TTNN layout

#layout = #ttnn.ttnn_layout<
  (d0, d1) -> (d0, d1),
  <1x1>,
  memref<2x4x!ttcore.tile<32x32, bf16>, #dram>,
  <interleaved>>

%0 = "ttnn.linear"(%arg0, %arg1, %arg2) <{
  transpose_a = false,
  transpose_b = false
}> : (...) -> tensor<64x64xbf16, #layout>

這裡可以看到幾個新資訊

• Logical tensor shape 是 64x64 BF16
• Physical representation 用 32x32 tile
• Memory space 是 DRAM
• Layout 是 interleaved
• linear 已經對到 TTNN API 的 operation 與 attribute

TTNN IR 裡還會出現 ttnn.deallocate，表示 compiler / runtime 介面已經需要處理 tensor lifetime。TTIR 的高階計算圖不需要暴露這些 runtime 細節，降階到 TTNN 後就必須把它們具體化

把官方兩段 simple_linear 並排後，差異可以濃縮成三點

1. 保留： function 的 logical shape、BF16 dtype 與 linear 語意都還在。
2. 新增： 每個 tensor 都帶有 #ttnn.ttnn_layout，ttnn.linear 加入 transpose_a / transpose_b，並且出現 ttnn.deallocate。
3. 改寫： TTIR 先用 ttir.empty 建立 output tensor；官方 TTNN 範例由 ttnn.linear 直接產生結果，所以原本的 output empty 不再出現。

這也說明 lowering 不只是把 operation prefix 從 ttir 換成 ttnn。Tensor type、operation attribute、結果產生方式與資源生命週期都改了。

接著換成官方 simple_matmul 範例

官方 Overview 的 TTIR 與 TTNN 分頁使用 simple_linear，D2M、TTKernel 與 TTMetal 分頁則改用 simple_matmul。下面採用第二個範例觀察抽象的變化，官網沒有提供同一份 simple_linear 逐行 lowering 到 TTMetal 的紀錄。

simple_matmul 的 logical tensor shape 是

A: 64x128
B: 128x96
C: 64x96

D2M：從一個 matmul 展開成 tiled generic compute

D2M 先用 d2m.to_layout 將三個 logical tensor 轉成 L1 內的 physical tiled tensor

64x128  → 1x1x2x4 tiles
128x96  → 1x1x4x3 tiles
64x96   → 1x1x2x3 tiles

每個 tile = 32x32 bf16

這些數字可以直接對回 shape，64 / 32 = 2、128 / 32 = 4、96 / 32 = 3。接著，d2m.generic 寫進執行網格、indexing map 與 iterator type。下面保留官方範例的關鍵 operation，省略完整 type signature

// 根據官方範例縮寫，不能直接執行
%7 = d2m.generic {
  grid = #ttcore.grid<1x1>,
  indexing_maps = [#map_a, #map_b, #map_c],
  iterator_types = [parallel, parallel, reduction],
  threads = [#d2m.thread]
} ... {
  %a = d2m.wait %cb0
  %b = d2m.wait %cb1
  %c = d2m.reserve %cb2
  %result = d2m.tile_matmul %a, %b, %c
  ...
}

兩個 parallel iterator 對應輸出的 tile，reduction iterator 對應 matmul 的 K 軸。wait 表示等輸入 CB 可讀，reserve 表示預留輸出 CB 空間。到了這一層，輸入是否就緒、輸出是否有位置可寫，已經成為 IR 的一部分。

TTKernel：展開成 device kernel 的同步與 tile instruction

官方 TTKernel 範例將 kernel 標記為 NoC thread 或 compute thread，並把 CB 編號作為 compile-time argument。下面是官方 compute kernel 的關鍵 operation，參數已經縮寫：

// 參數已縮寫，不能直接執行
ttkernel.cb_wait_front(%cb0, %one)
ttkernel.cb_wait_front(%cb1, %one)
ttkernel.cb_reserve_back(%cb2, %one)

ttkernel.tile_regs_acquire()
ttkernel.mm_block_init_short(...)
ttkernel.experimental.matmul_block(...)
ttkernel.pack_tile(...)

ttkernel.cb_pop_front(%cb0, %one)
ttkernel.cb_pop_front(%cb1, %one)
ttkernel.cb_push_back(%cb2, %one)

D2M 的 wait / reserve 現在展開成 CB front / back 同步，compute 也顯式使用 tile register、matmul instruction 與 pack。這些 operation 的形狀已經很接近昨天手寫的 TT-Metal compute kernel。

TTMetal：把 device kernels 包成 host 可啟動的 program

TTMetal dialect 處理 host 端編排，官方範例不只出現 enqueue_program，它的 attribute 還會指定處理器核範圍、NoC0 / NoC1 data-movement kernel、compute kernel 與 CB port。下面的縮寫版省略了完整 memref type 與 kernel argument

// 根據官方範例縮寫，不能直接執行
%buffer = ttmetal.create_buffer ...
ttmetal.enqueue_write_buffer ...

ttmetal.enqueue_program ... <{
  cb_ports = [0, 1, 2],
  kernelConfigs = [
    noc_config<@reader,  core_range<0x0, 2x3>, noc0>,
    noc_config<@writer,  core_range<0x0, 2x3>, noc1>,
    compute_config<@compute, core_range<0x0, 2x3>>
  ]
}>

ttmetal.enqueue_read_buffer ...
ttmetal.finish

此時的 IR 可以回答哪些 device kernel 在哪個處理器核網格上執行，也能表示 buffer address、semaphore 與 kernel argument。高階 matmul 語意已經分散到 buffer、reader、compute、writer 與 host command 之中。

這就是 TT-MLIR 和 TT-Metal 之間的具體連結，昨天從 API 和 kernel 分工往上看，今天從 tensor graph 往下追，兩邊會在 circular buffer、NoC operation、tile compute 與 host dispatch 匯合。

用 ttmlir-opt 觀察 pipeline

TT-MLIR 提供 ttmlir-opt 當 optimizer driver，用來對 .mlir 檔執行 compiler pass。官方文件給了兩條可直接對照的指令

./build/bin/ttmlir-opt \
  --ttir-to-ttnn-runtime-pipeline \
  test/ttmlir/Dialect/TTNN/simple_multiply.mlir

./build/bin/ttmlir-opt \
  --ttir-to-ttmetal-pipeline \
  test/ttmlir/Dialect/TTNN/simple_multiply.mlir

兩道指令的 input 可以相同，pipeline option 決定後面保留哪組 dialect 與要對接的 runtime 形狀。這也是 MLIR multi-level design 的一個具體例子，前面可以共用 TTIR 的 tensor semantics，後面依 backend 需求分流。

FlatBuffer 是 compiler 和 runtime 的邊界

TT-MLIR 的 Target 目錄定義 FlatBuffer schema，專案文件將它描述為 compiler 和 runtime 之間的 binary interface，目前列出三種副檔名

.ttsys 很適合用來理解 cross-compilation。Compiler 可以從 target machine 收集 system description，再將它帶回 development machine。這讓 target topology 和 compiler IR 保持分開，也讓上層 pipeline 不必直接查詢正在執行的裝置。

目前的限制

從官方 IR 與 pipeline 文件，可以確認 TT-MLIR 如何分開 high-level tensor IR、TTNN op、generic data movement、device kernel 與 host dispatch。這些材料還不足以回答

• 某個 model 能否在特定 Wormhole / Blackhole 系統上完整編譯與執行。
• TTNN 和 TTMetal 兩條 pipeline 對同一個 workload 的效能差異。
• Layout、sharding 與 grid 在特定 shape 下的最佳選擇。
• 官方 Overview 裡的每一條降階路徑是否都有完整 runtime support。

這些問題需要版本、build configuration、裝置、compiler output 和實測記錄

複習一下

今天用官方架構圖、simple_linear、simple_matmul、ttmlir-opt 和 FlatBuffer 把 TT-MLIR 的 compiler flow 走過一次

• TTIR 之後主要分成 TTNN 與 D2M 兩個方向，整體流程不能畫成單一直線
• ttir 保留高階 tensor compute graph
• ttnn 對應 TTNN API，並加入 physical tile、memory space、layout 與 tensor lifetime
• ttcore 提供 tile、grid、layout 與 memory space 等共用定義
• d2m 將 matmul 展開成 tiled generic compute、iterator、grid 與 CB data movement
• ttkernel 把 CB 同步與 tile instruction 寫進 device kernel，ttmetal 負責 host buffer、kernel config 與 program dispatch
• .ttnn 和 .ttb 代表不同 backend 的 binary interface，這個版本的 .ttb 仍標示 unsupported

明天回到 Triton，之前介紹主要是 Trtion GPU 的 compile 流程，但最近對於 Triton 在不同 NPU 上面是如何 compile 很有興趣，會整理目前有哪些專案在實作，看它們從 Triton 的哪一層 IR 分流，以及如何接進各自的 backend dialect 與 runtime。

參考資料

• Tenstorrent TT-MLIR repository
  https://github.com/tenstorrent/tt-mlir
• TT-MLIR Overview
  https://docs.tenstorrent.com/tt-mlir/overview.html
• TT-MLIR Dialects Overview
  https://docs.tenstorrent.com/tt-mlir/dialects-overview.html
• TT-MLIR ttmlir-opt
  https://docs.tenstorrent.com/tt-mlir/ttmlir-opt.html
• TT-MLIR FlatBuffers
  https://docs.tenstorrent.com/tt-mlir/flatbuffers.html
• TT-MLIR Project Structure
  https://docs.tenstorrent.com/tt-mlir/project-structure.html