Jacques Pienaar 如何用 MLIR 重整 AI 編譯器？從 TensorFlow 多層 IR 到 OpenXLA

Jacques Pienaar 如何用 MLIR 重整 AI 編譯器？從 TensorFlow 多層 IR 到 OpenXLA

AI 編譯器最難處理的，不是把單一矩陣乘法翻成機器指令，而是如何在模型框架、圖最佳化、張量運算、記憶體配置與各種加速器之間保存足夠語意。Jacques Pienaar 參與 MLIR 的創建與長期維護，並把這套多層中介表示方法帶進 TensorFlow 生態。他的價值不只是完成某個編譯 pass，而是協助產業從大量彼此斷裂的轉換器，走向能逐層驗證、逐層降低抽象的共享基礎設施。

讀 Jacques Pienaar 與 MLIR 時，可以先把問題拆成三個邊界：模型框架保留什麼語意，逐層 lowering 在哪裡做決策，以及編譯器如何證明轉換沒有改變結果。這樣看，MLIR 不只是另一種 IR 格式，而是讓不同抽象層能共同驗證、共同演進的工程方法。

本文把 Google Research 人物頁、MLIR／OpenXLA 官方文件與編輯分析分開。官方資料能核對人物角色、維護者與專案設計；至於效能、正確性和供應鏈風險，仍須回到實際模型、硬體、版本與測試條件，不把工具存在直接推成所有產品都已採用。

Jacques Pienaar 的位置：連接 TensorFlow 與多層編譯器

Google Research 官方人物頁將 Jacques Pienaar 的研究焦點放在機器學習系統與編譯器；MLIR 官方論文資訊也把他列為共同作者。這表示他的工作不能只用「TensorFlow 工程師」概括，而應放在模型框架如何銜接可重用編譯技術的脈絡中理解。

模型作者習慣思考 layer、tensor 與 shape，硬體端則面對向量寬度、記憶體階層、同步與指令。兩端若直接配對，每增加一個前端或後端就要增加一批專屬轉換。Pienaar 推動的路線，是保留中間各層的語意，讓轉換能重用，也讓錯誤可以定位。

N×M 轉換問題為何會拖垮 AI 軟體堆疊

假設有多個模型框架與多種 CPU、GPU、TPU 或 NPU，最直覺的作法是為每組組合寫 converter 與 code generator。短期看來快速，長期卻會讓相同的 constant folding、shape inference、fusion 與 layout 處理被重複實作，修正也無法同步。

更嚴重的是語意在轉換途中逐步流失。某個前端知道 operation 可以重排，後端卻只看到不透明函式；某個 accelerator 需要特殊 layout，高階圖又沒有表達位置。多層 IR 的目的，是把各種決策放在仍擁有必要資訊的階段，而不是用一個格式硬撐到底。

MLIR 不是單一格式，而是一套建立 IR 的基礎設施

MLIR 的名稱容易讓人誤以為它是一個固定中介語言，實際核心是 operation、type、attribute、region、block 與 dialect 等共同機制。不同領域能定義自己的語意，同時共享 parser、printer、verifier、pass manager、rewrite 與診斷工具。

這種結構讓模型層可以保留 tensor 與 domain operation，較低層再表達迴圈、向量、GPU 或 LLVM。每一層不必重新建一套 compiler framework，也不必過早犧牲資訊。對平台團隊而言，可重用基礎設施往往比單一最佳化更能降低多年維護成本。

Dialect 讓領域語意存在，但也需要清楚邊界

Dialect 可以表示 TensorFlow operation、線性代數、tensor、memory reference 或特定硬體功能。它的好處是 operation 能帶著明確 type、attribute 與 verifier，而不是躲在無法檢查的字串或 custom call 中；工具也能讀懂並轉換這些結構。

然而任何團隊都能新增 dialect，也可能造成碎片化。若兩個 dialect 對 broadcasting、rounding 或 dynamic dimension 的定義不同，只靠名稱相近不能互通。成熟平台必須指定 owner、語意文件、版本政策、canonical form 與明確 conversion path。

Progressive lowering 保留資訊到最適合決策的時刻

