簡介：隨著深度學習的不斷發展，AI模型結構在快速演化，底層計算硬件技術更是層出不窮，對于廣大開發者來說不僅要考慮如何在復雜多變的場景下有效的將算力發揮出來，還要應對計算框架的持續迭代。深度編譯器就成了應對以上問題廣受關注的技術方向，讓用戶僅需專注于上層模型開發，降低手工優化性能的人力開發成本，進一步壓榨硬件性能空間。阿里云機器學習PAI開源了業內較早投入實際業務應用的動態shape深度學習編譯器 BladeDISC，本文將詳解 BladeDISC的設計原理和應用。

作者 | 姜來
來源 | 阿里技術公眾號

一 導讀

隨著深度學習的不斷發展，AI模型結構在快速演化，底層計算硬件技術更是層出不窮，對于廣大開發者來說不僅要考慮如何在復雜多變的場景下有效的將算力發揮出來，還要應對計算框架的持續迭代。深度編譯器就成了應對以上問題廣受關注的技術方向，讓用戶僅需專注于上層模型開發，降低手工優化性能的人力開發成本，進一步壓榨硬件性能空間。阿里云機器學習PAI開源了業內較早投入實際業務應用的動態shape深度學習編譯器 BladeDISC，本文將詳解 BladeDISC的設計原理和應用。

二 BladeDISC是什么

BladeDISC是阿里最新開源的基于MLIR的動態shape深度學習編譯器。

1 主要特性

• 多款前端框架支持：TensorFlow，PyTorch
• 多后端硬件支持：CUDA，ROCM，x86
• 完備支持動態shape語義編譯
• 支持推理及訓練
• 輕量化API，對用戶通用透明
• 支持插件模式嵌入宿主框架運行，以及獨立部署模式

三 深度學習編譯器的背景

近幾年來，深度學習編譯器作為一個較新的技術方向異常地活躍，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，到后來呼聲很高的MLIR以及其不同領域的延伸項目IREE、mlir-hlo等等。能夠看到不同的公司、社區在這個領域進行著大量的探索和推進。

1 AI浪潮與芯片浪潮共同催生——從萌芽之初到蓬勃發展

深度學習編譯器近年來之所以能夠受到持續的關注，主要來自于幾個方面的原因：

框架性能優化在模型泛化性方面的需求

深度學習還在日新月異的發展之中，創新的應用領域不斷涌現，復雜多變的場景下如何有效的將硬件的算力充分發揮出來成為了整個AI應用鏈路中非常重要的一環。在早期，神經網絡部署的側重點在于框架和算子庫，這部分職責很大程度上由深度學習框架、硬件廠商提供的算子庫、以及業務團隊的手工優化工作來承擔。

上圖將近年的深度學習框架粗略分為三代。一個趨勢是在上層的用戶API層面，這些框架在變得越來越靈活，靈活性變強的同時也為底層性能問題提出了更大的挑戰。初代深度學習框架類似 Caffe 用 sequence of layer 的方式描述神經網絡結構，第二代類似 TensorFlow 用更細粒度的 graph of operators 描述計算圖，到第三代類似 PyTorch，TensorFlow Eager Mode 的動態圖。我們可以看到框架越來越靈活，描述能力越來越強，帶來的問題是優化底層性能變得越來越困難。業務團隊也經常需要補充完成所需要的手工優化，這些工作依賴于特定業務和對底層硬件的理解，耗費人力且難以泛化。而深度學習編譯器則是結合編譯時圖層的優化以及自動或者半自動的代碼生成，將手工優化的原理做泛化性的沉淀，以替代純手工優化帶來的各種問題，去解決深度學習框架的靈活性和性能之間的矛盾。

AI框架在硬件泛化性方面的需求

表面上看近些年AI發展的有目共睹、方興未艾，而后臺的硬件算力數十年的發展才是催化AI繁榮的核心動力。隨著晶體管微縮面臨的各種物理挑戰越來越大，芯片算力的增加越來越難，摩爾定律面臨失效，創新體系結構的各種DSA芯片迎來了一波熱潮，傳統的x86、ARM等都在不同的領域強化著自己的競爭力。硬件的百花齊放也為AI框架的發展帶來了新的挑戰。

