一 導(dǎo)讀

隨著深度學(xué)習(xí)的不斷發(fā)展，AI模型結(jié)構(gòu)在快速演化，底層計(jì)算硬件技術(shù)更是層出不窮，對(duì)于廣大開(kāi)發(fā)者來(lái)說(shuō)不僅要考慮如何在復(fù)雜多變的場(chǎng)景下有效的將算力發(fā)揮出來(lái)，還要應(yīng)對(duì)計(jì)算框架的持續(xù)迭代。深度編譯器就成了應(yīng)對(duì)以上問(wèn)題廣受關(guān)注的技術(shù)方向，讓用戶僅需專注于上層模型開(kāi)發(fā)，降低手工優(yōu)化性能的人力開(kāi)發(fā)成本，進(jìn)一步壓榨硬件性能空間。阿里云機(jī)器學(xué)習(xí)PAI開(kāi)源了業(yè)內(nèi)較早投入實(shí)際業(yè)務(wù)應(yīng)用的動(dòng)態(tài)shape深度學(xué)習(xí)編譯器 BladeDISC，本文將詳解 BladeDISC的設(shè)計(jì)原理和應(yīng)用。

二 BladeDISC是什么

BladeDISC是阿里最新開(kāi)源的基于MLIR的動(dòng)態(tài)shape深度學(xué)習(xí)編譯器。

1 主要特性

• 多款前端框架支持：TensorFlow，PyTorch
• 多后端硬件支持：CUDA，ROCM，x86
• 完備支持動(dòng)態(tài)shape語(yǔ)義編譯
• 支持推理及訓(xùn)練
• 輕量化API，對(duì)用戶通用透明
• 支持插件模式嵌入宿主框架運(yùn)行，以及獨(dú)立部署模式

三 深度學(xué)習(xí)編譯器的背景

近幾年來(lái)，深度學(xué)習(xí)編譯器作為一個(gè)較新的技術(shù)方向異常地活躍，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，到后來(lái)呼聲很高的MLIR以及其不同領(lǐng)域的延伸項(xiàng)目IREE、mlir-hlo等等。能夠看到不同的公司、社區(qū)在這個(gè)領(lǐng)域進(jìn)行著大量的探索和推進(jìn)。

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

深度學(xué)習(xí)編譯器近年來(lái)之所以能夠受到持續(xù)的關(guān)注，主要來(lái)自于幾個(gè)方面的原因：

框架性能優(yōu)化在模型泛化性方面的需求

深度學(xué)習(xí)還在日新月異的發(fā)展之中，創(chuàng)新的應(yīng)用領(lǐng)域不斷涌現(xiàn)，復(fù)雜多變的場(chǎng)景下如何有效的將硬件的算力充分發(fā)揮出來(lái)成為了整個(gè)AI應(yīng)用鏈路中非常重要的一環(huán)。在早期，神經(jīng)網(wǎng)絡(luò)部署的側(cè)重點(diǎn)在于框架和算子庫(kù)，這部分職責(zé)很大程度上由深度學(xué)習(xí)框架、硬件廠商提供的算子庫(kù)、以及業(yè)務(wù)團(tuán)隊(duì)的手工優(yōu)化工作來(lái)承擔(dān)。

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

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

表面上看近些年AI發(fā)展的有目共睹、方興未艾，而后臺(tái)的硬件算力數(shù)十年的發(fā)展才是催化AI繁榮的核心動(dòng)力。隨著晶體管微縮面臨的各種物理挑戰(zhàn)越來(lái)越大，芯片算力的增加越來(lái)越難，摩爾定律面臨失效，創(chuàng)新體系結(jié)構(gòu)的各種DSA芯片迎來(lái)了一波熱潮，傳統(tǒng)的x86、ARM等都在不同的領(lǐng)域強(qiáng)化著自己的競(jìng)爭(zhēng)力。硬件的百花齊放也為AI框架的發(fā)展帶來(lái)了新的挑戰(zhàn)。

