全面解析AI編譯器面臨的機(jī)遇和挑戰(zhàn)

本文分為兩個(gè)部分，第一部分為綜述（領(lǐng)域編譯器發(fā)展的前世今生 ? 綜述）；這部分重點(diǎn)討論面向AI領(lǐng)域的編譯技術(shù)。

隨著人工智能時(shí)代的來臨，AI領(lǐng)域應(yīng)用的大量出現(xiàn)也促進(jìn)著領(lǐng)域編譯的發(fā)展，最突出的表現(xiàn)就是多種AI編譯器的普及和應(yīng)用。AI領(lǐng)域有幾個(gè)重要的特征使得AI編譯器面臨很多新的機(jī)遇和挑戰(zhàn)：一是AI領(lǐng)域中編程框架對(duì)計(jì)算圖與算子分離的設(shè)計(jì)機(jī)制為編譯優(yōu)化提供了更多的機(jī)會(huì)和更廣闊的空間；二是AI領(lǐng)域中對(duì)張量的抽象為編譯優(yōu)化提供了具有鮮明領(lǐng)域特征的語(yǔ)義信息；三是以Python為主的動(dòng)態(tài)解釋器語(yǔ)言前端為其與AI編譯器的銜接帶了挑戰(zhàn)；四是面向AI的領(lǐng)域?qū)Ｓ眉軜?gòu)為應(yīng)用的可移植性帶來了挑戰(zhàn)。在這些因素的驅(qū)動(dòng)下，近年來學(xué)術(shù)界和工業(yè)界在AI編譯方面提出了一系列創(chuàng)新性的方法，也為編譯這一基礎(chǔ)學(xué)科的發(fā)展注入了新的活力。

01?圖算分離與圖算融合的優(yōu)化探索

為了讓開發(fā)者使用方便，框架前端(圖層)會(huì)盡量對(duì)Tensor計(jì)算進(jìn)行抽象封裝，開發(fā)者只要關(guān)注邏輯意義上的模型和算子；而在后端算子層的性能優(yōu)化時(shí)，又可以打破算子的邊界，從更細(xì)粒度的循環(huán)調(diào)度等維度，結(jié)合不同的硬件特點(diǎn)完成優(yōu)化。這種圖算分離的解耦設(shè)計(jì)大大簡(jiǎn)化了AI復(fù)雜系統(tǒng)的設(shè)計(jì)，因此，多層IR設(shè)計(jì)無疑是較好的選擇，目前的主流IR設(shè)計(jì)也是分為圖（TVM Relay，XLA HLO，MindSpore MindIR等）和算子（TVM tir，XLA LLO，MindSpore AKG等）兩層。以主流AI編譯器TVM[13]和TensorFlow XLA[14]為例。TVM 和 XLA 上層都采用了數(shù)據(jù)流圖的中間表示，用圖結(jié)點(diǎn)來表示計(jì)算，邊表示數(shù)據(jù)流的依賴。在下層TVM 和 XLA 都針對(duì)編譯器自動(dòng)生成不同平臺(tái)的高效代碼進(jìn)行了設(shè)計(jì)。其中TVM底層針對(duì)深度學(xué)習(xí)核心的張量處理設(shè)計(jì)的中間表示tir，它借鑒了Halide的中間表示來描述結(jié)點(diǎn)內(nèi)的計(jì)算，可以針對(duì)不同目標(biāo)平臺(tái)定制調(diào)度策略，從而實(shí)現(xiàn)平臺(tái)相關(guān)的深度優(yōu)化。TensorFlow XLA 則提出了一種基于代數(shù)表示的中間表示（XLA HLO），高層的數(shù)據(jù)流圖被轉(zhuǎn)換為XLA HLO 的中間表示，在此中間表示上可以實(shí)施支持jit的算子融合、內(nèi)存操作消除等優(yōu)化，優(yōu)化后的XLA HLO 可以被翻譯為L(zhǎng)LVM 中間表示或直接映射到TPU 平臺(tái)。TVM和TensorFlow XLA 的圖與張量（或代數(shù)）中間表示相結(jié)合的方法，一方面能夠適配人工智能領(lǐng)域用數(shù)據(jù)流圖來描述應(yīng)用的需求，另一方面又能夠兼顧應(yīng)用在不同硬件平臺(tái)之間的移植和優(yōu)化。

