前天我們看到 NVIDIA GPU 架構，提到 SM、warp、shared memory 與 Tensor Core，昨天看到 CUDA、CuTeDSL、cuTile 和 Triton 如何描述 GPU kernel。
我們之後主要會研究 Triton 的編譯流程，會經過 TTIR、TTGIR、LLVM IR、PTX 和 cubin。IR 是 intermediate representation（中間表示），可以先把它想成 compiler 在原始程式與機器指令之間使用的程式形式，在實際碰這些 IR 以前，今天先讀 Chris Lattner 等人撰寫的 MLIR 論文。我好奇為什麼現代 compiler 會需要多個抽象層級，以及這些 IR 為什麼需要共用一套 infrastructure。
MLIR: Scaling Compiler Infrastructure for Domain Specific Computation
今天主要依照 2020 年的 arXiv 版本閱讀，因為它完整交代了 MLIR 的起源、設計原則和早期案例，但是論文中的部分 dialect 名稱與語法已經過時。Dialect 是針對某個領域定義的一組 IR 詞彙與規則；例如當時常見的 std dialect 後來被拆分。現在寫 MLIR 應以官方文件和目前使用的 LLVM 版本為準。
大家常聽到的 Moore's Law 是 Intel 共同創辦人 Gordon Moore 在 1965 年對半導體產業的觀察，可容納的電晶體數量，大約每隔 18 至 24 個月就會增加一倍。
過去製程進步時，電壓與單顆電晶體的功耗也會下降，CPU 因此可以提高時脈頻率。很多程式不用重寫，換到新一代 CPU 就能跑得更快。
但是到了 2000 年代中期以後，電壓已經很難繼續降低，時脈頻率再往上提會遇到功耗與散熱限制，所以 End of Moore's Law 是說電晶體數量仍然增加，但程式不再只靠新製程就自然變快。
之後開始往多核、GPU 和各種專用加速器。像是前天看到的 SM、warp、shared memory 和 Tensor Core，就是 GPU 為平行運算與矩陣乘法設計的專用硬體。
專用硬體提供了更高效能的可能，程式卻不會自動知道如何使用它，所以 Compiler 要做的事情就是找出哪些運算可以平行執行、資料要如何切成 tiles、要放在 register 還是 shared memory，以及最後要選用哪些硬體指令。這些決定如果沒有做好，再多的運算單元也可能閒置，資料搬移也可能比計算花更多時間。
這也帶來新的困難，像是 TensorFlow graph 看得懂 matmul 和 tensor shape，GPU 低階 IR 看得懂 warp、shared memory 與硬體指令，但是單一 IR 很難同時清楚表達這些資訊。如果太早轉成低階指令，matmul、shape 和資料重用等高階結構便會消失；如果一直停在高階表示，又無法描述硬體如何執行。
MLIR 的重要性就在這裡，它讓 compiler 可以建立多個抽象層級，上層保留 tensor、shape 和 tile 等資訊，往下再逐步加入 loop、warp、register、shared memory 與特定指令。每一層都保留當下最佳化需要的資訊，並能繼續轉換到下一層，而且 MLIR 提供了一套組織這些多層 IR 的共用框架。
MLIR 源自 TensorFlow 軟體堆疊遇到的問題。當時一個 TensorFlow model 可能經過
TensorFlow Graph
├── Grappler
├── XLA HLO
├── TensorRT
├── nGraph
├── TensorFlow Lite
├── Core ML / NNAPI
├── TPU IR
└── LLVM IR 與其他 runtime
這些元件各自有 graph、IR、compiler pass 和 runtime。Compiler pass 是一次分析或轉換步驟，例如刪除無用程式碼或把高階運算改寫成低階運算。這些元件沒有共用一致的 infrastructure 和設計規則，會遇到一些問題像是
Chris後來發現，這個情況也存在於通用程式語言。LLVM 提供了非常成功的共同 backend，但 Swift、Rust、Julia、Fortran 等語言仍需要 SIL、MIR、Julia IR、FIR 這類高階或中階 IR，才能表達語言專屬語意與最佳化。
Swift AST ──> SIL ──┐
Rust AST ──> MIR ──┤
Julia AST ─> Julia IR ──> LLVM IR ──> machine code
Clang AST ──────────┘
這些專用 IR 都有自己的用途，每個團隊卻也要重做一套 compiler infrastructure。Chris 提議一套通用框架，讓各團隊快速定義新抽象，同時共享工程能力。
論文把 LLVM IR 簡單描述為 C with vectors，它擅長 scalar、vector、control flow 和低階 code generation，也讓不同語言能共用最佳化與多種 CPU backend。Backend 是 compiler 靠近目標硬體的後段，負責產生特定指令集的程式碼。
較高階的問題很難直接放進 LLVM IR
假設一個 matmul 已經被展開成 pointer arithmetic、load、multiply、add 和 branch，後面的 pass 要重新辨認這是一個矩陣乘法，成本高且容易受程式形狀影響。從低階表示中重建已經遺失的高階結構稱為 raising。
MLIR 允許 compiler 保留高階 operation。Operation 是 MLIR 的基本單位，可以表示一次運算，也可以表示 function 或 module 等結構。相關最佳化完成後，compiler 才逐步進行 lowering，也就是把高階 operation 改成更接近硬體的低階 operation。LLVM 仍然可以負責低階最佳化、register allocation 和 machine code generation。
不同語言、框架和加速器可以定義自己的 dialect，同時共用 IR 資料結構、改寫機制、錯誤訊息、輸入輸出與測試工具。
同一份 IR 可以同時包含多個抽象，某些 operation 保留 tensor 或 graph 語意，另一些 operation 已經 lower 到 loop、GPU 或特定 accelerator 指令，這讓 compiler 可以依硬體需求逐段處理。
開發新的 compiler 時，工程師可以把資源放在 domain operation、type、規則檢查和 lowering，省下重做通用 infrastructure 的成本。論文把新增 operation、type，並組成 dialect 視為常見的擴充方式。
MLIR 支援 dialect conversion，也就是把某種 dialect 的 operations 轉成另一種 dialect。逐步轉換後，IR 可以再交給 LLVM 或其他 backend。採用 MLIR 不要求一次換掉完整 stack，系統可以先在一段 graph optimizer 或一條 lowering path 中採用。
MLIR core 只固定少數通用概念，主要是 operation、type 和 attribute。Type 描述值的種類，例如 i32 是 32-bit integer，tensor<16xf32> 是包含 16 個 f32 的 tensor。Attribute 是附在 operation 上、編譯時期已知的靜態資訊，例如 convolution 的 stride 或 demo.scale 的 factor = 2.0。
來源程式的抽象語法樹（Abstract Syntax Tree，AST）則是語法剖析器（parser）依語法規則讀完原始程式後，建立的樹狀結構。例如 a + b 的 AST 會有一個加法節點，下面接著 a 和 b 兩個子節點。Tensor graph、AST、loop、control flow 與 LLVM instruction 等語意都可由 dialect 定義。
這個設計讓新領域能加入自己的 abstraction，代價是 dialect 可能各自發展出不相容的概念。論文承認技術機制無法單獨消除生態碎片化，因此 dialect 設計仍要考慮重用與互通。
MLIR 使用 Static Single Assignment（SSA）形式：每個值的名稱只會被定義一次。一般程式可能先寫 x = 1，後面再寫 x = x + 1；SSA 不會重新定義同一個 x，而是替新值取一個新名稱
%x0 = arith.constant 1 : i32
%one = arith.constant 1 : i32
%x1 = arith.addi %x0, %one : i32
%x0 和 %x1 是兩個不同的 SSA values。只要看名稱，compiler 就能追到值由哪個 operation 產生，又被哪些 operation 使用，這條關係稱為 use-def chain。很多最佳化都需要這種明確的資料流資訊。
MLIR 也把 region 設為一級結構。Region 是 operation 內可以容納子程式的區域，例如 loop body 或 function body。當 scf.for 表示一個 loop 時，loop 每次迭代要執行的 operations 就放在它的 region 裡。
Region 裡面再分成 blocks；block 是一段按順序執行的 operations，程式只會從 block 開頭進入。Block 的最後一個 operation 稱為 terminator，負責結束這個 block，並指定是否前往另一個 block。Operation 包含 region、region 包含 block、block 再包含 operation，因此能組成巢狀結構。
Operation
├── operands / results
├── attributes / types / location
└── Regions
└── Blocks
├── block arguments
└── Operations
└── Regions ...
LLVM IR 常用 Control Flow Graph（CFG，控制流程圖）表達程式如何執行。圖中的節點是 basic block，邊則表示 branch 可能跳到的下一個 block。一個 if / else 大致會形成這樣的 CFG
┌→ then block ─┐
entry block ─┤ ├→ merge block
└→ else block ─┘
entry block 檢查條件後選擇一條邊，then 與 else 執行完再合流到 merge block。MLIR 在 lowering 前可以繼續使用 structured loop 或 condition，之後再視 backend 需求轉成 CFG。
Compiler 把高階 representation 分成多個小步驟往下轉換，每一步只消去當下不再需要的抽象
tensor / graph op
↓
structured linear algebra
↓
loop / affine / vector
↓
GPU mapping 或 accelerator dialect
↓
LLVM dialect
↓
LLVM IR / target ISA
Pipeline 是依序執行的編譯步驟串。MLIR 沒有規定只能使用某一條固定 pipeline，compiler 可以加入 domain dialect、target dialect 與自己的 conversion。
高階結構應保留到相關分析與最佳化完成。Structured loop 還在時，可以直接做 tiling、fusion 或 interchange，tensor op 還在時，可以使用 shape 和 layout 資訊，降成一般 CFG 和 pointer operation 後，這些結構就難以可靠重建。
MLIR 也允許同一份 IR 混合不同抽象。例如 function 的一部分已經變成 accelerator-specific instruction，另一部分仍維持高階 operation。Progressive lowering 可以只轉換部分 operations，不必要求整個 module 同時下降一層。
可擴充 IR 需要明確檢查 invariant，也就是 IR 必須始終滿足的規則。Dialect 應該定義 operation 的 operand（輸入值）、result（輸出值）、type constraint、attribute 和結構規則，再由 verifier 檢查。Verifier 是 IR 的規則檢查器，很像每個 operation 自帶的輸入格式檢查。
例如 matmul 的兩個輸入都應該是 tensor，且 A 的 K 長度必須和 B 的 K 長度一致。如果 A 是 tensor<4x8xf32>，B 卻是 tensor<7x16xf32>，verifier 就能在 matmul 還清楚可辨識時報錯。等它降成一大串 load、multiply 和 add 後，同一個錯誤就較難指回矩陣尺寸。
許多 lowering 和最佳化可以寫成 rewrite pattern，當 IR 出現某種 operation 組合時，就把它換成另一組 operations。例如 x + 0 的結果一定是 x，可以寫成這條概念規則
match : arith.addi(%x, 0)
rewrite : %x
MLIR 的 rewrite framework 會把 pattern 和目前的 IR 比對。找到符合的 operations 後，framework 負責替換它們，並維護 SSA use-def chain 等關係。各 dialect 只需要提供自己的規則。
論文也介紹 Declarative Rewrite Rule（DRR），讓開發者寫下要找哪種 IR 形狀和要改成什麼，不用自己實作完整搜尋流程。Declarative 的重點就是寫下想要的匹配與結果，搜尋順序和改寫細節由 framework 處理。
DRR 使用 TableGen 描述，TableGen 是 LLVM 工具鏈中的宣告式語言。開發者用它列出 operation 的欄位、型別限制或 rewrite 前後的 operations，TableGen 再產生重複的 C++ 程式碼。它可以減少手寫配對邏輯與額外工程程式碼。
簡單 rewrite 可以宣告式描述，需要複雜分析時仍可寫 C++。這讓 common case 保持精簡，又能處理 target-specific lowering。
Operation 可以攜帶並傳遞 source location，也就是它對應的原始程式位置。當一個高階 op 經過多次 rewrite 和 lowering 後，compiler 仍有機會追查它從哪段 source、哪個 AST node 或哪條 transformation 產生。
論文把這項能力連到 diagnostic（錯誤與警告訊息）、debugging、compiler testing，以及安全關鍵程式需要的 transformation trace。多層 pipeline 如果沒有 provenance（來源紀錄），最終低階 IR 很難追回原始問題。
論文第三節把 operation 定義為 MLIR 的主要語意單位。Instruction、function、module 都可以建模成 operation。每個 operation 可以包含
arith.addf、affine.for。一個簡化的 MLIR operation 可以長成
%result = "demo.scale"(%input) {
factor = 2.0 : f32
} : (tensor<16xf32>) -> tensor<16xf32>
demo.scale 是 operation name；%input 是 operand；%result 是 SSA result；factor 是 attribute；前後的 tensor<16xf32> 是 operand 與 result type。
Dialect 把一組 operation、type 和 attribute 放進獨立 namespace。Namespace 是名稱字首，用來分組名稱並避免衝突，例如 arith.addf 的 arith 就是 namespace。
func : function
arith : arithmetic
scf : structured control flow
affine : affine loop / affine map
memref : memory reference
linalg : structured linear algebra
gpu : GPU execution mapping
llvm : LLVM-level representation
Dialect 本身不自動帶來新的執行語意，語意由其中的 operation 和 type 定義。為了讓通用 pass 理解這些自訂 operation，MLIR 使用 trait 與 interface。Trait 是可重用的性質標記，例如所有輸入與輸出使用相同型別，interface 則定義 pass 可以向不同 operation 詢問的共同問題，例如 operation 有沒有讀寫記憶體。
這項設計讓多個 dialect 出現在同一份 IR。func、arith、gpu 和自訂 accelerator dialect 可以逐步接在一起，不必先合併成一套全域 opcode 清單。
如果 MLIR 只允許新增 operation，各個團隊仍然要重做大量工具，論文第四節因此把 infrastructure 也列為貢獻。
Operation Definition Specification（ODS）是 MLIR 定義 operation 規格的方式。開發者用 TableGen 寫下 operation 的名稱、輸入、輸出、屬性、型別限制與 trait，工具再產生部分重複的 C++ 框架程式碼、規則檢查與文件。
用前面的 demo.scale 為例，ODS 可以規定它讀入一個 tensor、產生同型別 tensor，並要求 factor attribute 存在。從這份規格，MLIR 可以產生建立 operation 的 helper、取得 factor 的方法與基本 verifier。新增 operation 時，工程師就不用手寫每一段重複 C++。
ODS 讓 operation 規格成為工具可以讀取的資料，也減少宣告、實作與文件各寫一份後逐漸不一致的風險。
Canonicalization 是把等價的 IR 改成較統一、較容易處理的寫法。前面的 x + 0 → x 就是 canonicalization 範例：後面的 pass 只需要處理簡化後的 %x。Constant folding 則會把 2 + 3 這種只包含常數的運算提前算成 5。
Dialect conversion 把一種 dialect 的 operations 轉成另一種。這些工作和 lowering 都可以建立在 operation rewrite 上；Dialect 提供 pattern，driver 負責搜尋與套用。
開放的 operation 集合帶來一個問題：通用 pass 不可能事先認識未來出現的每個 dialect。MLIR 讓 dialect 藉由 trait、hook 和 interface 提供 pass 需要的性質。
例如刪除無用程式碼的 pass 需要知道 operation 有沒有 side effect。一個結果沒有被使用的加法可以刪除，寫檔案的 operation 即使沒有輸出被使用也不能刪除。Pass 可以藉由 interface 詢問 operation 的記憶體效果，不用事先認識每個 dialect 的所有 operations。
Inliner 則需要詢問某個 region 能不能被展開到呼叫處。缺少所需的 trait 或 interface 時，通用 pass 要跳過該 operation 或回報失敗，不能假設改寫一定安全。
這套機制讓同一個 canonicalization 或 inliner pass 能作用在多種 dialect，target-specific knowledge 仍由 dialect 自己維護。論文也保留 dialect-specific pass，讓需要完整 domain 語意的 transformation 直接針對特定 operation 工作。
Pass manager 負責組織一連串 compiler passes，包含執行順序與每個 pass 處理的 IR 範圍。MLIR 裡的 module 和 function 也是 operation，因此 pass 可以掛在適合的 operation type 與巢狀深度。
具有 isolated-from-above 性質的 operation 會形成獨立的 SSA 範圍，內部 operation 不能直接使用外層定義的 SSA value。邊界清楚後，pass manager 就能平行處理彼此獨立的 regions。
Round-trippable textual IR 表示「記憶體中的 IR → 文字 → 記憶體中的 IR」往返轉換後，重要資訊不會遺失。工程師因此可以把某個 pass 之前的 IR 存成檔案，之後再讀回來重現同一個 compiler bug。
Dialect 可以提供容易閱讀的 custom syntax，也保留任何 operation 都能使用的 generic syntax。Custom syntax 比較精簡，generic syntax 則會明確列出 operation 名稱、operands、attributes 與 types，即使 dialect 沒有專用的文字輸出邏輯也能保存和讀取。
每個 pass 都能用文字 IR 單獨測試
input.mlir
↓ mlir-opt --some-pass
output.mlir
↓ FileCheck
檢查 operation 是否按照預期改寫
論文的設計目標是讓 pass 不仰賴不可見的 pipeline state。只要輸入 IR、pass options、dialect / context 和外部環境相同，就能把某個階段的文字 IR 抽出來，單獨重現和測試該 pass。完整 pipeline 前面的 pass 如果已改變輸入 IR，兩次執行就不屬於相同條件，這項設計對除錯、測試和理解 lowering 過程很實用。
Lowering 是把某種 abstraction 換成較接近目標執行模型的 representation。以矩陣乘法為例，可能走過
linalg.matmul
↓ tiling / fusion
scf.for + vector operations
↓ GPU mapping
gpu operations
↓ dialect conversion
nvvm / llvm operations
↓ translation
LLVM IR / PTX
Pass 可以分成四種角色
順序會影響結果。如果 tensor operation 太早變成 pointer arithmetic，fusion pass 失去可辨識的結構。如果 target-specific operation 一直沒有 lower，backend 就無法產生指令。MLIR 提供組合 pipeline 的機制，但不會自動解出最佳 pass ordering。
這篇 paper 的 evaluation 用 use case 和工程經驗為主，沒有把 MLIR 當成單一 optimizer 做一組固定效能 benchmark。Chris用不同領域說明同一套 IR infrastructure 能否容納多種 abstraction。
TensorFlow graph 是帶有 dynamic execution semantics 的高階 dataflow graph。MLIR 可以表示 TensorFlow operation、tensor、resource 和 control dependency，再做 algebraic optimization、裝置分派、mobile lowering 或銜接 XLA。
Affine dialect 用 structured loop、affine map 和 multidimensional memory access 表示 polyhedral transformation 所需資訊。Loop 保留在 IR 中，compiler 可以做 transformation，又不必先把程式轉成差距很大的 polyhedral representation，再困難地 raise 回 loop。
Flang 使用 FIR 表示 Fortran 專屬的 array、dispatch 和語言語意。FIR 可以重用 MLIR 的基礎設施，也能和 OpenMP、GPU 等較通用 dialect 組合。
從 paper 的論證可以整理出 MLIR 已經提供的能力
論文也直接承認幾項風險與研究問題
MLIR 提供建造 compiler 的共用結構，compiler 的 domain modeling、cost model、最佳化策略與硬體知識仍要由各專案完成。
有了 paper 的背景後，Triton 的 pipeline 就能放進 MLIR
Triton Python
↓
TTIR：保留 tile-level tensor program
↓ progressive lowering
TTGIR：加入 GPU layout、warp / CTA mapping
↓
LLVM IR
↓
PTX / cubin
TTIR 和 TTGIR 是 Triton 為自己的 domain 與 GPU code generation 定義的 abstraction。它們能使用 MLIR 的 operation、dialect、SSA、pattern rewrite 和 pass infrastructure，同時保留 Triton 需要的 tile 與 layout 語意。
明天我們會拿 @triton.jit matmul，看它如何產生 TTIR、TTGIR、LLVM IR、PTX 和 cubin。
iThome鐵人賽
看影片追技術