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