而硬件的創新是一個問題，如何能把硬件的算力在真實的業務場景中發揮出來則是另外一個問題。新的AI硬件廠商不得不面臨除了要在硬件上創新，還要在軟件棧上做重度人力投入的問題。向下如何兼容硬件，成為了如今深度學習框架的核心難點之一，而兼容硬件這件事，則需要由編譯器來解決。

AI系統平臺對前端AI框架泛化性方面的需求

當今主流的深度學習框架包括Tensoflow、Pytorch、Keras、JAX等等，幾個框架都有各自的優缺點，在上層對用戶的接口方面風格各異，卻同樣面臨硬件適配以及充分發揮硬件算力的問題。不同的團隊常根據自己的建模場景和使用習慣來選擇不同的框架，而云廠商或者平臺方的性能優化工具和硬件適配方案則需要同時考慮不同的前端框架，甚至未來框架演進的需求。Google利用XLA同時支持TensorFlow和JAX，同時其它開源社區還演進出了Torch_XLA，Torch-MLIR這樣的接入方案，這些接入方案目前雖然在易用性和成熟程度等方面還存在一些問題，卻反應出AI系統層的工作對前端AI框架泛化性方面的需求和技術趨勢。

2 什么是深度學習編譯器

傳統編譯器是以高層語言作為輸入，避免用戶直接去寫機器碼，而用相對靈活高效的語言來工作，并在編譯的過程中引入優化來解決高層語言引入的性能問題，平衡開發效率與性能之間的矛盾。而深度學習編譯器的作用相仿，其輸入是比較靈活的，具備較高抽象度的計算圖描述，輸出包括CPU、GPU及其他異構硬件平臺上的底層機器碼及執行引擎。

傳統編譯器的使命之一是減輕編程者的壓力。作為編譯器的輸入的高級語言往往更多地是描述一個邏輯，為了便利編程者，高級語言的描述會更加抽象和靈活，至于這個邏輯在機器上是否能夠高效的執行，往往是考驗編譯器的一個重要指標。深度學習作為一個近幾年發展異常快速的應用領域，它的性能優化至關重要，并且同樣存在高層描述的靈活性和抽象性與底層計算性能之間的矛盾，因此專門針對深度學習的編譯器出現了。而傳統編譯器的另外一個重要使命是，需要保證編程者輸入的高層語言能夠執行在不同體系結構和指令集的硬件計算單元上，這一點也同樣反應在深度學習編譯器上。面對一個新的硬件設備，人工的話不太可能有精力對那么多目標硬件重新手寫一個框架所需要的全部算子實現，深度學習編譯器提供中間層的IR，將頂層框架的模型流圖轉化成中間層表示IR，在中間層IR上進行通用的圖層優化，而在后端將優化后的IR通用性的生成各個目標平臺的機器碼。

深度學習編譯器的目標是針對AI計算任務，以通用編譯器的方式完成性能優化和硬件適配。讓用戶可以專注于上層模型開發，降低用戶手工優化性能的人力開發成本，進一步壓榨硬件性能空間。

3 距離大規模應用面臨的瓶頸問題

深度學習編譯器發展到今天，雖然在目標和技術架構方面與傳統編譯器有頗多相似之處，且在技術方向上表現出了良好的潛力，然而目前的實際應用范圍卻仍然距離傳統編譯器有一定的差距，主要難點包括：

易用性

深度學習編譯器的初衷是希望簡化手工優化性能和適配硬件的人力成本。然而在現階段，大規模部署應用深度學習編譯器的挑戰仍然較大，能夠將編譯器用好的門檻較高，造成這個現象的主要原因包括：

