作者：朱凱
導讀
隨著深度學習的不斷發展，AI模型結構在快速演化，底層計算硬件技術更是層出不窮，對于廣大開發者來說不僅要考慮如何在復雜多變的場景下有效的將算力發揮出來，還要應對計算框架的持續迭代。深度編譯器就成了應對以上問題廣受關注的技術方向，讓用戶僅需專注于上層模型開發，降低手工優化性能的人力開發成本，進一步壓榨硬件性能空間。阿里云機器學習PAI開源了業內較早投入實際業務應用的動態shape深度學習編譯器 BladeDISC，本文將詳解 BladeDISC的設計原理和應用。
BladeDISC是什么
BladeDISC是阿里最新開源的基于MLIR的動態shape深度學習編譯器。
主要特性
· 多款前端框架支持：TensorFlow，PyTorch
· 多后端硬件支持：CUDA，ROCM，x86
· 完備支持動態shape語義編譯
· 支持推理及訓練
· 輕量化API，對用戶通用透明
· 支持插件模式嵌入宿主框架運行，以及獨立部署模式
開源地址：
https://github.com/alibaba/BladeDISC
深度學習編譯器的背景
近幾年來，深度學習編譯器作為一個較新的技術方向異常地活躍，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，到后來呼聲很高的MLIR以及其不同領域的延伸項目IREE、mlir-hlo等等。能夠看到不同的公司、社區在這個領域進行著大量的探索和推進。
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框架泛化性方面的需求和技術趨勢。
什么是深度學習編譯器
傳統編譯器是以高層語言作為輸入，避免用戶直接去寫機器碼，而用相對靈活高效的語言來工作，并在編譯的過程中引入優化來解決高層語言引入的性能問題，平衡開發效率與性能之間的矛盾。而深度學習編譯器的作用相仿，其輸入是比較靈活的，具備較高抽象度的計算圖描述，輸出包括CPU、GPU及其他異構硬件平臺上的底層機器碼及執行引擎。
傳統編譯器的使命之一是減輕編程者的壓力。作為編譯器的輸入的高級語言往往更多地是描述一個邏輯，為了便利編程者，高級語言的描述會更加抽象和靈活，至于這個邏輯在機器上是否能夠高效的執行，往往是考驗編譯器的一個重要指標。深度學習作為一個近幾年發展異常快速的應用領域，它的性能優化至關重要，并且同樣存在高層描述的靈活性和抽象性與底層計算性能之間的矛盾，因此專門針對深度學習的編譯器出現了。而傳統編譯器的另外一個重要使命是，需要保證編程者輸入的高層語言能夠執行在不同體系結構和指令集的硬件計算單元上，這一點也同樣反應在深度學習編譯器上。面對一個新的硬件設備，人工的話不太可能有精力對那么多目標硬件重新手寫一個框架所需要的全部算子實現，深度學習編譯器提供中間層的IR，將頂層框架的模型流圖轉化成中間層表示IR，在中間層IR上進行通用的圖層優化，而在后端將優化后的IR通用性的生成各個目標平臺的機器碼。
深度學習編譯器的目標是針對AI計算任務，以通用編譯器的方式完成性能優化和硬件適配。讓用戶可以專注于上層模型開發，降低用戶手工優化性能的人力開發成本，進一步壓榨硬件性能空間。
距離大規模應用面臨的瓶頸問題
深度學習編譯器發展到今天，雖然在目標和技術架構方面與傳統編譯器有頗多相似之處，且在技術方向上表現出了良好的潛力，然而目前的實際應用范圍卻仍然距離傳統編譯器有一定的差距，主要難點包括：
易用性
深度學習編譯器的初衷是希望簡化手工優化性能和適配硬件的人力成本。然而在現階段，大規模部署應用深度學習編譯器的挑戰仍然較大，能夠將編譯器用好的門檻較高，造成這個現象的主要原因包括：
· 與前端框架對接的問題。不同框架對深度學習任務的抽象描述和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等緩解性方案對用戶不友好，大幅降低應用的通用性和透明性，影響迭代效率。
星空人工智能技術網 倡導尊重與保護知識產權。如發現本站文章存在版權等問題，煩請30天內提供版權疑問、身份證明、版權證明、聯系方式等發郵件至1851688011@qq.com我們將及時溝通與處理。！：首頁 > 新聞 » 阿里 BladeDISC 深度學習編譯器正式開源