逐層 lowering 的關鍵，是不要在流程太早把張量計算壓平成低階指令。高階階段適合做模型語意相關的融合與形狀推導；中階可以處理 tiling、loop 與 vectorization；接近硬體時才決定記憶體空間、thread mapping 與實際指令。

每次轉換都應有合法輸入、合法輸出與失敗條件。若某個 operation 無法安全降低，compiler 應提供可讀診斷，而不是默默留下未支援節點或產生錯誤程式。這也讓平台能加入 verifier，在每個重要邊界阻止語意漂移。

TensorFlow 導入 MLIR 的真正意義

TensorFlow 生態曾同時存在圖最佳化、XLA、saved model、裝置特定路徑與各種轉換工具。MLIR 提供的不是一次性重寫，而是讓 TensorFlow operation 能進入共同架構，與 shape、tensor、linalg、GPU、LLVM 等較低層表示建立可治理的橋梁。

對使用者而言，成功不應以「內部用了 MLIR」衡量，而要看更多模型能否被正確編譯、錯誤是否更容易理解、不同硬體支援能否共享、升級是否減少回歸。內部架構只有轉化成可驗證的 coverage 與可靠性，才是產品價值。

Shape 是 AI compiler 最容易被低估的語意

靜態 shape 可讓 compiler 預先決定 buffer、tiling 與 kernel；動態 shape 則需要 guard、特化、runtime check 或通用 fallback。若在 IR 中只保留未知維度，很多最佳化無法進行；若把偶然觀測到的尺寸當成永遠成立，又會造成錯算或重編譯風暴。

可靠系統應記錄 shape constraint 的來源，區分模型保證、輸入協定與 profiling 猜測。特化版本要有命中範圍、容量上限與退回路徑；測試也要覆蓋邊界尺寸、空 tensor、大輸入與不規則 batch，而非只跑最常見樣本。

Operation definition 必須成為可執行契約

MLIR 透過宣告式機制描述 operation 的 operands、results、attributes 與約束，並能產生部分 boilerplate。這不只是節省 C++ 程式碼，而是讓 verifier、文件與工具使用同一份定義，減少實作和說明彼此偏離。

企業 dialect 應把必要 invariant 寫進 verifier，例如 rank、element type、layout 與 attribute 組合。不能驗證的假設就容易散落在 pass 裡，直到某個模型觸發 crash。錯誤訊息要指出具體 operation 與違反條件，才能讓模型與 compiler 團隊協作。

Pattern rewrite 需要終止性與語意保證

圖與 IR 最佳化大量依靠 rewrite pattern，例如把等價 operation 合併、刪除冗餘轉換或改成更適合後端的形式。局部規則若互相反轉，可能造成無限循環；若忽略浮點、overflow、NaN 或 side effect，也可能產生靜默錯算。

每個 pattern 應寫清楚前置條件，對語意保持轉換做 differential test，並限制可能擴張 IR 的規則。Canonicalization 是讓表示穩定，不是無條件追求最短。平台還應記錄 rewrite 統計，找出編譯時間暴增或規則競爭。

Conversion legality 讓 lowering 不再只是「盡量轉」

從一個 dialect 轉到另一個 dialect 時，真正重要的是定義哪些 operation 已合法、哪些必須消失、哪些可暫時保留。若 pipeline 結束後仍混有未支援 operation，部署端才發現問題，錯誤位置已離原始模型很遠。

MLIR conversion 架構讓 target legality 與 type conversion 成為流程的一部分。企業可在每個交接點建立硬 gate：輸出只能含允許 dialect、型別必須符合 target、未轉換節點立即失敗。這種嚴格性會增加早期工作，卻能降低 production 的不確定性。

StableHLO 與 OpenXLA 延伸了跨框架契約

OpenXLA 的 StableHLO 旨在提供可移植、具版本治理的高階 operation 集合，讓 framework 與 compiler 後端透過較穩定的介面交換模型。它和任意內部 IR 不同，因為序列化後的模型可能要跨時間、跨產品與跨組織使用。

穩定格式必須處理 operation 語意、版本、相容範圍與升級工具。使用者不能只確認模型今天能匯出，還要測舊 artifact 在新 runtime 是否被正確接受或明確拒絕。長期可攜性建立在規格與測試，不是副檔名。