• 與前端框架對接的問題。不同框架對深度學習任務的抽象描述和API接口各有不同，語義和機制上有各自的特點，且作為編譯器輸入的前端框架的算子類型數量呈開放性狀態。如何在不保證所有算子被完整支持的情況下透明化的支持用戶的計算圖描述，是深度學習編譯器能夠易于為用戶所廣泛使用的重要因素之一。
• 動態shape問題和動態計算圖問題。現階段主流的深度學習編譯器主要針對特定的靜態shape輸入完成編譯，此外對包含control flow語義的動態計算圖只能提供有限的支持或者完全不能夠支持。而AI的應用場景卻恰恰存在大量這一類的任務需求。這時只能人工將計算圖改寫為靜態或者半靜態的計算圖，或者想辦法將適合編譯器的部分子圖提取出來交給編譯器。這無疑加重了應用深度學習編譯器時的工程負擔。更嚴重的問題是，很多任務類型并不能通過人工的改寫來靜態化，這導致這些情況下編譯器完全無法實際應用。
• 編譯開銷問題。作為性能優化工具的深度學習編譯器只有在其編譯開銷對比帶來的性能收益有足夠優勢的情況下才真正具有實用價值。部分應用場景下對于編譯開銷的要求較高，例如普通規模的需要幾天時間完成訓練任務有可能無法接受幾個小時的編譯開銷。對于應用工程師而言，使用編譯器的情況下不能快速的完成模型的調試，也增加了開發和部署的難度和負擔。
• 對用戶透明性問題。部分AI編譯器并非完全自動的編譯工具，其性能表現比較依賴于用戶提供的高層抽象的實現模版。主要是為算子開發工程師提供效率工具，降低用戶人工調優各種算子實現的人力成本。但這也對使用者的算子開發經驗和對硬件體系結構的熟悉程度提出了比較高的要求。此外，對于新硬件的軟件開發者來說，現有的抽象卻又常常無法足夠描述創新的硬件體系結構上所需要的算子實現。需要對編譯器架構足夠熟悉的情況下對其進行二次開發甚至架構上的重構，門檻及開發負擔仍然很高。

魯棒性

目前主流的幾個AI編譯器項目大多數都還處于偏實驗性質的產品，但產品的成熟度距離工業級應用有較大的差距。這里的魯棒性包含是否能夠順利完成輸入計算圖的編譯，計算結果的正確性，以及性能上避免coner case下的極端badcase等。

性能問題

編譯器的優化本質上是將人工的優化方法，或者人力不易探究到的優化方法通過泛化性的沉淀和抽象，以有限的編譯開銷來替代手工優化的人力成本。然而如何沉淀和抽象的方法學是整個鏈路內最為本質也是最難的問題。深度學習編譯器只有在性能上真正能夠代替或者超過人工優化，或者真正能夠起到大幅度降低人力成本作用的情況下才能真正發揮它的價值。

然而達到這一目標卻并非易事，深度學習任務大都是tensor級別的計算，對于并行任務的拆分方式有很高的要求，但如何將手工的優化泛化性的沉淀在編譯器技術內，避免編譯開銷爆炸以及分層之后不同層級之間優化的聯動，仍然有更多的未知需要去探索和挖掘。這也成為以MLIR框架為代表的下一代深度學習編譯器需要著重思考和解決的問題。

四 BladeDISC的主要技術特點

項目最早啟動的初衷是為了解決XLA和TVM當前版本的靜態shape限制，內部命名為 DISC (DynamIc Shape Compiler)，希望打造一款能夠在實際業務中使用的完備支持動態shape語義的深度學習編譯器。

從團隊四年前啟動深度學習編譯器方向的工作以來，動態shape問題一直是阻礙實際業務落地的嚴重問題之一。彼時，包括XLA在內的主流深度學習框架，都是基于靜態shape語義的編譯器框架。典型的方案是需要用戶指定輸入的shape，或是由編譯器在運行時捕捉待編譯子圖的實際輸入shape組合，并且為每一個輸入shape組合生成一份編譯結果。

靜態shape編譯器的優勢顯而易見，編譯期完全已知靜態shape信息的情況下，Compiler可以作出更好的優化決策并得到更好的CodeGen性能，同時也能夠得到更好的顯存/內存優化plan和調度執行plan。然而，其缺點也十分明顯，具體包括：

• 大幅增加編譯開銷。引入離線編譯預熱過程，大幅增加推理任務部署過程復雜性；訓練迭代速度不穩定甚至整體訓練時間負優化。
• 部分業務場景shape變化范圍趨于無窮的，導致編譯緩存永遠無法收斂，方案不可用。
• 內存顯存占用的增加。編譯緩存額外占用的內存顯存，經常導致實際部署環境下的內存/顯存OOM，直接阻礙業務的實際落地。
• 人工padding為靜態shape等緩解性方案對用戶不友好，大幅降低應用的通用性和透明性，影響迭代效率。

