AI與傳統編譯器

AI與傳統編譯器

至於TVM，現在有很多框架（TF，Pytorch），然后會部署到不同平台（CPU、GPU、TPU），神經網絡編譯器呢就是把不同框架里寫的東西編譯成一樣的格式再生成到某一平台的代碼

再來看傳統編譯器（更偏向於LLVM），現在有許多語言（C、ObjC、C++），也有許多平台（x86、arm），編譯器做的就是把不同語言編譯到同樣的中間代碼再生成某一平台的代碼

這兩個就是把前端的表示進行統一再生成硬件相關的程序，只不過一個前端表示的是神經網絡，一個是大家都熟悉的代碼，結構類似但實際內部工作大相徑庭

傳統編譯器：輸入高級語言輸出低級語言

神經網絡編譯器：輸入計算圖/算子，輸出低級語言

相同點是都做了類似語言轉換的工作

不同點

傳統編譯器解決的主要問題是降低編程難度，其次是優化程序性能

神經網絡編譯器解決的主要問題是優化程序性能，其次是降低編程難度

問題

1. 對於神經網絡編譯器，如果沒有體系結構相關信息的輸入是否能生成高效代碼（涉及完全自動化的問題）
2. 當前的技術投入比來看，神經網絡編譯器和人工算子實現的哪個性價比更高

神經網絡編譯器，編譯時考慮神經網絡有關的特性來優化程序。

程序優化，解釋器vs編譯器，JVM，JIT，llvm, Halide，TensorFlow, XLA, ONNX, TVM, MLIR

AI編譯器和傳統編譯器的本質是一樣的，都是一類能夠將不同的編程語言所表達code進行轉換的program。這也是AI編譯器之所以被稱之為“編譯器”的原因。

兩者的聯系

因為AI編譯器出現的比較晚，所以在設計的時候往往會借鑒傳統編譯器的思路：

• 兩者的理念比較類似。兩者都力求通過一種更加通用，更加自動化的方式進行程序優化和代碼生成，從而降低手工優化的effort。
• 兩者的軟件結構比較類似。一般都分成前端，IR，后端等模塊。其中前端負責講不同的語言的描述轉換成統一的IR表述，后端通常會對IR表示進行優化，最終生成可執行的code。IR層用來解耦前端和后端，降低集成的effort。
• 兩者的優化方式比較類似。通常編譯器都會對code進行一系列的優化，從而提高performance或者減少memory footprint等。AI編譯器和傳統編譯器都是通過在IR上面，run各種各樣的pass，進行優化的。而且，AI編譯器往往還會借鑒傳統編譯器中的一些pass，比如constant folding, dead code elimination等
• AI編譯器通常會依賴於傳統編譯器。AI編譯器在IR上面，對model進行優化之后，通常會有lowering的過程，將優化后的high-level IR轉換成傳統編譯器的low-level IR，然后依賴傳統編譯器，做最終的機器碼生成。

兩者的區別

兩者最根本的區別是應用場景的區別：

• AI編譯器是把一個深度學習模型轉換成executable。這里可以把一個深度學習模型理解成一段用DSL（Domain Specific Language）描述的code，而executable就是一段用硬件能理解的機器碼描述的code。這正好能對應到compiler的定義。
• 傳統編譯器是把一段用高級語言編寫的code轉換成executable。這里的高級語言可能是C/C++等。這也能夠對應到compiler的定義。

應用場景的區別導致了兩者在設計上不同：

• 兩者的IR表達層次有區別。AI編譯器一般會有一套high-level的IR，用來更抽象的描述深度學習模型中常用的high-level的運算，比如convolution，matmul等。而傳統編譯器的IR更偏low-level，用於描述一些更加基本的運算，比如load，store，arithmetic等。有了high-level的IR，AI編譯器在描述深度學習模型的時候會更加方便。
• 兩者的優化策略有區別。AI編譯器因為是面向AI領域的，在優化時，可以引入更多領域特定的先驗知識，從而進行更加high-level，更加aggressive的優化。比如說：
• AI編譯器可以在high-level的IR上面做operator fusion等，而傳統編譯器在做類似的loop fusion的時候往往更加保守。
• AI編譯器可以降低計算的精度，比如int8, bf16等，因為深度學習模型對計算精度不那么敏感。但傳統編譯器一般不會做這種優化。