而硬件的創(chuàng)新是一個(gè)問(wèn)題，如何能把硬件的算力在真實(shí)的業(yè)務(wù)場(chǎng)景中發(fā)揮出來(lái)則是另外一個(gè)問(wèn)題。新的AI硬件廠商不得不面臨除了要在硬件上創(chuàng)新，還要在軟件棧上做重度人力投入的問(wèn)題。向下如何兼容硬件，成為了如今深度學(xué)習(xí)框架的核心難點(diǎn)之一，而兼容硬件這件事，則需要由編譯器來(lái)解決。

AI系統(tǒng)平臺(tái)對(duì)前端AI框架泛化性方面的需求

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

2 什么是深度學(xué)習(xí)編譯器

傳統(tǒng)編譯器是以高層語(yǔ)言作為輸入，避免用戶直接去寫(xiě)機(jī)器碼，而用相對(duì)靈活高效的語(yǔ)言來(lái)工作，并在編譯的過(guò)程中引入優(yōu)化來(lái)解決高層語(yǔ)言引入的性能問(wèn)題，平衡開(kāi)發(fā)效率與性能之間的矛盾。而深度學(xué)習(xí)編譯器的作用相仿，其輸入是比較靈活的，具備較高抽象度的計(jì)算圖描述，輸出包括CPU、GPU及其他異構(gòu)硬件平臺(tái)上的底層機(jī)器碼及執(zhí)行引擎。

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

深度學(xué)習(xí)編譯器的目標(biāo)是針對(duì)AI計(jì)算任務(wù)，以通用編譯器的方式完成性能優(yōu)化和硬件適配。讓用戶可以專注于上層模型開(kāi)發(fā)，降低用戶手工優(yōu)化性能的人力開(kāi)發(fā)成本，進(jìn)一步壓榨硬件性能空間。

3 距離大規(guī)模應(yīng)用面臨的瓶頸問(wèn)題

深度學(xué)習(xí)編譯器發(fā)展到今天，雖然在目標(biāo)和技術(shù)架構(gòu)方面與傳統(tǒng)編譯器有頗多相似之處，且在技術(shù)方向上表現(xiàn)出了良好的潛力，然而目前的實(shí)際應(yīng)用范圍卻仍然距離傳統(tǒng)編譯器有一定的差距，主要難點(diǎn)包括：

易用性

深度學(xué)習(xí)編譯器的初衷是希望簡(jiǎn)化手工優(yōu)化性能和適配硬件的人力成本。然而在現(xiàn)階段，大規(guī)模部署應(yīng)用深度學(xué)習(xí)編譯器的挑戰(zhàn)仍然較大，能夠?qū)⒕幾g器用好的門檻較高，造成這個(gè)現(xiàn)象的主要原因包括：

與前端框架對(duì)接的問(wèn)題。不同框架對(duì)深度學(xué)習(xí)任務(wù)的抽象描述和API接口各有不同，語(yǔ)義和機(jī)制上有各自的特點(diǎn)，且作為編譯器輸入的前端框架的算子類型數(shù)量呈開(kāi)放性狀態(tài)。如何在不保證所有算子被完整支持的情況下透明化的支持用戶的計(jì)算圖描述，是深度學(xué)習(xí)編譯器能夠易于為用戶所廣泛使用的重要因素之一。

動(dòng)態(tài)shape問(wèn)題和動(dòng)態(tài)計(jì)算圖問(wèn)題?，F(xiàn)階段主流的深度學(xué)習(xí)編譯器主要針對(duì)特定的靜態(tài)shape輸入完成編譯，此外對(duì)包含control flow語(yǔ)義的動(dòng)態(tài)計(jì)算圖只能提供有限的支持或者完全不能夠支持。而AI的應(yīng)用場(chǎng)景卻恰恰存在大量這一類的任務(wù)需求。這時(shí)只能人工將計(jì)算圖改寫(xiě)為靜態(tài)或者半靜態(tài)的計(jì)算圖，或者想辦法將適合編譯器的部分子圖提取出來(lái)交給編譯器。這無(wú)疑加重了應(yīng)用深度學(xué)習(xí)編譯器時(shí)的工程負(fù)擔(dān)。更嚴(yán)重的問(wèn)題是，很多任務(wù)類型并不能通過(guò)人工的改寫(xiě)來(lái)靜態(tài)化，這導(dǎo)致這些情況下編譯器完全無(wú)法實(shí)際應(yīng)用。

