2：MLIR的引入)
系列文章目錄章節(jié)1什么是MLIR章節(jié)2MLIR的引入章節(jié)3MLIR完成一個端到端的流程目錄系列文章目錄一、tensorflow執(zhí)行方式1. 前端構建計算圖與 Op 注冊2. 后端準備Kernel 注冊與設備匹配3. 核心調度Executor 執(zhí)行器4. 底層執(zhí)行Kernel Compute二、使用XLA加速tensorflow1.核心原理從“解釋執(zhí)行”到“編譯執(zhí)行”2.XLA 加速帶來的主要優(yōu)勢3.XLA和LLVM的關系三、XLA存在的問題四、MLIR 針對性解決的核心痛點提示以下是本篇文章正文內容下面案例可供參考一、tensorflow執(zhí)行方式執(zhí)行方式流程如下其中Graph? 定義了“做什么”?Executor? 決定了“何時做”和“在哪做”而 ?Kernel? 負責“具體怎么做”。三者協(xié)同工作實現(xiàn)了 TensorFlow 的張量計算能力。TensorFlow 從計算圖Graph的操作節(jié)點到執(zhí)行器Executor再到內核Kernel的執(zhí)行流程是一個典型的?“定義-優(yōu)化-調度-執(zhí)行”?的流水線過程。以下是這一過程的詳細拆解1. 前端構建計算圖與 Op 注冊當用戶在 Python 中調用如tf.matmul或tf.add時并沒有立即進行數(shù)值計算而是完成了以下工作?Op 定義與注冊?TensorFlow 內部維護了一個全局的OpDefLibrary。每個操作如 MatMul都有一個對應的 ?OpDefinition?定義了它的輸入輸出類型、屬性如 transpose_a等。這些定義通過宏如REGISTER_OP在 C 層注冊。?節(jié)點插入 Graph?Python API如gen_math_ops.py會將用戶傳入的參數(shù)轉換為底層 C 調用最終在當前的默認Graph對象中創(chuàng)建一個 ?NodeDef?。?NodeDef? 是圖的序列化表示僅包含操作的名稱、類型、輸入依賴和屬性?不包含具體的計算邏輯?。此時圖中的邊Edge建立了節(jié)點之間的數(shù)據(jù)依賴關系。2. 后端準備Kernel 注冊與設備匹配計算邏輯是由 ?Kernel? 實現(xiàn)的而不是 Op 本身。一個 Op 可以有多個 Kernel 實現(xiàn)分別對應不同的硬件設備CPU/GPU/TPU和數(shù)據(jù)類型float32/int32。?Kernel 注冊?開發(fā)者使用REGISTER_KERNEL_BUILDER宏將具體的 C 計算類繼承自OpKernel注冊到系統(tǒng)中。例如MatMul在 CPU 上有一個 kernel在 GPU 上有另一個基于 cuBLAS 的 kernel。?設備放置Device Placement?在會話Session初始化或圖執(zhí)行前TensorFlow 的運行時會分析圖中的每個 Node根據(jù)設備的可用性、內存限制和用戶指定的with tf.device()上下文決定每個節(jié)點由哪個設備Device執(zhí)行。?實例化 Kernel?一旦確定了設備運行時就會為該節(jié)點查找并實例化對應的OpKernel對象。這個對象包含了執(zhí)行計算所需的所有狀態(tài)和方法主要是Compute()函數(shù)。3. 核心調度Executor 執(zhí)行器?Executor? 是 TensorFlow 運行時的核心組件負責驅動整個計算圖的執(zhí)行。它的主要職責是將靜態(tài)的圖結構轉化為動態(tài)的執(zhí)行流。?圖優(yōu)化與分區(qū)?Executor 接收優(yōu)化后的 Graph。如果涉及多設備如 CPU 和 GPU它會將圖分割成多個子圖Partition每個子圖對應一個特定的設備。?構建執(zhí)行計劃?Executor 會對節(jié)點進行?拓撲排序?確定執(zhí)行的先后順序。它會識別哪些節(jié)點可以并行執(zhí)行即沒有數(shù)據(jù)依賴關系的節(jié)點。?異步調度模型?Executor 維護一個任務隊列。它不會等待一個節(jié)點完全結束才啟動下一個而是采用異步機制當一個節(jié)點的所有輸入張量Tensor就緒時該節(jié)點變?yōu)椤翱蓤?zhí)行”狀態(tài)。Executor 將該節(jié)點的計算任務封裝為一個回調函數(shù)提交給目標設備的線程池Thread Pool。對于 GPU這通常意味著將 CUDA Kernel 啟動指令放入 GPU 流Stream中。4. 底層執(zhí)行Kernel Compute當調度器將任務分發(fā)給具體設備后真正的計算發(fā)生在 ?Kernel? 層面?調用 Compute()?設備線程池中的工作線程會調用對應OpKernel的Compute(OpKernelContext* context)方法。?獲取輸入與分配輸出?OpKernelContext提供了訪問輸入張量數(shù)據(jù)的接口并負責為輸出張量分配內存通常在設備的顯存或內存中。?執(zhí)行具體算法??CPU Kernel?直接執(zhí)行 C 代碼可能調用 Eigen 庫或 MKL 進行矩陣運算。?GPU Kernel?通常是一個輕量級的包裝器負責準備參數(shù)并啟動實際的 CUDA/HIP 核函數(shù)如cublasSgemm。?結果回傳與同步?計算完成后Kernel 將結果寫入輸出張量。Executor 會監(jiān)控這些完成信號一旦某個節(jié)點的輸出就緒它會通知所有依賴該輸出的后續(xù)節(jié)點觸發(fā)它們的執(zhí)行。二、使用XLA加速tensorflow使用 XLA 加速?是指利用 ?XLAAccelerated Linear Algebra加速線性代數(shù)? 編譯器對深度學習模型的數(shù)學運算進行即時編譯和優(yōu)化從而顯著提升模型在特定硬件如 TPU、GPU上的訓練和推理速度。簡單來說它就像是一個“超級翻譯官”兼“優(yōu)化大師”將你寫的 Python/PyTorch/JAX 代碼轉換成針對硬件高度優(yōu)化的機器碼而不是讓硬件去逐條解釋執(zhí)行原始的 Python 指令。執(zhí)行方式如下1.核心原理從“解釋執(zhí)行”到“編譯執(zhí)行”傳統(tǒng)框架如原生 PyTorch 的 eager 模式或早期 TensorFlow通常采用?命令式執(zhí)行?每行代碼發(fā)出一個操作指令硬件立即執(zhí)行。這種方式靈活但效率低因為每個小操作都有啟動開銷且無法看到全局優(yōu)化空間。XLA 加速的核心機制是“惰性求值”與“圖編譯”??構建計算圖Tracing/Lazy Evaluation?當你運行模型時XLA 不會立即執(zhí)行每個算子而是先“記錄”下所有的操作序列形成一個中間表示IR的計算圖。這個過程被稱為“追蹤”Tracing。?整體編譯與優(yōu)化?XLA 編譯器拿到整個計算圖后會進行全局優(yōu)化。最關鍵的技術是?算子融合Operator Fusion?它將多個小的、連續(xù)的操作如Add-Relu-BatchNorm合并成一個單一的 GPU/TPU 內核Kernel。?生成高效機器碼?優(yōu)化后的圖被編譯成針對特定硬件如 NVIDIA GPU 的 CUDA 代碼或 Google TPU 的指令集的高效二進制代碼。?執(zhí)行?最后硬件直接執(zhí)行這段優(yōu)化后的代碼。由于減少了內核啟動次數(shù)和數(shù)據(jù)在顯存中的讀寫次數(shù)性能大幅提升。2.XLA 加速帶來的主要優(yōu)勢?算子融合Operator Fusion?這是 XLA 最核心的優(yōu)化手段。通過減少 GPU/TPU 內核啟動開銷和中間張量的內存讀寫顯著降低延遲并提高吞吐量。?內存優(yōu)化?XLA 能更智能地管理顯存復用緩沖區(qū)從而允許你訓練更大的模型或使用更大的 Batch Size。?硬件特異性優(yōu)化?XLA 針對不同硬件后端CPU、GPU、TPU生成專屬的最優(yōu)代碼。例如在 TPU 上XLA 能充分利用其矩陣乘法單元MXU的高并行能力在 NVIDIA GPU 上它能更好地利用 Tensor Cores。3.XLA和LLVM的關系XLA 和 LLVM 是典型的「高層領域專用編譯器」與「通用底層編譯器基礎設施」的深度協(xié)作關系?二者定位完全互補XLA 不重復造輪子而是基于 LLVM 成熟的底層能力專注于AI張量計算的專屬優(yōu)化。XLA 的完整編譯流水線天然以 LLVM 作為核心后端基座?XLA 前端?接收來自 TensorFlow/JAX 的高層張量計算圖生成 XLA 專屬的中間表示 ?HLOHigh Level Optimizer?完成算子融合、自動并行、內存復用等AI領域專屬的高層優(yōu)化。?HLO 轉譯層?將高層的 HLO 張量操作逐步展開、映射為符合 LLVM 規(guī)范的中間表示LLVM IR比如一個HLO矩陣乘法會被拆解為一系列LLVM的加載、乘加、存儲指令序列 。?LLVM 后端接管?生成的LLVM IR直接送入LLVM完整的工具鏈經(jīng)過指令調度、寄存器分配、目標平臺專屬優(yōu)化后最終輸出x86、ARM、CUDA等不同硬件架構的原生機器碼 。三、XLA存在的問題XLA 原生架構的核心不足?IR層級斷層優(yōu)化空間受限?XLA 只有HLO這一個核心中間表示高層張量圖優(yōu)化完成后直接硬著陸到LLVM IR或TPU硬件指令中間缺失了多層級過渡IR。這導致循環(huán)變換、多維數(shù)組分塊、仿射優(yōu)化這類關鍵的中間層變換無法優(yōu)雅實現(xiàn)只能靠手寫大量特殊Case的Emitter代碼優(yōu)化能力被嚴重鎖死 。?硬件適配成本極高擴展性差?XLA 為每一種新硬件后端如某款新ASIC、邊緣端NPU都需要從零開始手寫整套代碼生成邏輯復用率極低。傳統(tǒng)XLA的GPU代碼生成甚至重度依賴LLVM的元數(shù)據(jù)hack來傳遞優(yōu)化信息維護成本極高新硬件接入周期長達數(shù)月甚至數(shù)年 。?代碼庫碎片化維護難度指數(shù)級上升?XLA的所有變換邏輯都是用C硬編碼實現(xiàn)的沒有統(tǒng)一的IR基礎設施。隨著算子數(shù)量、硬件種類增加手寫的Emitter代碼量爆炸式增長公共子表達式消除、常量折疊這類基礎編譯器優(yōu)化都要重復造輪子無法復用成熟的通用編譯器生態(tài)能力 。?跨框架兼容性差生態(tài)封閉?早期XLA深度綁定TensorFlow沒有標準化的中間交換格式PyTorch、PaddlePaddle等第三方框架很難無縫接入XLA的加速能力無法形成跨框架的統(tǒng)一編譯生態(tài)嚴重限制了XLA的普及范圍。簡單來說組合爆炸很多組件無法共用各個層次優(yōu)化難以遷移就是從XLA HLO直接到LLVM IR跨度太大四、MLIR 針對性解決的核心痛點MLIR通過“多層級方言Dialect”架構完美補全了XLA的短板新增Linalg、Affine等中間層方言打通了從HLO張量圖到底層循環(huán)優(yōu)化的完整鏈路讓XLA可以直接復用MLIR生態(tài)中成熟的循環(huán)分塊、向量化優(yōu)化Pass無需從零手寫代碼 。提供統(tǒng)一的IR基礎設施新硬件廠商只需實現(xiàn)對應MLIR方言的Lowering邏輯就能自動復用整條編譯流水線的所有優(yōu)化能力硬件接入成本降低一個數(shù)量級?；贛LIR推出StableHLO標準化方言成為跨框架的統(tǒng)一中間表示讓PyTorch、JAX等所有主流框架都能無縫接入XLA的加速能力徹底打破生態(tài)封閉的局面。簡單來說XLA解決了“在TPU上把AI模型跑快”的問題而MLIR解決了“讓XLA能適配所有硬件、所有框架同時大幅降低編譯器維護成本”的問題二者結合形成了現(xiàn)在工業(yè)界最成熟的AI編譯棧。簡單來說語義統(tǒng)一組件可以復用共享生態(tài)各個層次可以協(xié)調優(yōu)化層級之間跨度小方便實現(xiàn) 章節(jié)2：MLIR的引入
系列文章目錄章節(jié)1什么是MLIR章節(jié)2MLIR的引入章節(jié)3MLIR完成一個端到端的流程目錄系列文章目錄一、tensorflow執(zhí)行方式1. 前端構建計算圖與 Op 注冊2. 后端準備Kernel 注冊與設備匹配3. 核心調度Executor 執(zhí)行器4. 底層執(zhí)行Kernel Compute二、使用XLA加速tensorflow1.核心原理從“解釋執(zhí)行”到“編譯執(zhí)行”2.XLA 加速帶來的主要優(yōu)勢3.XLA和LLVM的關系三、XLA存在的問題四、MLIR 針對性解決的核心痛點提示以下是本篇文章正文內容下面案例可供參考一、tensorflow執(zhí)行方式執(zhí)行方式流程如下其中Graph? 定義了“做什么”?Executor? 決定了“何時做”和“在哪做”而 ?Kernel? 負責“具體怎么做”。三者協(xié)同工作實現(xiàn)了 TensorFlow 的張量計算能力。TensorFlow 從計算圖Graph的操作節(jié)點到執(zhí)行器Executor再到內核Kernel的執(zhí)行流程是一個典型的?“定義-優(yōu)化-調度-執(zhí)行”?的流水線過程。以下是這一過程的詳細拆解1. 前端構建計算圖與 Op 注冊當用戶在 Python 中調用如tf.matmul或tf.add時并沒有立即進行數(shù)值計算而是完成了以下工作?Op 定義與注冊?TensorFlow 內部維護了一個全局的OpDefLibrary。每個操作如 MatMul都有一個對應的 ?OpDefinition?定義了它的輸入輸出類型、屬性如 transpose_a等。這些定義通過宏如REGISTER_OP在 C 層注冊。?節(jié)點插入 Graph?Python API如gen_math_ops.py會將用戶傳入的參數(shù)轉換為底層 C 調用最終在當前的默認Graph對象中創(chuàng)建一個 ?NodeDef?。?NodeDef? 是圖的序列化表示僅包含操作的名稱、類型、輸入依賴和屬性?不包含具體的計算邏輯?。此時圖中的邊Edge建立了節(jié)點之間的數(shù)據(jù)依賴關系。2. 后端準備Kernel 注冊與設備匹配計算邏輯是由 ?Kernel? 實現(xiàn)的而不是 Op 本身。一個 Op 可以有多個 Kernel 實現(xiàn)分別對應不同的硬件設備CPU/GPU/TPU和數(shù)據(jù)類型float32/int32。?Kernel 注冊?開發(fā)者使用REGISTER_KERNEL_BUILDER宏將具體的 C 計算類繼承自OpKernel注冊到系統(tǒng)中。例如MatMul在 CPU 上有一個 kernel在 GPU 上有另一個基于 cuBLAS 的 kernel。?設備放置Device Placement?在會話Session初始化或圖執(zhí)行前TensorFlow 的運行時會分析圖中的每個 Node根據(jù)設備的可用性、內存限制和用戶指定的with tf.device()上下文決定每個節(jié)點由哪個設備Device執(zhí)行。?實例化 Kernel?一旦確定了設備運行時就會為該節(jié)點查找并實例化對應的OpKernel對象。這個對象包含了執(zhí)行計算所需的所有狀態(tài)和方法主要是Compute()函數(shù)。3. 核心調度Executor 執(zhí)行器?Executor? 是 TensorFlow 運行時的核心組件負責驅動整個計算圖的執(zhí)行。它的主要職責是將靜態(tài)的圖結構轉化為動態(tài)的執(zhí)行流。?圖優(yōu)化與分區(qū)?Executor 接收優(yōu)化后的 Graph。如果涉及多設備如 CPU 和 GPU它會將圖分割成多個子圖Partition每個子圖對應一個特定的設備。?構建執(zhí)行計劃?Executor 會對節(jié)點進行?拓撲排序?確定執(zhí)行的先后順序。它會識別哪些節(jié)點可以并行執(zhí)行即沒有數(shù)據(jù)依賴關系的節(jié)點。?異步調度模型?Executor 維護一個任務隊列。它不會等待一個節(jié)點完全結束才啟動下一個而是采用異步機制當一個節(jié)點的所有輸入張量Tensor就緒時該節(jié)點變?yōu)椤翱蓤?zhí)行”狀態(tài)。Executor 將該節(jié)點的計算任務封裝為一個回調函數(shù)提交給目標設備的線程池Thread Pool。對于 GPU這通常意味著將 CUDA Kernel 啟動指令放入 GPU 流Stream中。4. 底層執(zhí)行Kernel Compute當調度器將任務分發(fā)給具體設備后真正的計算發(fā)生在 ?Kernel? 層面?調用 Compute()?設備線程池中的工作線程會調用對應OpKernel的Compute(OpKernelContext* context)方法。?獲取輸入與分配輸出?OpKernelContext提供了訪問輸入張量數(shù)據(jù)的接口并負責為輸出張量分配內存通常在設備的顯存或內存中。?執(zhí)行具體算法??CPU Kernel?直接執(zhí)行 C 代碼可能調用 Eigen 庫或 MKL 進行矩陣運算。?GPU Kernel?通常是一個輕量級的包裝器負責準備參數(shù)并啟動實際的 CUDA/HIP 核函數(shù)如cublasSgemm。?結果回傳與同步?計算完成后Kernel 將結果寫入輸出張量。Executor 會監(jiān)控這些完成信號一旦某個節(jié)點的輸出就緒它會通知所有依賴該輸出的后續(xù)節(jié)點觸發(fā)它們的執(zhí)行。二、使用XLA加速tensorflow使用 XLA 加速?是指利用 ?XLAAccelerated Linear Algebra加速線性代數(shù)? 編譯器對深度學習模型的數(shù)學運算進行即時編譯和優(yōu)化從而顯著提升模型在特定硬件如 TPU、GPU上的訓練和推理速度。簡單來說它就像是一個“超級翻譯官”兼“優(yōu)化大師”將你寫的 Python/PyTorch/JAX 代碼轉換成針對硬件高度優(yōu)化的機器碼而不是讓硬件去逐條解釋執(zhí)行原始的 Python 指令。執(zhí)行方式如下1.核心原理從“解釋執(zhí)行”到“編譯執(zhí)行”傳統(tǒng)框架如原生 PyTorch 的 eager 模式或早期 TensorFlow通常采用?命令式執(zhí)行?每行代碼發(fā)出一個操作指令硬件立即執(zhí)行。這種方式靈活但效率低因為每個小操作都有啟動開銷且無法看到全局優(yōu)化空間。XLA 加速的核心機制是“惰性求值”與“圖編譯”??構建計算圖Tracing/Lazy Evaluation?當你運行模型時XLA 不會立即執(zhí)行每個算子而是先“記錄”下所有的操作序列形成一個中間表示IR的計算圖。這個過程被稱為“追蹤”Tracing。?整體編譯與優(yōu)化?XLA 編譯器拿到整個計算圖后會進行全局優(yōu)化。最關鍵的技術是?算子融合Operator Fusion?它將多個小的、連續(xù)的操作如Add-Relu-BatchNorm合并成一個單一的 GPU/TPU 內核Kernel。?生成高效機器碼?優(yōu)化后的圖被編譯成針對特定硬件如 NVIDIA GPU 的 CUDA 代碼或 Google TPU 的指令集的高效二進制代碼。?執(zhí)行?最后硬件直接執(zhí)行這段優(yōu)化后的代碼。由于減少了內核啟動次數(shù)和數(shù)據(jù)在顯存中的讀寫次數(shù)性能大幅提升。2.XLA 加速帶來的主要優(yōu)勢?算子融合Operator Fusion?這是 XLA 最核心的優(yōu)化手段。通過減少 GPU/TPU 內核啟動開銷和中間張量的內存讀寫顯著降低延遲并提高吞吐量。?內存優(yōu)化?XLA 能更智能地管理顯存復用緩沖區(qū)從而允許你訓練更大的模型或使用更大的 Batch Size。?硬件特異性優(yōu)化?XLA 針對不同硬件后端CPU、GPU、TPU生成專屬的最優(yōu)代碼。例如在 TPU 上XLA 能充分利用其矩陣乘法單元MXU的高并行能力在 NVIDIA GPU 上它能更好地利用 Tensor Cores。3.XLA和LLVM的關系XLA 和 LLVM 是典型的「高層領域專用編譯器」與「通用底層編譯器基礎設施」的深度協(xié)作關系?二者定位完全互補XLA 不重復造輪子而是基于 LLVM 成熟的底層能力專注于AI張量計算的專屬優(yōu)化。XLA 的完整編譯流水線天然以 LLVM 作為核心后端基座?XLA 前端?接收來自 TensorFlow/JAX 的高層張量計算圖生成 XLA 專屬的中間表示 ?HLOHigh Level Optimizer?完成算子融合、自動并行、內存復用等AI領域專屬的高層優(yōu)化。?HLO 轉譯層?將高層的 HLO 張量操作逐步展開、映射為符合 LLVM 規(guī)范的中間表示LLVM IR比如一個HLO矩陣乘法會被拆解為一系列LLVM的加載、乘加、存儲指令序列 。?LLVM 后端接管?生成的LLVM IR直接送入LLVM完整的工具鏈經(jīng)過指令調度、寄存器分配、目標平臺專屬優(yōu)化后最終輸出x86、ARM、CUDA等不同硬件架構的原生機器碼 。三、XLA存在的問題XLA 原生架構的核心不足?IR層級斷層優(yōu)化空間受限?XLA 只有HLO這一個核心中間表示高層張量圖優(yōu)化完成后直接硬著陸到LLVM IR或TPU硬件指令中間缺失了多層級過渡IR。這導致循環(huán)變換、多維數(shù)組分塊、仿射優(yōu)化這類關鍵的中間層變換無法優(yōu)雅實現(xiàn)只能靠手寫大量特殊Case的Emitter代碼優(yōu)化能力被嚴重鎖死 。?硬件適配成本極高擴展性差?XLA 為每一種新硬件后端如某款新ASIC、邊緣端NPU都需要從零開始手寫整套代碼生成邏輯復用率極低。傳統(tǒng)XLA的GPU代碼生成甚至重度依賴LLVM的元數(shù)據(jù)hack來傳遞優(yōu)化信息維護成本極高新硬件接入周期長達數(shù)月甚至數(shù)年 。?代碼庫碎片化維護難度指數(shù)級上升?XLA的所有變換邏輯都是用C硬編碼實現(xiàn)的沒有統(tǒng)一的IR基礎設施。隨著算子數(shù)量、硬件種類增加手寫的Emitter代碼量爆炸式增長公共子表達式消除、常量折疊這類基礎編譯器優(yōu)化都要重復造輪子無法復用成熟的通用編譯器生態(tài)能力 。?跨框架兼容性差生態(tài)封閉?早期XLA深度綁定TensorFlow沒有標準化的中間交換格式PyTorch、PaddlePaddle等第三方框架很難無縫接入XLA的加速能力無法形成跨框架的統(tǒng)一編譯生態(tài)嚴重限制了XLA的普及范圍。簡單來說組合爆炸很多組件無法共用各個層次優(yōu)化難以遷移就是從XLA HLO直接到LLVM IR跨度太大四、MLIR 針對性解決的核心痛點MLIR通過“多層級方言Dialect”架構完美補全了XLA的短板新增Linalg、Affine等中間層方言打通了從HLO張量圖到底層循環(huán)優(yōu)化的完整鏈路讓XLA可以直接復用MLIR生態(tài)中成熟的循環(huán)分塊、向量化優(yōu)化Pass無需從零手寫代碼 。提供統(tǒng)一的IR基礎設施新硬件廠商只需實現(xiàn)對應MLIR方言的Lowering邏輯就能自動復用整條編譯流水線的所有優(yōu)化能力硬件接入成本降低一個數(shù)量級?；贛LIR推出StableHLO標準化方言成為跨框架的統(tǒng)一中間表示讓PyTorch、JAX等所有主流框架都能無縫接入XLA的加速能力徹底打破生態(tài)封閉的局面。簡單來說XLA解決了“在TPU上把AI模型跑快”的問題而MLIR解決了“讓XLA能適配所有硬件、所有框架同時大幅降低編譯器維護成本”的問題二者結合形成了現(xiàn)在工業(yè)界最成熟的AI編譯棧。簡單來說語義統(tǒng)一組件可以復用共享生態(tài)各個層次可以協(xié)調優(yōu)化層級之間跨度小方便實現(xiàn)