對神經網絡優化，盡量減少邏輯判斷，一算到底是最好的。對內存要盡可能優化，降低內存占用。

神經網就是一組矩陣計算。神經網編譯器就是將這組計算針對平台盡可能加速。

編譯神經網絡，把一張張計算圖編譯成cpu的gpu的，或者是某些專用的AI計算設備，比如google 的TPU的指令集合。具體說來就是先來個圖剪枝，再來個拓撲序遍歷計算圖，一邊遍歷，一邊映射為中間表示。

后面從中間表示到指令集就大同小異了。

神經網絡編譯器大概有TVM/Glow/TensorRT/TensorComprehension/XLA/Tiramisu。這些針對的都是神經網絡模型推理階段的優化，從神經網絡模型到機器代碼的編譯。一般過程是 神經網絡模型->圖優化->中間代碼生成（例如Halide)->中間代碼優化（例如TC/Tiramisu使用多面體模型進行變換）->機器代碼。編譯的是神經網絡的模型，優化的是網絡模型本身，各層數據數據存儲的方式（如分塊存儲，nchw，nhcw），各個算子（如mlp，conv）的計算方式（如向量化，分塊）等等。

傳統編譯器（GCC，Clang這些）的編譯范圍更廣，是從源代碼到機器代碼的編譯，輸入是一段完整的代碼，經過了詞法分析，語法分析，語義分析，中間代碼生成，優化，最后到機器代碼。

聯系：

首先是神經網絡編譯器，從中間代碼到機器代碼的過程，可能就對應了傳統編譯器的整個編譯過程，比如Halide->機器代碼

然后，目標都是都要針對目標處理器進行的優化。無論是什么代碼/模型，最后的優化，就是如何最大化利用硬件，比如cache的命中率，計算速度啥的，最終目標都是生成好的機器代碼。

神經網絡編譯器，可以對應用做很多很強的假設，主要以嵌套循環的計算為主，所以可以針對性的進行優化。

傳統編譯器的前端也非常厚重，都是以編程語言為輸入來生成IR的。而神經網絡編譯器的主要問題，還是性能優化和適配，所以基本都不做前端，直接用代碼手動構造IR。

針對deep learning的編譯器，把應用限制在tensor operator上，做domain specific optimization。傳統編譯器面向的程序更加general。前者更偏上層，只需要考慮deep models，流行的deep models基本算子就卷積和矩陣乘，后者更偏底層。

以TVM和LLVM舉例，TVM拿到模型的計算圖，先用Relay做一下圖切分，算子融合，conv-bn-relu之類的，也有人做multiple conv fusion，這一步是graph-level的優化；之后再到算子層面，現在的deep compiler側重於循環優化，這部分在傳統編譯器里研究的很多，即使是deep learning領域，能做的domain specific的優化也沒多少，auto tuning做的主要還是tiling的參數 （AutoTVM / FlexTensor (ASPLOS 2020) / Ansor (OSDI 2020)）。做完operator-level的優化，TVM IR轉成LLVM IR，再借助LLVM的各種后端生成可執行代碼。

要部署一個模型，后端可以選擇使用手調庫，比如廠商庫，MKLDNN, CuDNN，某些廠商的，或者第三方的Blas庫，算子庫，比如阿里的MNN；另外一條路就是選擇deep compilers，做代碼生成。

先說deep compiler的缺點。首先編譯器能做的工作比較有限，實際的部署，要考慮到模型設計，模型壓縮之類的。另外因為比較偏上層，代碼生成部分交給了black-box compiler, 很難做到匯編級的調優，能在tuning中避免shared memory bank conflicts，但是，並不能優化掉register bank conflicts，在現有的DSL中也缺乏底層的表達，相比於某些手調庫，最終性能不太行。比如說，某些人專門做Winograd Conv的優化，性能都快接近理論極限了 （ppopp 2020）。能想到的缺點都非常細節，覺得未來很容易解決，比如GPU的prefetch，現在TVM里面，用prefetch怎么選它的size和offset基本都會導致性能變差。