硬體特化應靠逐層決策，而不是污染模型前端

TPU、GPU 與 CPU 的記憶體階層、向量能力與同步方式差異巨大。若把 target 細節直接塞進高階模型，移植成本會快速上升；若完全忽略 target，則可能無法達到可用效能。多層 IR 讓硬體資訊在適當階段進入，而不必穿透所有上層。

平台應為每個 target 定義 feature profile 與 fallback，而非假設同一 pipeline 適用所有裝置。對 tiling、fusion 與 layout 選擇保存 cost model、硬體版本與量測證據；驅動程式或晶片 revision 改變後，重新驗證而不是沿用舊結論。

Pass pipeline 是需要版本控制的產品

相同一組 pass 只要順序不同，就可能改變正確性、編譯時間與輸出效能。生產環境不能只保存 compiler 版本，還要保存完整 pipeline、選項、target feature、輸入 artifact digest 與產出 digest，才能重現一次部署。

新 pipeline 先用代表模型做 shadow compile，比較輸出、數值與效能，再逐步放量。若出現 regression，系統應能回到上一個已驗證 artifact，而不是在服務啟動時臨時重新編譯。編譯流程的可回退性和應用程式發布同樣重要。

診斷與 IR 可視性決定平台能否被一般團隊使用

Compiler 專家能閱讀數千行 IR，但模型開發者需要知道是哪個模型節點、輸入 shape 或不支援 operation 導致失敗。好的工具應保留 source location，能輸出前後 IR、最小化問題案例，並把內部 dialect 名稱翻成可採取行動的說明。

同時不能把完整模型、權重路徑或使用者資料任意寫進 log。除錯 artifact 要分級、去識別並限制存取。可觀測性不是把所有內容留下，而是在安全邊界內保存足夠 provenance，使錯誤可重現又不造成資料外洩。

數值正確性比 compiler 成功返回更重要

AI compiler 常做 reassociation、fusion、低精度轉換與近似數學，輸出不一定逐位元相同。測試要依模型任務設定容差、分布與品質指標，並保留高精度 reference。只比較幾個輸出樣本，很難發現罕見輸入下的偏差。

每次新增後端或最佳化，先覆蓋 operation-level 測試，再跑模型級 differential test，最後以實際資料做品質 gate。對安全或高風險用途，還要檢查 subgroup 表現與極端值。速度提升不能掩蓋無法量化的正確性損失。

Artifact cache 需要身分、相容與完整性鎖

編譯大型模型昂貴，因此平台會重用 cache；但只用模型名稱當 key，可能把不同權重、shape、compiler 或硬體的產物混在一起。完整 key 至少應包含模型內容、pipeline、target、工具鏈與重要選項的 digest。

部署端要驗證簽章、schema 與 target 相容性，損壞或來源不明的 cache 不可直接執行。Cache miss 可以造成效能下降，錯誤命中則可能造成錯算或安全問題。平台還應有容量、淘汰與回溯政策，避免舊 artifact 無限堆積。

開放式 dialect 生態也帶來供應鏈風險

Custom dialect、pass plugin 與外部模型 parser 可能處理不可信輸入，甚至載入 native code。Build service 應隔離檔案、網路、CPU 與記憶體，對 plugin 建立來源 allowlist、code review、版本鎖與簽章，不讓 production worker 任意下載依賴。

Parser、verifier 與 rewrite 要做 fuzz、資源上限與惡意尺寸測試。對未知 operation 明確拒絕，不以「忽略後繼續」掩蓋問題。編譯器位於模型供應鏈中心，它的擴充性必須和最小權限一起設計。

MLIR 上游治理是技術能力的一部分

MLIR 官方維護者文件把 Jacques Pienaar 列在核心維護群中。這類角色不只提交功能，也要評估抽象是否通用、介面是否能長期維護、測試是否足夠，以及下游需求是否適合放進核心。共享 infrastructure 的品質來自這些日常治理。

企業若大量依賴 MLIR，應投入 upstream，而不是永久維護龐大私有 fork。通用修正盡量送回社群，內部擴充則縮小邊界並持續 rebase。Fork 初期看似取得控制，數年後常變成無法升級的 merge debt 與人才風險。