但圖層和算子層獨(dú)立優(yōu)化無法充分發(fā)揮芯片性能。近來面向圖算融合的優(yōu)化也日益成為學(xué)術(shù)界重要的研究方向。EasyView[15]提出了針對(duì)在網(wǎng)絡(luò)實(shí)現(xiàn)中高頻出現(xiàn)的tensor view類算子的端到端在線編譯自動(dòng)融合方法，包含view lowering，內(nèi)存活動(dòng)追蹤，讀寫關(guān)系一致的算子拓?fù)湫蛄蝎@取，以及計(jì)算內(nèi)存優(yōu)化策略等內(nèi)容。Apollo[16]設(shè)計(jì)了一個(gè)開放式多層規(guī)約式融合架構(gòu)以實(shí)現(xiàn)不同算子融合方式的協(xié)同組合。將不同的融合方式實(shí)現(xiàn)為不同的Layer，在各級(jí)layer分別做基于polyhedral優(yōu)化的循環(huán)融合，通過計(jì)算圖算子級(jí)別依賴和元素級(jí)別依賴的分析對(duì)訪存密集型算子盡可能融合，識(shí)別無計(jì)算依賴的算子并行化等優(yōu)化，然后通過對(duì)不同Layer進(jìn)行逐層規(guī)約合并，從而得到最終的融合算子子圖，并獲得最佳的融合性能收益。

Apollo架構(gòu)：子圖切分，融合，優(yōu)化[16]

02?面向張量的極致編譯優(yōu)化

AI編譯器的核心抽象是張量（矩陣的高維推廣）。在AI領(lǐng)域，各種數(shù)據(jù)，如圖片，文字，視頻等，都被抽象成張量，而原本對(duì)這些數(shù)據(jù)的處理也被抽象成對(duì)張量的計(jì)算操作，如卷積，轉(zhuǎn)置，池化，歸一化等，這些對(duì)張量的操作按照順序組合就組成了張量計(jì)算的數(shù)據(jù)流圖。做這層抽象的意義在筆者看來是因?yàn)閭鹘y(tǒng)編譯器的IR，如LLVM IR太底層了，其對(duì)數(shù)據(jù)的處理粒度在標(biāo)量至多向量級(jí)別，而編譯器針對(duì)如此底層的IR分析能力有限。畢竟通過嵌套在多層循環(huán)里的load，store和alu操作已經(jīng)很難還原出其原本的計(jì)算信息（如矩陣轉(zhuǎn)置等）。從而像是矩陣轉(zhuǎn)置再轉(zhuǎn)置這樣的pattern就基本無法在LLVM IR的粒度分析出來，而通過張量和張量計(jì)算的抽象層，這種優(yōu)化就可以很容易實(shí)現(xiàn)。