但是，手調庫的缺點更加明顯，除了耗費人力外，做的優化也是general的，無法cover到具體的input configuration。即使是針對某些input，選擇調用不同的kernel，也非常有限。比如MKL-DNN，CuDNN雖然是廠商庫，代表了手調的state-of-the-art，可能對3 * 3的卷積做了特殊優化，但對於某些大的feature map或者大的kernel size性能就很差。在某個具體網絡上，通過auto-tuning，超過MKL-DNN和CuDNN並不難。AMD的就更不用說了，性能太差了，針對CUDA做的調優，用hipify工具轉到ROCm上，性能都強。

自動調優最重要的是調優之后的性能，其次是調優的時間。

對TVM了解比較深，對其他的deep compiler了解不多。至少相比於主流框架Torch/TensorFlow來看，當然考慮了這些框架用的底層庫，在某個網絡上，比如ResNet-18，針對Input大小為(1, 3, 224, 224)做調優，超過還不算太難。因為做的就是inference optimization，實際部署模型的時候，input size都是運行時不再變的，所以這條路可行。

調優時間上，Ansor調一個網絡大概一天左右，比較短了。Facebook有工作做貪心搜索，能把調優時間降到一分鍾以內，最終性能也不算差。

如果指的是針對神經網絡的編譯器，相對傳統編譯器最大的不同，引入了multi-level IR。

傳統編譯器里分為前端，優化和后端，其中前端和語言打交道，后端和機器打交道，現代編譯器的的前端和后端分的很開，共同橋梁就是IR。IR可以說是一種膠水語言，注重邏輯，去掉了平台相關的各種特性，這樣為了支持一種新語言或新硬件都會非常方便。

由於神經網絡結構領域的特殊性，這類編譯器不光要解決跨平台，還有解決對神經網絡本身的優化問題，原先一層的IR就顯得遠遠不夠，原因在於如果設計一個方便硬件優化的low level的語言，幾乎很難從中推理一些NN中高階的概念進行優化。比如說LLVM，很難把一連串的循環理解成卷積。一個完善的High level IR，至少需要包括對計算圖的表示（DAG, let binding），滿足對tensor和operator的支持。

 

 神經網絡編譯器或者深度學習編譯器（下稱 DL 編譯器），屬於一種領域特定編譯器，專門用於將神經網絡的訓練/推理部署到 CPU、GPU、NPU 上。與傳統的編譯器有着類似的結構，有很多共用的部分，同時也有自己的側重點。

關於 DL 編譯器，更多談一下 edge 端 DL 編譯器。

1. DL 編譯器產生的背景

早期神經網絡部署的側重點在於框架和算子庫。神經網絡可以由數據流圖來表示，圖上的節點就是算子（比如 Conv2D、BatchNorm、Softmax），節點之間的連接代表 Tensor。由於數據流圖很直觀，很多框架的 Runtime 采用了類似 Caffe 的方式，運行時通過一定的順序（例如直接 Post order DFS）分配 Tensor、調用算子庫就行了。優化重點在於優化算子庫的性能。

但隨着時間的發展，這種直觀的部署方式也逐漸暴露出一些問題。

• 越來越多的新算子被提出，算子庫的開發和維護工作量越來越大

比如提出一個新的 Swish，算子庫就要新增 Swish 的實現，還要有優化、測試。Swish由一些基礎的一元二元算子組成。

• NPU 的爆發導致性能可移植性成為一種剛需

大多數 NPU 作為一種 ASIC 在神經網絡場景對計算、存儲和 data movement 做了特殊優化，對能效比相對 CPU、GPU 要好很多。在移動端和 edge 端，越來越多的 NPU 開始出現。同時 NPU 的 ISA 千奇百怪，一般也缺乏 GCC、LLVM 等工具鏈，使得已有的針對 CPU 和 GPU 優化的算子庫，很難短期移植到 NPU 上，充分利用硬件的能力達到較好的性能。