編譯開(kāi)銷問(wèn)題。作為性能優(yōu)化工具的深度學(xué)習(xí)編譯器只有在其編譯開(kāi)銷對(duì)比帶來(lái)的性能收益有足夠優(yōu)勢(shì)的情況下才真正具有實(shí)用價(jià)值。部分應(yīng)用場(chǎng)景下對(duì)于編譯開(kāi)銷的要求較高，例如普通規(guī)模的需要幾天時(shí)間完成訓(xùn)練任務(wù)有可能無(wú)法接受幾個(gè)小時(shí)的編譯開(kāi)銷。對(duì)于應(yīng)用工程師而言，使用編譯器的情況下不能快速的完成模型的調(diào)試，也增加了開(kāi)發(fā)和部署的難度和負(fù)擔(dān)。

對(duì)用戶透明性問(wèn)題。部分AI編譯器并非完全自動(dòng)的編譯工具，其性能表現(xiàn)比較依賴于用戶提供的高層抽象的實(shí)現(xiàn)模版。主要是為算子開(kāi)發(fā)工程師提供效率工具，降低用戶人工調(diào)優(yōu)各種算子實(shí)現(xiàn)的人力成本。但這也對(duì)使用者的算子開(kāi)發(fā)經(jīng)驗(yàn)和對(duì)硬件體系結(jié)構(gòu)的熟悉程度提出了比較高的要求。此外，對(duì)于新硬件的軟件開(kāi)發(fā)者來(lái)說(shuō)，現(xiàn)有的抽象卻又常常無(wú)法足夠描述創(chuàng)新的硬件體系結(jié)構(gòu)上所需要的算子實(shí)現(xiàn)。需要對(duì)編譯器架構(gòu)足夠熟悉的情況下對(duì)其進(jìn)行二次開(kāi)發(fā)甚至架構(gòu)上的重構(gòu)，門檻及開(kāi)發(fā)負(fù)擔(dān)仍然很高。

魯棒性

目前主流的幾個(gè)AI編譯器項(xiàng)目大多數(shù)都還處于偏實(shí)驗(yàn)性質(zhì)的產(chǎn)品，但產(chǎn)品的成熟度距離工業(yè)級(jí)應(yīng)用有較大的差距。這里的魯棒性包含是否能夠順利完成輸入計(jì)算圖的編譯，計(jì)算結(jié)果的正確性，以及性能上避免coner case下的極端badcase等。

性能問(wèn)題

編譯器的優(yōu)化本質(zhì)上是將人工的優(yōu)化方法，或者人力不易探究到的優(yōu)化方法通過(guò)泛化性的沉淀和抽象，以有限的編譯開(kāi)銷來(lái)替代手工優(yōu)化的人力成本。然而如何沉淀和抽象的方法學(xué)是整個(gè)鏈路內(nèi)最為本質(zhì)也是最難的問(wèn)題。深度學(xué)習(xí)編譯器只有在性能上真正能夠代替或者超過(guò)人工優(yōu)化，或者真正能夠起到大幅度降低人力成本作用的情況下才能真正發(fā)揮它的價(jià)值。

然而達(dá)到這一目標(biāo)卻并非易事，深度學(xué)習(xí)任務(wù)大都是tensor級(jí)別的計(jì)算，對(duì)于并行任務(wù)的拆分方式有很高的要求，但如何將手工的優(yōu)化泛化性的沉淀在編譯器技術(shù)內(nèi)，避免編譯開(kāi)銷爆炸以及分層之后不同層級(jí)之間優(yōu)化的聯(lián)動(dòng)，仍然有更多的未知需要去探索和挖掘。這也成為以MLIR框架為代表的下一代深度學(xué)習(xí)編譯器需要著重思考和解決的問(wèn)題。