AI編譯器的優(yōu)化目標(biāo)主要是為了提升AI模型的端到端性能，這個(gè)性能會(huì)受到包括計(jì)算訪存比，并行性，資源占用率等多方面因素影響，因而很難通過一個(gè)通用的策略涵蓋大量不同的后端而都能達(dá)到非常優(yōu)秀的性能。AI編譯器通常通過搜索調(diào)度空間的方式，來尋找適配后端的極致張量?jī)?yōu)化策略。這里以TVM為例。TVM將計(jì)算和調(diào)度分離，計(jì)算通過張量表達(dá)式表示，張量表達(dá)式在設(shè)計(jì)上借鑒了 Halide、Darkroom 和 TACO。調(diào)度則是針對(duì)計(jì)算的一系列變換，為了在許多后端實(shí)現(xiàn)高性能，必須要支持足夠多的調(diào)度原語(yǔ)來涵蓋不同硬件后端的各種優(yōu)化而包括tile/fuse/reorder/bind/compute_at等等，通過調(diào)度可以挖掘張量計(jì)算在特定硬件后端下的極致性能。TVM陸續(xù)發(fā)展了從基于模版的AutoTVM到基于搜索的Ansor，再到通過DSL tensorIR結(jié)合二者優(yōu)點(diǎn)的Meta Schedule逐漸發(fā)展起來的自動(dòng)調(diào)度搜索策略，可以針對(duì)不同的目標(biāo)硬件平臺(tái)來自動(dòng)搜索的更優(yōu)卸載方式，從而實(shí)現(xiàn)平臺(tái)相關(guān)的深度優(yōu)化。

TVM 張量表達(dá)式優(yōu)化流程，多后端調(diào)度支持[17]

03?Python為主的動(dòng)態(tài)解釋器語(yǔ)言前端銜接

與傳統(tǒng)編譯器不同，AI編譯器通常不需要Lexer/Parser，而是基于前端語(yǔ)言（主要是Python）的AST將模型解析并構(gòu)造為計(jì)算圖IR，側(cè)重于保留shape、layout等Tensor計(jì)算特征信息，當(dāng)然部分編譯器還能保留控制流的信息。這里的難點(diǎn)在于，Python是一種靈活度極高的解釋執(zhí)行的語(yǔ)言，AI編譯器需要把它轉(zhuǎn)到靜態(tài)的IR上。針對(duì)這一難點(diǎn)，目前一個(gè)比較突出的工作是pytorch2.0[18]提出的TorchDynamo[19]這一 JIT 編譯接口，傳統(tǒng)pytorch編程無論是trace還是eager模式都沒有辦法簡(jiǎn)單地通過python代碼獲取模型的圖結(jié)構(gòu)，導(dǎo)致模型導(dǎo)出、算子融合優(yōu)化、模型量化等工作都出現(xiàn)困難。TorchDynamo支持在運(yùn)行時(shí)修改python動(dòng)態(tài)執(zhí)行邏輯，修改的時(shí)機(jī)是 在CPython 解釋器的 ByteCode 執(zhí)行前，從而可以使通過在運(yùn)行時(shí)設(shè)置一個(gè)自定義的 Frame，該 Frame 中的 ByteCode 支持 CallBack 到 Python 層去修改。供用戶自定義計(jì)算圖。利用這一機(jī)制TorchDynamo支持了動(dòng)態(tài)圖的特性，即只需要通過python的執(zhí)行機(jī)制即可自動(dòng)調(diào)用后方對(duì)應(yīng)的靜態(tài)子模型。以下圖為例，TorchDynamo方式寫的bar模塊具有一個(gè)動(dòng)態(tài)信息的條件分支（對(duì)b.sum()<0的判斷），這是傳統(tǒng)trace和eager模式都無法支持的描述，但通過TorchDynamo方式，python執(zhí)行到條件分支時(shí)可以根據(jù)動(dòng)態(tài)信息自動(dòng)根據(jù)條件調(diào)用有if語(yǔ)句的子圖或沒有if語(yǔ)句的子圖，從而完成了對(duì)python if語(yǔ)句描述的動(dòng)態(tài)圖信息的支持。

通過TorchDynamo支持動(dòng)態(tài)的控制依賴計(jì)算圖[19]

04?DSA芯片架構(gòu)的支持