• 更多可優化的點得到關注

早期 CPU 和 GPU 上帶寬問題不是很明顯，大家更多關注單個算子的性能。但在移動端和 edge 端的應用中，逐漸遇到了帶寬跟不上算力的問題，在這些 target 上增大帶寬，意味着功耗和成本的上升，利用算子間的 fusion 和調度，節省帶寬開始被重視起來。

2. 與傳統編譯器前端的異同

傳統編譯器多接受文本類型的編程語言，通過 lexer 和 parser 構造 token 和 AST。

DL 編譯器接收的一般是 DL 框架的模型文件，例如 TensorFlow 的 pb、PyTorch 的 pth，還有 ONNX 等。DL 編譯器一般把模型的導入模塊叫做 importer，將 DL 框架的模型轉換為 DL 編譯器的 IR，只跟模型文件格式和 IR 表示耦合，要支持新的框架，只需要新增一個 importer 就行了。

3. 與傳統編譯器中后端的異同

DL 編譯器和傳統編譯器一樣，使用 Constant Folding、DCE、CSE 等對 IR 進行優化。

除此之外，DL 編譯器還會有一些領域特定的圖優化：

• 合並冗余、消除無意義的 Transpose、Reshape、Pad
• 合並 BatchNorm 到 Conv2D、MatMul
• 對於先 Add 后激活的殘差結構，可以將一路輸入作為另一路 Conv2D 的初始值

目前大多數圖優化，還是根據經驗人工編寫 rules，同樣有着工作量越來越大，容易陷入局部最優的問題。有一些研究已經開始解決這些問題。應用了傳統編譯器界研究了很多年的 Equality Saturation 技術。

圖優化之后 DL 編譯器，還要進行一些 ISA 相關的優化：

• Layout

選擇 NCHW 還是 NHWC 還是 NCHW16c 等等，對於算子在特定 ISA 上的效率，會產生影響，需要納入 cost-model

• Tiling

一些 NPU 利用高速片上內存進行計算，容量一般都很有限，編譯器需要對大塊的計算進行 tiling。對於 Conv2D 這類數據復用很多的計算，如何進行 tiling 對性能和帶寬，也有很大影響，選擇 tiling 參數，也需要納入 cost-model

• Fusion

一些 NPU 可以 fusion Conv2D 和激活，甚至 fusion 一段一元二元算子組成的計算圖。編譯器需要根據硬件，提供的能力和 cost-model 選擇合適的 fusion 區域，如果貪心去匹配，也容易產生次優結果。

• Partition

對於 CPU、DSP、GPU、NPU 組成的異構系統，編譯器需要考慮算力、帶寬、數據交換的代價，對計算圖進行合理地切分。

這幾個優化有時候也需要同時考慮，比如 fusion 多層 Conv2D 時的 tiling 和單層又有不同。

很多場景下計算圖中的 Shape 是已知的，在方便了上述優化的同時，還解鎖了下面幾個優化：

• 峰值最小的內存分配

因為分配釋放序列和每次分配的 Buffer 大小已知，可以找到每個 Buffer 的最優分配位置，使得內存峰值占用最小

• Concat 消除

對於一些特殊情況，可以通過將幾個算子輸出的 Buffer 分配到一起，從而避免運行時 Concat 的發生。比較常見的是 densenet 中 Concat 的消除。

4. DL 編譯器特別的地方

DL 編譯器因為領域特定，還包含一些特別的功能。

• 稀疏

稀疏存儲 Tensor 可以降低帶寬。一些 NPU 還可以通過跳過無用計算的方式加速稀疏 Tensor 的計算。

DL 編譯器需要根據數據、Weights 的分布合理選擇，對某個 Tensor 是否進行稀疏。

• 量化