四 BladeDISC的主要技術(shù)特點(diǎn)

項(xiàng)目最早啟動(dòng)的初衷是為了解決XLA和TVM當(dāng)前版本的靜態(tài)shape限制，內(nèi)部命名為 DISC (DynamIc Shape Compiler)，希望打造一款能夠在實(shí)際業(yè)務(wù)中使用的完備支持動(dòng)態(tài)shape語(yǔ)義的深度學(xué)習(xí)編譯器。

從團(tuán)隊(duì)四年前啟動(dòng)深度學(xué)習(xí)編譯器方向的工作以來(lái)，動(dòng)態(tài)shape問(wèn)題一直是阻礙實(shí)際業(yè)務(wù)落地的嚴(yán)重問(wèn)題之一。彼時(shí)，包括XLA在內(nèi)的主流深度學(xué)習(xí)框架，都是基于靜態(tài)shape語(yǔ)義的編譯器框架。典型的方案是需要用戶指定輸入的shape，或是由編譯器在運(yùn)行時(shí)捕捉待編譯子圖的實(shí)際輸入shape組合，并且為每一個(gè)輸入shape組合生成一份編譯結(jié)果。

靜態(tài)shape編譯器的優(yōu)勢(shì)顯而易見(jiàn)，編譯期完全已知靜態(tài)shape信息的情況下，Compiler可以作出更好的優(yōu)化決策并得到更好的CodeGen性能，同時(shí)也能夠得到更好的顯存/內(nèi)存優(yōu)化plan和調(diào)度執(zhí)行plan。然而，其缺點(diǎn)也十分明顯，具體包括：

• 大幅增加編譯開(kāi)銷。引入離線編譯預(yù)熱過(guò)程，大幅增加推理任務(wù)部署過(guò)程復(fù)雜性;訓(xùn)練迭代速度不穩(wěn)定甚至整體訓(xùn)練時(shí)間負(fù)優(yōu)化。
• 部分業(yè)務(wù)場(chǎng)景shape變化范圍趨于無(wú)窮的，導(dǎo)致編譯緩存永遠(yuǎn)無(wú)法收斂，方案不可用。
• 內(nèi)存顯存占用的增加。編譯緩存額外占用的內(nèi)存顯存，經(jīng)常導(dǎo)致實(shí)際部署環(huán)境下的內(nèi)存/顯存OOM，直接阻礙業(yè)務(wù)的實(shí)際落地。
• 人工padding為靜態(tài)shape等緩解性方案對(duì)用戶不友好，大幅降低應(yīng)用的通用性和透明性，影響迭代效率。

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

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

五 關(guān)鍵技術(shù)

BladeDISC的整體架構(gòu)，及其在阿里云相關(guān)產(chǎn)品中的上下文關(guān)系如下圖所示：

1 MLIR基礎(chǔ)架構(gòu)

MLIR是由Google在2019年發(fā)起的項(xiàng)目，MLIR 的核心是一套靈活的多層IR基礎(chǔ)設(shè)施和編譯器實(shí)用工具庫(kù)，深受 LLVM 的影響，并重用其許多優(yōu)秀理念。這里我們選擇基于MLIR的主要原因包括其比較豐富的基礎(chǔ)設(shè)施支持，方便擴(kuò)展的模塊化設(shè)計(jì)架構(gòu)以及MLIR較強(qiáng)的膠水能力。

2 動(dòng)態(tài)shape編譯

上圖為BladeDISC的主體Pass Pipeline設(shè)計(jì)。對(duì)比目前主流的深度學(xué)習(xí)編譯器項(xiàng)目，主要技術(shù)特點(diǎn)如下：

圖層IR設(shè)計(jì)