隨著AI應(yīng)用的算力需求日益增長(zhǎng)，以CPU為主的通用計(jì)算算力越來越難以滿足大規(guī)模AI模型的性能和時(shí)延需求，因此AI應(yīng)用往往需要通過高性能的針對(duì)AI的DSA架構(gòu)的加速器后端來加速算子性能。這不僅包括對(duì)基于FPGA、ASIC、類腦計(jì)算等專用加速器的，也有在GPGPU上擴(kuò)展的tensor core等專用加速核心。這些加速器在計(jì)算通用性，可配置性，功耗，性能，性價(jià)比，易用性等方面各有千秋，同時(shí)在實(shí)現(xiàn)架構(gòu)上差異巨大，根據(jù)這些架構(gòu)約定的軟硬件接口層次的不同，如是否將內(nèi)部cache操作暴露給編譯器，是否硬件自動(dòng)內(nèi)存管理等，編譯器能做的優(yōu)化層次也不同。這也導(dǎo)致AI應(yīng)用部署在不同的后端架構(gòu)時(shí)需要生成出各種粒度的最終計(jì)算負(fù)載，這種粒度對(duì)應(yīng)硬件/runtime規(guī)定的接口。如在GPU后端下，一個(gè)或多個(gè)圖上的算子，在經(jīng)過算子融合等優(yōu)化步驟后，最終匹配上了cudnn/miopen等GPU算子庫(kù)的一個(gè)庫(kù)函數(shù)實(shí)現(xiàn)，最終編譯器通過codegen生成以cuMalloc，cuMemcpy，cuLaunchKernel等cuda driver級(jí)別的api調(diào)用，借助CUDA runtime軟件棧分別完成GPU內(nèi)存空間管理，CPU到GPU的數(shù)據(jù)傳輸，算子的底層庫(kù)函數(shù)調(diào)用，GPU數(shù)據(jù)回傳等操作，最終完成了計(jì)算步驟。由于DSA架構(gòu)的差異性和多樣性，每一種架構(gòu)都可能對(duì)應(yīng)一種codegen方式和runtime支持。因而如何設(shè)計(jì)擴(kuò)展DSA架構(gòu)以實(shí)現(xiàn)對(duì)新的DSA架構(gòu)的快速支持也是AI編譯器的難點(diǎn)。

不同類型AI芯片后端體系結(jié)構(gòu)的比較[20]

05?面向神經(jīng)網(wǎng)絡(luò)的全局優(yōu)化

AI應(yīng)用的優(yōu)化主要是為了提升端到端的模型性能，其受到包括計(jì)算訪存比，并行性，資源占用率等多方面因素影響。在通用的編譯器優(yōu)化（如LLVM的優(yōu)化pass）以外，AI編譯器引入的面向神經(jīng)網(wǎng)絡(luò)應(yīng)用的特定優(yōu)化主要包括三類：張量?jī)?yōu)化，自動(dòng)微分，自動(dòng)并行。在AI領(lǐng)域，計(jì)算被抽象成張量的計(jì)算，基于張量的計(jì)算原語(yǔ)和計(jì)算優(yōu)化是AI編譯器的重要關(guān)注點(diǎn)。自動(dòng)微分是支持AI網(wǎng)絡(luò)訓(xùn)練的重要基礎(chǔ)，因?yàn)槟Ｐ陀?xùn)練的基礎(chǔ)就是梯度下降和誤差反向傳播，都需要自動(dòng)微分技術(shù)的支持。目前主要有基于計(jì)算圖的自動(dòng)微分、基于Tape和運(yùn)算符重載的自動(dòng)微分、基于source2source的自動(dòng)微分等主流方案。自動(dòng)并行技術(shù)則主要包括：數(shù)據(jù)并行、算子級(jí)模型并行、Pipeline模型并行、優(yōu)化器模型并行和重計(jì)算等，以提高整體并行度的方式優(yōu)化網(wǎng)絡(luò)性能。 ?

06?領(lǐng)域定制架構(gòu)下的編譯基礎(chǔ)設(shè)施