很多場景下神經網絡的推理，不需要太高的數據精度。int8 甚至 int4 已經在工業界落地。模型量化分為訓練感知量化（QAT）和訓練后量化（PTQ）。大部分用戶使用 PTQ，編譯器需要利用用戶提供的校准集（calibration dataset），得出需要量化的 Tensor 的數據分布，選擇非飽和或者飽和量化。

為了簡化，將”面向神經網絡的編譯器“簡稱為"AI編譯器"。

1. 關於AI編譯器和傳統編譯器的區別和聯系，從形式上可以理解為是輸入和輸出的區別。AI編譯器的輸入是建模的DSL描述(可能是python，比如TensorFlow/PyTorch，也可能是Lua，比如上一代的Torch，還可能是Caffe時代的PB描述文件，如果自己手寫一個AI框架的自定義DSL)，輸出通常是傳統編譯器的輸入（LLVM IR也可以視為是廣義的傳統編譯器的輸入）。傳統編譯器的輸入是傳統編程語言描述的代碼，輸出的是硬件可執行碼。

2. 透過形式，再深究一下背后的東西。AI編譯器和傳統編譯器的優化原理，有很多共通的地方，比如:

• 計算圖層面的循環不變量優化(Loop Invariant Node Motion)，高級語言層面的循環不變量優化(Loop Invariant Code Motion)。
• 計算圖層面的常量折疊和高級語言層面的常量折疊。
• 計算圖層面的peep hole optimization(模板匹配)，高級語言層面的peep hole optimization。
• 計算圖層面的strength reduction優化(比如針對Transformer模型的冗余padding計算消除優化，在LightSeq，Faster Transformer的開源代碼里都可以看到)，高級語言層面的strength reduction優化。
• 將大量計算零碎算子進行fusion&codegen優化，減少AI框架和訪存overhead的優化，將多條高級語言指令進行融合，減少中間變量的訪存操作，通過寄存器中轉優化，目的上是相似的(細節原理上是不同的)
• 還有類似TASO這樣的工作等等。

本質上都是在一種或多種表達形式上進行變換，變換的目的是為了優化，優化的標的可能是性能、顯存/內存，通信量、功耗等等，在計算圖上面結合不同的約束條件，進行變換工作。從這個層面來看，大量的傳統編譯領域技術，在AI編譯領域的應用，只是施加的層次不同。

3. 與此同時，也會存在一些細節層面的區別。最大的一個區別，AI編譯器作為一個domain specific的compiler，其實多了不少可以利用這個domain特性使巧勁的地方，舉幾個例子：

• 自動分布式並行。自動分布式並行，可以在不同層面來進行推進，一種方式是在更靠近編譯的IR層(比如HLO IR以及TorchScript的IR)來完成自動並行策略的探索。另一種方式是在更靠近建模層的圖表示層來做，比如TF Graph/JAX Graph/PyTorch NN module。從系統極致角度來考慮，前者更為究竟，這是看到G-shard，MindSpore的作法，從實現的工程量/效果回報速度來看，后者更為practical，這是看到Horovod/DeepSpeed/Megatron的作法。
• 關於算子優化，也有不同的作法。一種是通過自動codegen的作法，進行批量化生成，另一種是通過手寫（或半手工，類似ATLAS這種計算庫里的作法)開發精細的kernel，獲得極致的性能。如果AI workload高度diversified，前者更有效率，如果AI workload呈現半收斂態，其實后者反而效率更高。對於新硬件，多出了show case和長尾case的不同考慮，讓這個問題變得更復雜了。
• 結合一些workload甚至業務層面的特點，可以起到“四兩撥千斤”的優化效果。幾個比較具體的例子，推薦類模型涉及到ID類特征的處理，可能涉及到對字符串類源特征的處理，提前在預處理環節對字符串做ID化，還是在模型里做ID化，對性能影響會非常明顯，而這個優化其實不需要復雜的系統優化技術就能達到。另一個例子，如果能夠對一些重要的建模庫進行干預，在模型寫法上，對后端AI框架更為友好，實際上能大大簡化后端優化的復雜性，Google開源出的Transformer的代碼其實就有TPU-friendly的痕跡。