導入 MLIR 的合理路徑不是一次重寫全部 compiler

團隊可先選一個邊界清楚、現有痛點明確的子流程，例如匯入某類 operation、建立 shape dialect 或替換單一 lowering stage。先定義輸入輸出契約、golden test、效能基線與 fallback，再比較新舊路徑。

第二階段才擴大 operation coverage，建立 dialect owner、版本與診斷準則。若一開始就同時更換模型格式、最佳化、runtime 與硬體後端，回歸很難定位。漸進採用不只降低交付風險，也能驗證抽象是否真的適合。

企業驗收矩陣應同時看正確性、效能與營運

代表模型要覆蓋 static 與 dynamic shape、不同 dtype、control flow、custom op 與不規則輸入。量測 compile time、steady-state latency、throughput、peak memory、artifact size 與能耗，並保存硬體、driver、資料與命令。

營運測試則包括不支援 operation、compiler timeout、cache 損壞、版本不相容、worker crash 與回退。只有在失敗可診斷、可隔離且可恢復時，compiler 才能成為 production 基礎設施。單一 benchmark 領先不是完整驗收。

Jacques Pienaar 留給 AI 基礎設施的核心方法

Jacques Pienaar 的代表性，在於協助把模型框架與 compiler research 之間的接縫變成可重用系統。MLIR 讓多種抽象並存，以 dialect 保存領域語意，以 progressive lowering 控制資訊何時消失，再以 verifier、rewrite 與 conversion framework 建立可檢查的流程。

這套方法不保證每個硬體自動變快，卻提供一種可持續的工程紀律：先定義語意，再做轉換；保存 provenance，再談效能；為失敗建立診斷與 fallback；把共享問題送回上游。當模型與硬體持續變動時，這些能力比一次性的最佳化更有長期價值。

同一套 MLIR／AI compiler 生態的延伸閱讀：Jacques Pienaar 文章先從 TensorFlow 多層 IR 與 OpenXLA 看編譯器轉換，再接著對照 Mehdi Amini 如何打造可擴充 MLIR？從 LLVM 基礎設施到 AI 編譯器治理 如何從 LLVM 基礎設施與治理角度建立可擴充邊界；兩篇能把 IR 設計與企業導入策略接在同一條技術實體路徑上。

延伸閱讀

若想把MLIR與OpenXLA放進AI硬體平台脈絡，可以接著閱讀黃仁勳如何把GPU變成AI權力中心，對照編譯器契約、GPU與模型部署。

若要把 MLIR 的多層 IR 放回更大的 AI compiler 脈絡，也可閱讀Chris Lattner、LLVM、MLIR 與 Modular與陳天奇從 XGBoost、TVM 到 MLC LLM 的路線；兩篇是技術對照，不把不同人物的職涯與專案成果互相代用。

官方資料與延伸閱讀

• Google Research：Jacques Pienaar 官方人物頁
• MLIR 官方 FAQ 與論文作者資訊
• MLIR 官方核心維護者文件
• MLIR 官方開放專案與導師名單
• Google Research：MLIR in TensorFlow Ecosystem
• OpenXLA：StableHLO 官方文件
• Jacques Pienaar GitHub 官方身分頁

增量：多層 IR 把 AI 編譯器的變化隔離在可重用介面內

MLIR 的核心價值可以再用「保留資訊—逐步降低—目標生成」來理解。高層 IR 保留模型與運算語意，中層 dialect 讓不同硬體與框架有可組合的表示，低層 IR 再逐步接近特定 runtime 或指令集。這使編譯器不必為每個模型與每種硬體建立一條完全獨立的管線。

• 保留：越早降低抽象，越可能失去融合、形狀與佈局資訊。
• 組合：dialect、pass與轉換規則能否被不同專案重用。
• 驗證：每次 lowering 後，語意與數值是否仍符合原模型。
• 部署：OpenXLA、StableHLO與硬體後端的介面是否清楚。

因此，MLIR 不是一個「自動變快」的黑箱，而是一套讓編譯器工程團隊能管理複雜度、交換中間表示並測試每次轉換的基礎設施。