BladeDISC選擇基于HLO作為核心圖層IR來(lái)接入不同的前端框架，但是HLO是原本為XLA設(shè)計(jì)的純靜態(tài)shape語(yǔ)義的IR。靜態(tài)場(chǎng)景下，HLO IR中的shape表達(dá)會(huì)被靜態(tài)化，所有的shape計(jì)算會(huì)被固化為編譯時(shí)常量保留在編譯結(jié)果中;而在動(dòng)態(tài)shape場(chǎng)景下，IR本身需要有足夠的能力表達(dá)shape計(jì)算和動(dòng)態(tài)shape信息的傳遞。BladeDISC從項(xiàng)目建立開(kāi)始一直與MHLO社區(qū)保持緊密的合作，在XLA的HLO IR基礎(chǔ)上，擴(kuò)展了一套具有完備動(dòng)態(tài)shape表達(dá)能力的IR，并增加了相應(yīng)的基礎(chǔ)設(shè)施以及前端框架的算子轉(zhuǎn)換邏輯。這部分實(shí)現(xiàn)目前已經(jīng)完整upstream至MHLO社區(qū)，確保后續(xù)其它MHLO相關(guān)項(xiàng)目中IR的一致性。

運(yùn)行時(shí)Shape計(jì)算、存儲(chǔ)管理和Kernel調(diào)度

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

動(dòng)態(tài)shape下的性能問(wèn)題

在shape未知或者部分未知的情況下，深度學(xué)習(xí)編譯器在性能上面臨的挑戰(zhàn)被進(jìn)一步放大。在大多數(shù)主流硬件backend上，BladeDISC采用區(qū)分計(jì)算密集型部分和訪存密集型部分的策略，以期在性能與復(fù)雜性和編譯開(kāi)銷之間獲取更好的平衡。

對(duì)于計(jì)算密集型部分，不同的shape要求更加精細(xì)的schedule實(shí)現(xiàn)來(lái)獲得更好的性能，pass pipeline在設(shè)計(jì)上的主要考慮是需要支持在運(yùn)行時(shí)根據(jù)不同的具體shape選擇合適的算子庫(kù)實(shí)現(xiàn)，以及處理動(dòng)態(tài)shape語(yǔ)義下的layout問(wèn)題。

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

大顆粒度算子融合

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

在動(dòng)態(tài)shape語(yǔ)義下實(shí)現(xiàn)fusion-stitching對(duì)比靜態(tài)shape語(yǔ)義下同樣需要處理更大的復(fù)雜性，動(dòng)態(tài)shape語(yǔ)義下的shape constraint抽象一定程度上簡(jiǎn)化了這一復(fù)雜性，使整體性能進(jìn)一步接近甚至超過(guò)手工算子實(shí)現(xiàn)。

3 多前端框架支持

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

4 運(yùn)行時(shí)環(huán)境適配

為將編譯的結(jié)果能夠配合TensorFlow/PyTorch等宿主在各自的運(yùn)行環(huán)境中執(zhí)行起來(lái)，以及管理運(yùn)行時(shí)IR層不易表達(dá)的狀態(tài)信息等等，我們?yōu)椴煌倪\(yùn)行時(shí)環(huán)境實(shí)現(xiàn)了一套統(tǒng)一的Compiler架構(gòu)，并引入了運(yùn)行時(shí)抽象層，即RAL(Runtime Abstraction Layer)層。

RAL實(shí)現(xiàn)了多種運(yùn)行環(huán)境的適配支持，用戶可以根據(jù)需要進(jìn)行選擇，具體包括：

• 全圖編譯，獨(dú)立運(yùn)行。當(dāng)整個(gè)計(jì)算圖都支持編譯時(shí)，RAL提供了一套簡(jiǎn)易的runtime以及在此之上RAL Driver的實(shí)現(xiàn)，使得compiler編譯出來(lái)結(jié)果可以脫離框架直接運(yùn)行，減少框架overhead。
• TF中子圖編譯運(yùn)行。
• Pytorch中子圖編譯運(yùn)行。