這些巧勁得以發揮的一個關鍵原因，當視野集中在AI domain的關鍵workload時，可以結合這些workload的特性做一些看起來"overfit"，但實現效率更高的設計妥協。而傳統編譯器，因為打擊的workload多樣性更強（通用域編譯器和domain-specific編譯器的區別），所以在leverage workload特性上會更為謹慎，通常會以workload-agnostic的角度來提供優化手段，workload-specific的優化就往往上推到各自domain里了，比如在數據庫領域利用編譯思想進行JIT優化的工作。

4.應該如何看待AI編譯器在AI系統中的地位和作用。觀點是"no silver bullet"。這就好比傳統系統領域，存在編譯器、庫(STL/glibc/...)，運行時這若干個component進行組合協同一樣，當然可以不使用STL，期望編譯器足夠的優秀，對於一個普通版本的STL alike的實現，能通過編譯手段獲得極致性能，但這樣決策涉及到在編譯器上投入的effort，是否值得就要仔細考慮了。在AI system領域，同樣會有類似的分工。對於一個workload，一族workload，整個AI worload的全場景，應該如何在AI編譯器、AI底層庫、運行時、AI建模庫之間進行職能划分，很考驗系統設計能力的事情。如果再有機會對硬件設計也有干預，影響到programming model，device compiler的設計，更具挑戰，也更有意思的事情了。

對於社區AI編譯領域的一些作法，比如需要用戶手工打標簽，標識哪段子圖可以進行JIT優化是有些diss的，時過境遷，現在覺得這種作法和AI編譯的通用優化，不是互斥矛盾的，反而可能是另一種看到了整體工作復雜性以后的trade-off考慮罷了。從系統設計的角度來說，對MLIR的設計理念的認同度也更高了，當然MLIR社區里的聲音不夠清晰統一。

1.神經網絡編譯器背景和歷史

1、早期深度學習框架，重點是框架和庫，與編譯器關系相對較弱

比如Tensorflow早期版本，在神經網絡/深度學習的編程模型上，主要進行了graph/圖和op/算子兩層抽象

• 圖層通過聲明式的編程方式，然后通過靜態圖的方式進行執行，這里其實也做了一些編譯器的事情，這里包括硬件無關和硬件相關的優化：硬件無關的優化包括編譯器通用的優化，如表達式化簡、常量折疊，也包括與深度學習/神經網絡強相關的，如自動微分等；硬件相關的優化包括簡單的算子融合、內存分配優化等。
• 算子層通常采用手寫的方式，比如GPU上基於CUDA/cuDNN。

這種方式遇到幾個問題：

• 表達上，語法不是Python原生的，算法工程師使用的易用性不夠好
• 更多的Transform出現，比如並行、量化、混合精度等
• 算子粒度和邊界提前確定后，無法充分發揮硬件的性能
• 硬件廠商提供的算子庫也不一定是性能最優的，在SIMT和SIMD的架構中，scheduling、tilling都是有很大的空間，在具體到一個模型，shape確定的情況下，開發者還有可能開發出性能更高的算子。
• AI專用芯片出現（Google TPU、華為Ascend等），3與4的情況加劇。

2、后期引入大量編譯器的技術進行改進

• 表達上的改進(Pytorch/TorchScript、JAX)

Pytorch的Eager Model是一種解決易用性的方案，雖然基本上還是圖層和算子兩層的抽象，但是整個語法基本上是Python Native的，讓算法工程師比較容易上手；不過這個方案在運行的時候，基於Python解釋器的能力，不是一種高性能的解決方案，本身與神經網絡的編譯器關系不大；但是其表達的方式成為后面框架參考的標桿，圖層的神經網絡編譯器主要就是考慮如何把這樣表達轉換到圖層的IR進行優化，目前主要有兩種方式：

AST-Based：以Pytorch TorchScript為例，主要通過Python的修飾符，把Python代碼的AST拿到，然后變換成圖層的IR，進行編譯優化。

Tracing-Based：以JAX為例，主要把Python代碼假執行一遍，保存執行序列，基於執行序列，變換到圖層IR進行編譯優化。