在2020年夏天，DISC完成了僅支持TensorFlow前端以及Nvidia GPU后端的初版，并且正式在阿里內部上線投入實際應用。最早在幾個受困于動態shape問題已久的業務場景上投入使用，并且得到了預期中的效果。即在一次編譯且不需要用戶對計算圖做特殊處理的情況下，完備支持動態shape語義，且性能幾乎與靜態shape編譯器持平。對比TensorRT等基于手工算子庫為主的優化框架，DISC基于編譯器自動codegen的技術架構在經常為非標準開源模型的實際業務上獲得了明顯的性能和易用性優勢。

從2020年第二季度開始至今，DISC持續投入研發力量，針對前文提到的從云端平臺方視角看到的深度學習編譯器距離大規模部署和應用的幾個瓶頸問題，在性能、算子覆蓋率和魯棒性、CPU及新硬件支持、前端框架支持等方面逐漸完善。目前在場景覆蓋能力和性能等方面，已經逐漸替換掉團隊過往基于XLA和TVM等靜態shape框架上的工作，成為PAI-Blade支持阿里內部及阿里云外部業務的主要優化手段。2021年后，DISC在CPU及GPGPU體系結構的后端硬件上的性能有了顯著的提升，同時在新硬件的支持上面投入了更多的技術力量。2021年底，為了吸引更多的技術交流和合作共建需要，以及更大范圍的用戶反饋，正式更名為BladeDISC并完成了初版開源。

五 關鍵技術

BladeDISC的整體架構，及其在阿里云相關產品中的上下文關系如下圖所示：

1 MLIR基礎架構

MLIR是由Google在2019年發起的項目，MLIR 的核心是一套靈活的多層IR基礎設施和編譯器實用工具庫，深受 LLVM 的影響，并重用其許多優秀理念。這里我們選擇基于MLIR的主要原因包括其比較豐富的基礎設施支持，方便擴展的模塊化設計架構以及MLIR較強的膠水能力。

2 動態shape編譯

上圖為BladeDISC的主體Pass Pipeline設計。對比目前主流的深度學習編譯器項目，主要技術特點如下：

圖層IR設計

BladeDISC選擇基于HLO作為核心圖層IR來接入不同的前端框架，但是HLO是原本為XLA設計的純靜態shape語義的IR。靜態場景下，HLO IR中的shape表達會被靜態化，所有的shape計算會被固化為編譯時常量保留在編譯結果中；而在動態shape場景下，IR本身需要有足夠的能力表達shape計算和動態shape信息的傳遞。BladeDISC從項目建立開始一直與MHLO社區保持緊密的合作，在XLA的HLO IR基礎上，擴展了一套具有完備動態shape表達能力的IR，并增加了相應的基礎設施以及前端框架的算子轉換邏輯。這部分實現目前已經完整upstream至MHLO社區，確保后續其它MHLO相關項目中IR的一致性。

運行時Shape計算、存儲管理和Kernel調度

動態shape編譯的主要挑戰來自于需要在靜態的編譯過程中能夠處理動態的計算圖語義。為完備支持動態shape，編譯結果需要能夠在運行時做實時的shape推導計算，不僅要為數據計算，同時也需要為shape計算做代碼生成。計算后的shape信息用于做內存/顯存管理，以及kernel調度時的參數選擇等等。BladeDISC的pass pipeline的設計充分考慮了上述動態shape語義支持的需求，采用了host-device聯合codegen的方案。以GPU Backend為例，包括shape計算、內存/顯存申請釋放、硬件管理、kernel launch運行時流程全部為自動代碼生成，以期得到完備的動態shape端到端支持方案和更為極致的整體性能。

動態shape下的性能問題

在shape未知或者部分未知的情況下，深度學習編譯器在性能上面臨的挑戰被進一步放大。在大多數主流硬件backend上，BladeDISC采用區分計算密集型部分和訪存密集型部分的策略，以期在性能與復雜性和編譯開銷之間獲取更好的平衡。

對于計算密集型部分，不同的shape要求更加精細的schedule實現來獲得更好的性能，pass pipeline在設計上的主要考慮是需要支持在運行時根據不同的具體shape選擇合適的算子庫實現，以及處理動態shape語義下的layout問題。