以上環(huán)境中在諸如資源管理，API語(yǔ)義等上存在差異，RAL通過(guò)抽象出一套最小集合的API ，并清晰的定義出它們的語(yǔ)義，將編譯器與運(yùn)行時(shí)隔離開(kāi)來(lái)，來(lái)達(dá)到在不同的環(huán)境中都能夠執(zhí)行編譯出來(lái)的結(jié)果的目的。此外RAL層實(shí)現(xiàn)了無(wú)狀態(tài)編譯，解決了計(jì)算圖的編譯之后，編譯的結(jié)果可能被多次執(zhí)行時(shí)的狀態(tài)信息處理問(wèn)題。一方面簡(jiǎn)化了代碼生成的復(fù)雜度，另一方面也更容易支持多線程并發(fā)執(zhí)行(比如推理)的場(chǎng)景，同時(shí)在錯(cuò)誤處理，回滾方面也更加容易支持。

六 應(yīng)用場(chǎng)景

BladeDISC的典型應(yīng)用場(chǎng)景可以不太嚴(yán)格的分為兩類：其一是在主流的硬件平臺(tái)上(包括Nvidia GPU，x86 CPU等)上作為通用、透明的性能優(yōu)化工具，降低用戶部署AI作業(yè)的人力負(fù)擔(dān)，提高模型迭代效率;另一個(gè)重要的應(yīng)用場(chǎng)景是幫助新硬件做AI場(chǎng)景的適配和接入支持。

目前BladeDISC已經(jīng)廣泛應(yīng)用在阿里內(nèi)部和阿里云上外部用戶的多個(gè)不同應(yīng)用場(chǎng)景下，覆蓋模型類型涉及NLP、機(jī)器翻譯、語(yǔ)音類ASR、語(yǔ)音TTS、圖像檢測(cè)、識(shí)別、AI for science等等多種典型AI應(yīng)用;覆蓋行業(yè)包括互聯(lián)網(wǎng)、電商、自動(dòng)駕駛、安全行業(yè)、在線娛樂(lè)、醫(yī)療和生物等等。

在推理場(chǎng)景下，BladeDISC與TensorRT等廠商提供的推理優(yōu)化工具形成良好的技術(shù)互補(bǔ)，其主要差異性優(yōu)勢(shì)包括：

• 應(yīng)對(duì)動(dòng)態(tài)shape業(yè)務(wù)完備的動(dòng)態(tài)shape語(yǔ)義支持
• 基于compiler based的技術(shù)路徑的模型泛化性在非標(biāo)準(zhǔn)模型上的性能優(yōu)勢(shì)
• 更為靈活的部署模式選擇，以插件形式支持前端框架的透明性優(yōu)勢(shì)

下圖為Nvidia T4硬件上幾個(gè)真實(shí)的業(yè)務(wù)例子的性能收益數(shù)字：

在新硬件支持方面，目前普遍的情況是除了積累比較深厚的Nvidia等頭部廠商之外，包括ROCM等其它GPGPU硬件普遍存在的情況是硬件的指標(biāo)已經(jīng)具備相當(dāng)?shù)母?jìng)爭(zhēng)力，但廠商受制于AI軟件棧上的積累相對(duì)較少，普遍存在硬件算力無(wú)法發(fā)揮出來(lái)導(dǎo)致硬件落地應(yīng)用困難的問(wèn)題。如前文所述，基于編譯器的技術(shù)路徑下天然對(duì)于硬件的后端具備一定的泛化能力，且與硬件廠商的技術(shù)儲(chǔ)備形成比較強(qiáng)的互補(bǔ)。BladeDISC目前在GPGPU和通用CPU體系結(jié)構(gòu)上的儲(chǔ)備相對(duì)比較成熟。以GPGPU為例，在Nvidia GPU上的絕大部分技術(shù)?？梢赃w移至海光DCU和AMD GPU等體系結(jié)構(gòu)相近的硬件上。BladeDISC較強(qiáng)的硬件泛化能力配合硬件本身較強(qiáng)的通用性，很好的解決了新硬件適配的性能和可用性問(wèn)題。