兩種方案各有優缺點，第一種方案實現復雜，第二種方案在一些處理上有限制(比如控制流的處理)。

• 性能上的優化(XLA/TVM/TC)

性能上的優化思路其實比較統一，就是打開圖和算子的邊界，進行重新組合優化。

XLA：基本上的思路，把圖層下發的子圖中的算子全部打開成小算子，然后基於這張小算子組成的子圖，進行編譯優化，包括buffer fusion、水平融合等，這里的關鍵是大算子怎樣打開、小算子如何重新融合、新的大的算子(kernel)怎樣生成，整體設計主要通過HLO/LLO/LLVM層lowering實現，所有規則都是手工提前指定。

TVM：分為Relay和TVM兩層，Relay主要關注圖層，TVM主要關注算子層，總體思路與XLA是類似的，也是拿到前端給一張子圖進行優化，Relay關注算子間的融合、TVM關注新的算子和kernel的生成，區別在於TVM是一個開放的架構，Relay目標是可以接入各種前端，TVM也是一個可以獨立使用的算子開發和編譯的工具（基於Halide IR，最新演進到自己定義的TIR），TVM在算子實現方面采用了compute和schedule分離的方案，開發人員通過compute，設計計算的邏輯，通過schedule來指定調度優化的邏輯。

TC(Tensor Comprehensions)：開發者發現算子的計算邏輯的開發，比較容易的，但是schedule的開發非常困難，既要了解算法的邏輯，又要熟悉硬件的體系架構，更重要的是，圖算邊界打開后，小算子融合后，生成新的算子和kernel，這些新的算子compute是容易確定的（小算子compute的組合），但是schedule卻很難生成，傳統的方法就是事先定義一大堆schedule模板，萬一組合的新算子不在模板之內，性能就可能比較差，甚至出錯； TC則希望通過Polyhedra model，實現auto schedule，降低開發門檻，當然這個項目基本已經停更了，類似的工作在MLIR、MindSpore上，還在不停發展。

• 圖層和算子層的IR表達

在神經網絡編譯器發展過程中，有多種IR的出現，各有特點：

圖層IR：朴素的DataflowIR、函數式IR、函數式圖IR、SSA風格IR

算子層IR：HalideIR、LLVM等

圖算融合表達：MLIR

以前分析過圖層IR，供參考：

2．回到問題

1、神經網絡編譯器與傳統編譯器的相同點

神經網絡編譯器和傳統編譯器一樣，也是有前端表達、硬件無關優化和硬件相關優化、最后的codegen等，整體結構是類似的。

2、神經網絡編譯器與傳統編譯器的區別

主要體現在神經網絡編譯器，像數據庫的SQL引擎/向量化引擎一樣，一個特定領域的編譯器，這些領域特征包括：以Python為主的動態解釋器語言的前端、多層IR設計（圖層/算子層/codegen）、面向神經網絡的特定優化（自動微分、量化/混合精度、大規模並行、張量運算/循環優化等）。

• 編譯前端解析

與傳統編譯器不同，神經網絡編譯器通常不需要lexer/parser，而是基於前端語言（如Python）的AST，將模型解析，構造為計算圖IR，側重於保留shape、layout等Tensor計算特征信息，當然部分編譯器還能保留控制流的信息。

這里的難點在於，Python是一種靈活度極高的解釋執行的語言，像弱類型、靈活的數據結構等，而神經網絡編譯器本質上是偏靜態，兩者之間的完全轉化是不大可能的。

• 多層IR設計

為什么需要多層IR設計，主要是為了同時滿足易用性與高性能這兩類需求。為了讓開發者使用方便，框架前端(圖層)會盡量對Tensor計算進行抽象封裝，開發者只要關注模型和粗粒度OP；而在后端算子性能優化時，又可以打破算子的邊界，從更細粒度的循環調度等維度，結合不同的硬件特點完成優化。因此，多層IR設計無疑是較好的選擇。