隨著傳統(tǒng)通用處理器的性能提升越來越困難，近年來領(lǐng)域定制硬件（DSA）成為體系結(jié)構(gòu)設(shè)計(jì)中新的增長(zhǎng)點(diǎn)，也是計(jì)算機(jī)體系結(jié)構(gòu)黃金時(shí)代重要的發(fā)展。多種類型的加速器不斷涌現(xiàn)。例如GPU、NPU、FGPA、ASIC 等。而為了適配大量新生的DSA，領(lǐng)域編譯開發(fā)需要通過共用來降低開發(fā)成本，并重點(diǎn)關(guān)注于核心商業(yè)邏輯的實(shí)現(xiàn)，這使得領(lǐng)域編譯也像電力，網(wǎng)絡(luò)，公共云等技術(shù)一樣，走向了技術(shù)演進(jìn)的自然終點(diǎn)：基礎(chǔ)設(shè)施化。其中最有代表性的基礎(chǔ)設(shè)施是MLIR[21，24]。

MLIR是由Google 提出的一個(gè)能夠快速構(gòu)建領(lǐng)域編譯器的基礎(chǔ)設(shè)施，提出了一種構(gòu)建可重用、可擴(kuò)展編譯器基礎(chǔ)結(jié)構(gòu)的新方法。其核心思想是利用多層次中間表示來解決軟件的碎片化問題，減少構(gòu)建Domain Specific Compiler的開銷。MLIR雖然目前主要用于機(jī)器學(xué)習(xí)領(lǐng)域，但在設(shè)計(jì)上是通用的編譯器框架，比如也有FLANG（llvm中的FORTRAN編譯器）[22]，CIRCT（用于硬件設(shè)計(jì)）[23]等與ML無關(guān)的項(xiàng)目。MLIR 提供一系列可復(fù)用的易擴(kuò)展的基礎(chǔ)組件，從而使得不同領(lǐng)域的編譯開發(fā)人員能夠快速的搭建領(lǐng)域?qū)Ｓ镁幾g器，而且不同領(lǐng)域的編譯分析及優(yōu)化可以被復(fù)用。

與 LLVM IR 唯一的中間表示不同，MLIR能通過多層方言（dialect）的設(shè)計(jì)，表示更高層次的結(jié)構(gòu)和操作，比如神經(jīng)網(wǎng)絡(luò)的圖結(jié)構(gòu)，張量的計(jì)算等。MLIR將一些領(lǐng)域的特性抽象為方言并允許用戶自定義新的方言，與此同時(shí)，MLIR 提供了不同方言間的轉(zhuǎn)換機(jī)制來實(shí)現(xiàn)不同方言上編譯分析和優(yōu)化的復(fù)用。MLIR執(zhí)行過程和LLVM一樣，IR會(huì)通過由Pass組成的優(yōu)化Pipeline，不斷地方言內(nèi)，方言間變換直到生成最終的IR，然后被lower到底層的通用IR上進(jìn)行代碼生成。MLIR不僅僅是一個(gè)中間表示，而是一個(gè)新的編譯器基礎(chǔ)設(shè)施。近年來，學(xué)術(shù)界和工業(yè)界也在MLIR 上開展了很多領(lǐng)域編譯優(yōu)化的工作。如下圖所示的將TensorFlow XLA 接入到了MLIR 的例子。上層的模型輸入為TF Graph，在MLIR架構(gòu)下逐層變換到HLO，LHLO，Affine，Vector等更低層次的方言上，在每級(jí)方言上都有對(duì)應(yīng)層次和粒度的優(yōu)化和調(diào)度，如在最高層的HLO適合做融合等。最終生成LLVM IR或SPIR-V等通用中間表示，再調(diào)用后端通用編譯器完成最終代碼生成。

MLIR在Tensorflow XLA上的實(shí)現(xiàn)[25]

07?國(guó)內(nèi)學(xué)界和業(yè)界的工作

AI生態(tài)的火熱也促使著學(xué)術(shù)界和工業(yè)界在AI編譯系統(tǒng)和AI異構(gòu)加速器等方面全面發(fā)展。國(guó)內(nèi)在AI生態(tài)的建設(shè)方面開始較早，投入眾多，催生著大量的學(xué)術(shù)成果和工業(yè)界基礎(chǔ)設(shè)施的發(fā)展。