而訪存密集型部分的自動算子融合作為深度學習編譯器主要的性能收益來源之一，同樣面臨shape未知情況下在性能上的挑戰。許多靜態shape語義下比較確定性的問題，例如指令層的向量化，codegen模版選擇，是否需要implicit broadcast等等在動態shape場景下都會面臨更大的復雜性。針對這些方面的問題，BladeDISC選擇將部分的優化決策從編譯時下沉到運行時。即在編譯期根據一定的規則生成多個版本的kernel實現，在運行時根據實際shape自動選擇最優的實現。這一機制被稱作speculation，在BladeDISC內基于host-device的聯合代碼生成來實現。此外，在編譯期沒有具體shape數值的情況下，會很容易在各個層級丟失掉大量的優化機會，從圖層的線性代數簡化、fusion決策到指令層級的CSE、常數折疊等。BladeDISC在IR及pass pipeline的設計過程中著重設計了shape constraint在IR中的抽象和在pass pipeline中的使用，例如編譯期未知的不同dimension size之間的約束關系等。在優化整體性能方面起到了比較明顯的作用，保證能夠足夠接近甚至超過靜態shape編譯器的性能結果。

大顆粒度算子融合

團隊在開啟BladeDISC項目之前，曾經基于靜態shape編譯器在大顆粒度算子融合及自動代碼生成方面有過若干探索3，其基本思想可以概括為借助于GPU硬件中低訪存開銷的shared memory或CPU中低訪存開銷的Memory Cache，將不同schedule的計算子圖縫合進同一個kernel內，實現多個parallel loop復合，這種codegen方法稱之為fusion-stitching。這種訪存密集型子圖的自動代碼生成打破了常規的loop fusion，input/output fusion對fusion顆粒度的限制。在保證代碼生成質量的同時，大幅增加fusion顆粒度，同時避免復雜性及編譯開銷爆炸。且整個過程完全對用戶透明，無需人工指定schedule描述。

在動態shape語義下實現fusion-stitching對比靜態shape語義下同樣需要處理更大的復雜性，動態shape語義下的shape constraint抽象一定程度上簡化了這一復雜性，使整體性能進一步接近甚至超過手工算子實現。

3 多前端框架支持

AICompiler框架在設計時也包含了擴展支持不同前端框架的考慮。PyTorch側通過實現一個輕量的Converter將TorchScript轉換為DHLO IR實現了對PyTorch推理作業的覆蓋。MLIR相對完備的IR基礎設施也為Converter的實現提供了便利。BladeDISC包含Compiler以及適配不同前端框架的Bridge側兩部分。其中Bridge進一步分為宿主框架內的圖層pass以及運行時Op兩部分，以插件的方式接入宿主框架。這種工作方式使BladeDISC可以透明化的支持前端計算圖，可以適配用戶各種版本的宿主框架。

4 運行時環境適配

為將編譯的結果能夠配合TensorFlow/PyTorch等宿主在各自的運行環境中執行起來，以及管理運行時IR層不易表達的狀態信息等等，我們為不同的運行時環境實現了一套統一的Compiler架構，并引入了運行時抽象層，即RAL(Runtime Abstraction Layer)層。

RAL實現了多種運行環境的適配支持，用戶可以根據需要進行選擇，具體包括：

• 全圖編譯，獨立運行。當整個計算圖都支持編譯時，RAL提供了一套簡易的runtime以及在此之上RAL Driver的實現，使得compiler編譯出來結果可以脫離框架直接運行，減少框架overhead。
• TF中子圖編譯運行。
• Pytorch中子圖編譯運行。

以上環境中在諸如資源管理，API語義等上存在差異，RAL通過抽象出一套最小集合的API ，并清晰的定義出它們的語義，將編譯器與運行時隔離開來，來達到在不同的環境中都能夠執行編譯出來的結果的目的。此外RAL層實現了無狀態編譯，解決了計算圖的編譯之后，編譯的結果可能被多次執行時的狀態信息處理問題。一方面簡化了代碼生成的復雜度，另一方面也更容易支持多線程并發執行（比如推理）的場景，同時在錯誤處理，回滾方面也更加容易支持。