High-level IR（圖層IR），如XLA的HLO，TVM的Relay IR，MindSpore的MindIR等，重點關注非循環相關的優化。除了傳統編譯器中常見的常量折疊、代數化簡、公共子表達式等優化外，還會完成Layout轉換，算子融合等優化，通過分析和優化現有網絡計算圖邏輯，對原有計算邏輯進行拆分、重組、融合等操作，減少算子執行間隙的開銷，提升設備計算資源利用率，實現網絡整體執行時間的優化。

Low-level IR，如TVM的TIR，HalideIR，以及isl schedule tree[7]等。針對Low-level IR主要有循環變換、循環切分等調度相關的優化，與硬件intrinsic映射、內存分配等后端pass優化。自動調度優化，主要包含了基於搜索的自動調度優化（如ansor）和基於polyhedral編譯技術的自動調度優化（如TC和MindAKG）。

有人可能會問，圖層和算子層的表達和編譯能否放在一起？也許可以，但是明顯看到這樣做面臨幾個挑戰：

1、整圖展開到原子算子，看上去編譯的規模/復雜度，指數級上升

2、顯然圖編譯優化的問題和算子編譯優化的問題，有明顯的區別，一個關注變換和融合，另外一個關注循環優化，放在一起對編譯器實現的復雜度是個比較大的挑戰

3、要看到硬件供應商和框架供應商目前是分開的，兩者總是需要一個邊界。

• 面向神經網絡的特定優化

自動微分：BP是深度學習/神經網絡最有代表的部分，目前相對已經比較成熟，基於計算圖的自動微分、基於Tape和運算符重載的自動微分方案、基於source2source的自動微分，都是現在主流的方案。

並行優化：隨着深度學習的模型規模越來越大，模型的並行優化，也成為編譯優化的一部分，包括：數據並行、算子級模型並行、Pipeline模型並行、優化器模型並行和重計算等

張量計算/循環優化：循環優化其實是一個古老的編譯器的難題，在高性能計算領域，循環優化已經研究了幾十年，一直沒有很好的解決，但是看上去，深度學習/神經網絡領域的問題，要簡單一點，原因是這個領域大量的以Dense的矩陣運算為主，不像高性能計算領域那么復雜(大量稀疏/非規則的矩陣和向量運算)，這為循環優化帶來了很大的空間，不過即便是這樣，自動scheduling、自動tilling、自動向量化這些理想的方案和技術也還遠遠沒有成熟。

量化/....：推理側常用的一些變換。

3）神經網絡編譯器未來的方向探討

編譯器形態：也許需要兩類編譯器同時存在，一類是面向極致高性能的AOT編譯器，同時這類編譯器對NPU更加友好；另外一類是JIT編譯器，適合與動態圖配合；

IR形態：需不需要MLIR這種統一的形態？

自動並行：配合Cost model，提供自動並行優化的能力；

自動Scheduling/Tilling/Tensorizing：可能很難全部做到，能支持大部分也可以。

 

參考鏈接：

https://www.zhihu.com/question/396105855

參考文獻：

[1] Tensorflow. https://github.com/tensorflow/tensorflow

[2] MindSpore. https://gitee.com/mindspore/mindspore

[3] Tvm. https://tvm.apache.org/

[4] XLA. https://www.tensorflow.org/xla

[5]TensorComprehensions.https://github.com/facebookresearch/TensorComprehensions

[6] Li, Mingzhen and Liu, Yi, etc. The deep learning compiler: A comprehensive survey. IEEE Transactions on Parallel and Distributed Systems. 2020

[7] Polyhedral Compilation. https://polyhedral.info. Accessed February 4, 2020.

[8] Zheng, Lianmin, etc. Ansor: Generating high-performance tensor programs for deep learning. Symposium on Operating Systems Design and Implementation. 2020

[9] MindAKG. https://gitee.com/mindspore/akg

[10] ^TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions https://cs.stanford.edu/~zhihao/papers/sosp19.pdf

[11] ^^abEquality Saturation for Tensor Graph Superoptimization https://arxiv.org/pdf/2101.01332.pdf