下圖為海光DCU上幾個(gè)真實(shí)業(yè)務(wù)例子上的性能數(shù)字：

七 開(kāi)源生態(tài)——構(gòu)想和未來(lái)

我們決定建設(shè)開(kāi)源生態(tài)主要有如下的考慮：

• BladeDISC發(fā)源于阿里云計(jì)算平臺(tái)團(tuán)隊(duì)的業(yè)務(wù)需求，在開(kāi)發(fā)過(guò)程中與MLIR/MHLO/IREE等社區(qū)同行之間的討論和交流給了我們很好的輸入和借鑒。在我們自身隨著業(yè)務(wù)需求的迭代逐漸完善的同時(shí)，也希望能夠開(kāi)源給社區(qū)，在目前整個(gè)AI編譯器領(lǐng)域?qū)嶒?yàn)性項(xiàng)目居多，偏實(shí)用性強(qiáng)的產(chǎn)品偏少，且不同技術(shù)棧之間的工作相對(duì)碎片化的情況下，希望能夠?qū)⒆陨淼慕?jīng)驗(yàn)和理解也同樣回饋給社區(qū)，希望和深度學(xué)習(xí)編譯器的開(kāi)發(fā)者和AI System的從業(yè)者之間有更多更好的交流和共建，為這個(gè)行業(yè)貢獻(xiàn)我們的技術(shù)力量。
• 我們希望能夠借助開(kāi)源的工作，收到更多真實(shí)業(yè)務(wù)場(chǎng)景下的用戶反饋，以幫助我們持續(xù)完善和迭代，并為后續(xù)的工作投入方向提供輸入。

后續(xù)我們計(jì)劃以兩個(gè)月為單位定期發(fā)布Release版本。BladeDISC近期的Roadmap如下：

• 持續(xù)的魯棒性及性能改進(jìn)
• x86后端補(bǔ)齊計(jì)算密集型算子的支持，端到端完整開(kāi)源x86后端的支持
• GPGPU上基于Stitching的大顆粒度自動(dòng)代碼生成
• AMD rocm GPU后端的支持
• PyTorch訓(xùn)練場(chǎng)景的支持

此外，在中長(zhǎng)期，我們?cè)谙旅鎺讉€(gè)探索性的方向上會(huì)持續(xù)投入精力，也歡迎各種維度的反饋和改進(jìn)建議以及技術(shù)討論，同時(shí)我們十分歡迎和期待對(duì)開(kāi)源社區(qū)建設(shè)感興趣的同行一起參與共建。

• 更多新硬件體系結(jié)構(gòu)的支持和適配，以及新硬件體系結(jié)構(gòu)下軟硬件協(xié)同方法學(xué)的沉淀
• 計(jì)算密集型算子自動(dòng)代碼生成和動(dòng)態(tài)shape語(yǔ)義下全局layout優(yōu)化的探索
• 稀疏子圖的優(yōu)化探索
• 動(dòng)態(tài)shape語(yǔ)義下運(yùn)行時(shí)調(diào)度策略、內(nèi)存/顯存優(yōu)化等方面的探索
• 模型壓縮與編譯優(yōu)化聯(lián)合的技術(shù)探索
• 圖神經(jīng)網(wǎng)絡(luò)等更多AI作業(yè)類型的支持和優(yōu)化等

開(kāi)源地址：

https://github.com/alibaba/BladeDISC

參考文獻(xiàn)

1. "DISC: A Dynamic Shape Compiler for Machine Learning Workloads", Kai Zhu, Wenyi Zhao, Zhen Zheng, Tianyou Guo, Pengzhan Zhao, Feiwen Zhu, Junjie Bai, Jun Yang, Xiaoyong Liu, Lansong Diao, Wei Lin
2. Presentations on MLIR Developers' Weekly Conference: 1, 2
3. "AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures", Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Jidong Zhai, Shuaiwen Leon Song, and Wei Lin. The 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2022. [to appear]
4. "FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads", Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang, and Wei Lin. arXiv preprint