六 應用場景

BladeDISC的典型應用場景可以不太嚴格的分為兩類：其一是在主流的硬件平臺上（包括Nvidia GPU，x86 CPU等）上作為通用、透明的性能優化工具，降低用戶部署AI作業的人力負擔，提高模型迭代效率；另一個重要的應用場景是幫助新硬件做AI場景的適配和接入支持。

目前BladeDISC已經廣泛應用在阿里內部和阿里云上外部用戶的多個不同應用場景下，覆蓋模型類型涉及NLP、機器翻譯、語音類ASR、語音TTS、圖像檢測、識別、AI for science等等多種典型AI應用；覆蓋行業包括互聯網、電商、自動駕駛、安全行業、在線娛樂、醫療和生物等等。

在推理場景下，BladeDISC與TensorRT等廠商提供的推理優化工具形成良好的技術互補，其主要差異性優勢包括：

• 應對動態shape業務完備的動態shape語義支持
• 基于compiler based的技術路徑的模型泛化性在非標準模型上的性能優勢
• 更為靈活的部署模式選擇，以插件形式支持前端框架的透明性優勢

下圖為Nvidia T4硬件上幾個真實的業務例子的性能收益數字：

在新硬件支持方面，目前普遍的情況是除了積累比較深厚的Nvidia等頭部廠商之外，包括ROCM等其它GPGPU硬件普遍存在的情況是硬件的指標已經具備相當的競爭力，但廠商受制于AI軟件棧上的積累相對較少，普遍存在硬件算力無法發揮出來導致硬件落地應用困難的問題。如前文所述，基于編譯器的技術路徑下天然對于硬件的后端具備一定的泛化能力，且與硬件廠商的技術儲備形成比較強的互補。BladeDISC目前在GPGPU和通用CPU體系結構上的儲備相對比較成熟。以GPGPU為例，在Nvidia GPU上的絕大部分技術棧可以遷移至海光DCU和AMD GPU等體系結構相近的硬件上。BladeDISC較強的硬件泛化能力配合硬件本身較強的通用性，很好的解決了新硬件適配的性能和可用性問題。

下圖為海光DCU上幾個真實業務例子上的性能數字：

七 開源生態——構想和未來

我們決定建設開源生態主要有如下的考慮：

• BladeDISC發源于阿里云計算平臺團隊的業務需求，在開發過程中與MLIR/MHLO/IREE等社區同行之間的討論和交流給了我們很好的輸入和借鑒。在我們自身隨著業務需求的迭代逐漸完善的同時，也希望能夠開源給社區，在目前整個AI編譯器領域實驗性項目居多，偏實用性強的產品偏少，且不同技術棧之間的工作相對碎片化的情況下，希望能夠將自身的經驗和理解也同樣回饋給社區，希望和深度學習編譯器的開發者和AI System的從業者之間有更多更好的交流和共建，為這個行業貢獻我們的技術力量。
• 我們希望能夠借助開源的工作，收到更多真實業務場景下的用戶反饋，以幫助我們持續完善和迭代，并為后續的工作投入方向提供輸入。

后續我們計劃以兩個月為單位定期發布Release版本。BladeDISC近期的Roadmap如下：

• 持續的魯棒性及性能改進
• x86后端補齊計算密集型算子的支持，端到端完整開源x86后端的支持
• GPGPU上基于Stitching的大顆粒度自動代碼生成
• AMD rocm GPU后端的支持
• PyTorch訓練場景的支持

此外，在中長期，我們在下面幾個探索性的方向上會持續投入精力，也歡迎各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社區建設感興趣的同行一起參與共建。

• 更多新硬件體系結構的支持和適配，以及新硬件體系結構下軟硬件協同方法學的沉淀
• 計算密集型算子自動代碼生成和動態shape語義下全局layout優化的探索
• 稀疏子圖的優化探索
• 動態shape語義下運行時調度策略、內存/顯存優化等方面的探索
• 模型壓縮與編譯優化聯合的技術探索
• 圖神經網絡等更多AI作業類型的支持和優化等

原文鏈接

本文為阿里云原創內容，未經允許不得轉載。?

總結

以上是生活随笔為你收集整理的阿里 BladeDISC 深度学习编译器正式开源的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。