本文分為兩個部分，第一部分為綜述（領域編譯器發展的前世今生 ? 綜述）；這部分重點討論面向AI領域的編譯技術。

隨著人工智能時代的來臨，AI領域應用的大量出現也促進著領域編譯的發展，最突出的表現就是多種AI編譯器的普及和應用。AI領域有幾個重要的特征使得AI編譯器面臨很多新的機遇和挑戰：一是AI領域中編程框架對計算圖與算子分離的設計機制為編譯優化提供了更多的機會和更廣闊的空間；二是AI領域中對張量的抽象為編譯優化提供了具有鮮明領域特征的語義信息；三是以Python為主的動態解釋器語言前端為其與AI編譯器的銜接帶了挑戰；四是面向AI的領域專用架構為應用的可移植性帶來了挑戰。在這些因素的驅動下，近年來學術界和工業界在AI編譯方面提出了一系列創新性的方法，也為編譯這一基礎學科的發展注入了新的活力。

01?圖算分離與圖算融合的優化探索

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

但圖層和算子層獨立優化無法充分發揮芯片性能。近來面向圖算融合的優化也日益成為學術界重要的研究方向。EasyView[15]提出了針對在網絡實現中高頻出現的tensor view類算子的端到端在線編譯自動融合方法，包含view lowering，內存活動追蹤，讀寫關系一致的算子拓撲序列獲取，以及計算內存優化策略等內容。Apollo[16]設計了一個開放式多層規約式融合架構以實現不同算子融合方式的協同組合。將不同的融合方式實現為不同的Layer，在各級layer分別做基于polyhedral優化的循環融合，通過計算圖算子級別依賴和元素級別依賴的分析對訪存密集型算子盡可能融合，識別無計算依賴的算子并行化等優化，然后通過對不同Layer進行逐層規約合并，從而得到最終的融合算子子圖，并獲得最佳的融合性能收益。

Apollo架構：子圖切分，融合，優化[16]

02?面向張量的極致編譯優化

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

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

TVM 張量表達式優化流程，多后端調度支持[17]

03?Python為主的動態解釋器語言前端銜接

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

通過TorchDynamo支持動態的控制依賴計算圖[19]

04?DSA芯片架構的支持

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

不同類型AI芯片后端體系結構的比較[20]

05?面向神經網絡的全局優化

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

06?領域定制架構下的編譯基礎設施

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

MLIR是由Google 提出的一個能夠快速構建領域編譯器的基礎設施，提出了一種構建可重用、可擴展編譯器基礎結構的新方法。其核心思想是利用多層次中間表示來解決軟件的碎片化問題，減少構建Domain Specific Compiler的開銷。MLIR雖然目前主要用于機器學習領域，但在設計上是通用的編譯器框架，比如也有FLANG（llvm中的FORTRAN編譯器）[22]，CIRCT（用于硬件設計）[23]等與ML無關的項目。MLIR 提供一系列可復用的易擴展的基礎組件，從而使得不同領域的編譯開發人員能夠快速的搭建領域專用編譯器，而且不同領域的編譯分析及優化可以被復用。

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

MLIR在Tensorflow XLA上的實現[25]

07?國內學界和業界的工作

AI生態的火熱也促使著學術界和工業界在AI編譯系統和AI異構加速器等方面全面發展。國內在AI生態的建設方面開始較早，投入眾多，催生著大量的學術成果和工業界基礎設施的發展。

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

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

領域編譯器作為通用編譯器的重要補充，在發揮極致性能和提升開發效率方面一直發揮著重要的作用，近年來在AI領域更是備受關注。隨著AI技術的快速發展，DSA硬件在AI計算中的大量使用，AI軟件棧也相應的日趨復雜，對編譯技術提出更高的要求，這也大大促進了編譯學科的快速發展。我們相信，在強烈的需求驅動下，通過學界和業界的共同努力，領域編譯技術在各種類型的計算系統中將扮演越來越重要的角色。

編輯：黃飛

?