一方面，國(guó)內(nèi)學(xué)術(shù)界在AI編譯系統(tǒng)和AI加速器體系結(jié)構(gòu)等方面有很多研究突破，包括但不限于語(yǔ)言，編譯，軟硬件系統(tǒng)設(shè)計(jì)等方面。例如：針對(duì)目前的張量?jī)?yōu)化只考慮了完全等價(jià)變換，為張量程序優(yōu)化引入了部分等價(jià)變換的優(yōu)化和對(duì)應(yīng)的糾正機(jī)制的PET[26]；通過語(yǔ)言層面引入對(duì)tensor的細(xì)粒度控制，包括tensor的不規(guī)則索引等，避免了大量冗余計(jì)算的FreeTensor[27]；通過特定于GNN網(wǎng)絡(luò)的圖運(yùn)算符抽象做優(yōu)化的uGrapher[28]；將硬件抽象設(shè)計(jì)為IR以支持更多intrinsic 原語(yǔ)的AMOS[29]；通過基于殘差的精度細(xì)化技術(shù)，控制量化誤差，在性能和精度之間進(jìn)行權(quán)衡的QUANTENSOR[30]；根據(jù)優(yōu)化目標(biāo)配置實(shí)現(xiàn)云-移動(dòng)部署的性能功耗綜合優(yōu)化的DNNTune[31]；面向任意精度計(jì)算（Arbitrary Precision Computing：APC）的Cambricon-P[32]體系結(jié)構(gòu)，等等。

另一方面，工業(yè)界結(jié)合各廠家自身業(yè)務(wù)的需求，在AI基礎(chǔ)設(shè)施和系統(tǒng)技術(shù)上不斷優(yōu)化和探索，在不同的維度持續(xù)發(fā)力，也貢獻(xiàn)了很多開源項(xiàng)目。在AI編譯基礎(chǔ)設(shè)施方面，華為的MindSpore社區(qū)[33]提供了一個(gè)主要面向華為的昇騰處理器后端的云邊端全場(chǎng)景開放AI推理和訓(xùn)練框架。其在包括圖算聯(lián)合優(yōu)化，大規(guī)模分布式自動(dòng)并行，自動(dòng)算子生成等多項(xiàng)技術(shù)上做出了探索和貢獻(xiàn)。阿里巴巴的PAI團(tuán)隊(duì)也專注于編譯優(yōu)化，探索了XLA，TVM，MLIR等多條技術(shù)路線，目前在大顆粒算子融合技術(shù)，以及GPU上訪存密集型算子的融合優(yōu)化上也取得不錯(cuò)的效果，并在MLIR這條技術(shù)路線上擴(kuò)充了框架應(yīng)對(duì)動(dòng)態(tài)輸入shape上的能力[34]。此外，國(guó)內(nèi)的大型互聯(lián)網(wǎng)公司，AI技術(shù)和芯片公司等都在面向AI的編譯技術(shù)上有越來越多的投入，極大推進(jìn)了相關(guān)技術(shù)的發(fā)展。

領(lǐng)域編譯器作為通用編譯器的重要補(bǔ)充，在發(fā)揮極致性能和提升開發(fā)效率方面一直發(fā)揮著重要的作用，近年來在AI領(lǐng)域更是備受關(guān)注。隨著AI技術(shù)的快速發(fā)展，DSA硬件在AI計(jì)算中的大量使用，AI軟件棧也相應(yīng)的日趨復(fù)雜，對(duì)編譯技術(shù)提出更高的要求，這也大大促進(jìn)了編譯學(xué)科的快速發(fā)展。我們相信，在強(qiáng)烈的需求驅(qū)動(dòng)下，通過學(xué)界和業(yè)界的共同努力，領(lǐng)域編譯技術(shù)在各種類型的計(jì)算系統(tǒng)中將扮演越來越重要的角色。

編輯：黃飛

?