AI編譯器與傳統編譯器的區別與聯系
總結整理自知乎問題 針對神經網絡的編譯器和傳統編譯器的區別和聯系是什么?。
文中提到的答主的知乎主頁:@金雪鋒、@楊軍、@藍色、@SunnyCase、@貝殼與知了、@工藤福爾摩
筆者本人理解
為了不用直接手寫機器碼,我們可以用匯編語言;為了不用手寫匯編,我們開發出了高級語言,并用編譯器將我們寫的高級語言編譯成匯編。因此,傳統編譯器主要解決的問題是要降低編程的難度,其次是優化程序性能。其輸入是高級語言,輸出是硬件可執行碼。
而對于神經網絡編譯器,其輸入是一個深度學習模型(這里可以看作是一種 DSL),輸出也是硬件可執行碼。即神經網絡編譯器通常是不需要傳統編譯器的厚重的前端部分(即詞法分析、語法分析、語義分析等)的,其輸入直接就是一種描述深度學習模型的IR,我們對這個輸入進行優化,并針對特定的硬件目標生成可執行代碼。對于深度學習模型,我們設計神經網絡編譯器主要是為了提高推理時的速度。即神經網絡編譯器主要解決的問題是要優化模型的推理性能。
另外,為了同時前端簡單方便地使用 Python 代碼和后端對算子進行高效優化,神經網絡編譯器通常采用多層IR的形式。
由于神經網絡編譯器是近幾年才開始大規模發展的領域,因此,其在設計上會借用一些傳統編譯器的通用代碼優化方式,如表達式化簡,常量折疊等。而更關鍵的是,作為一種專用于深度模型推理的DSL,神經網絡編譯器可以根據AI模型的常見計算范式(如矩陣乘法、卷積等),做一些更強、更激進的假設,有更 domain specific 的優化。
以下是各位知乎大佬的回答:
@金雪鋒
一、神經網絡編譯器出現的背景和歷史
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
二、神經網絡編譯器與傳統編譯器的聯系與區別
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 tree7等。針對Low-level IR主要有循環變換、循環切分等調度相關的優化,與硬件intrinsic映射、內存分配等后端pass優化。其中,當前的自動調度優化主要包含了基于搜索的自動調度優化(如ansor)和基于polyhedral編譯技術的自動調度優化(如TC和MindAKG)。
有人可能會問,圖層和算子層的表達和編譯能否放在一起?也許可以,但是明顯看到這樣做面臨幾個挑戰:
1、整圖展開到原子算子,看上去編譯的規模/復雜度指數級上升
2、顯然圖編譯優化的問題和算子編譯優化的問題是有明顯的區別,一個關注變換和融合,另外一個關注循環優化,放在一起對編譯器實現的復雜度是個比較大的挑戰
3、要看到硬件供應商和框架供應商目前是分開的,兩者總是需要一個邊界。
-
面向神經網絡的特定優化
自動微分:BP是深度學習/神經網絡最有代表的部分,目前相對已經比較成熟,基于計算圖的自動微分、基于Tape和運算符重載的自動微分方案、基于source2source的自動微分都是現在主流的方案。
并行優化:隨著深度學習的模型規模越來越大,模型的并行優化也成為編譯優化的一部分,包括:數據并行、算子級模型并行、Pipeline模型并行、優化器模型并行和重計算等
張量計算/循環優化:循環優化其實是一個古老的編譯器的難題,在高性能計算領域,循環優化已經研究了幾十年,一直沒有很好的解決,但是看上去,深度學習/神經網絡領域的問題要簡單一點,原因是這個領域大量的以Dense的矩陣運算為主,不像高性能計算領域那么復雜(大量稀疏/非規則的矩陣和向量運算),這為循環優化帶來了很大的空間,不過即便是這樣,自動scheduling、自動tilling、自動向量化這些理想的方案和技術也還遠遠沒有成熟。
量化 /…:推理側常用的一些變換,不展開了
三、神經網絡編譯器未來的方向探討
編譯器形態:也許需要兩類編譯器同時存在,一類是面向極致高性能的AOT編譯器,同時這類編譯器對NPU更加友好;另外一類是JIT編譯器,適合與動態圖配合;
IR形態:需不需要MLIR這種統一的形態?
自動并行:配合Cost model,提供自動并行優化的能力;
自動Scheduling/Tilling/Tensorizing:可能很難全部做到,能支持大部分也可以。
@楊軍
- 關于AI編譯器和傳統編譯器的區別和聯系,藍色的圖(下圖)比較形象, 從形式上可以理解為是輸入和輸出的區別。AI編譯器的輸入是建模的DSL描述(可能是python,比如TensorFlow/PyTorch,也可能是Lua,比如上一代的Torch,還可能是Caffe時代的PB描述文件,以及如果自己手寫一個AI框架的自定義DSL),輸出通常是傳統編譯器的輸入(LLVM IR也可以視為是廣義的傳統編譯器的輸入)。傳統編譯器的輸入是傳統編程語言描述的代碼,輸出的是硬件可執行碼。
- 透過形式,再深究一下背后的東西。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編譯領域的應用,只是施加的層次不同。
-
與此同時,也會存在一些細節層面的區別。最大的一個區別,我個人認為是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優化的工作。
-
-
一個更重要的問題我覺得是應該如何看待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的設計,那就是一個更具挑戰,也更有意思的事情了。
@SunnyCase
神經網絡編譯器或者深度學習編譯器(下稱 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 的數據分布,選擇非飽和或者飽和量化(具體細節不再細說)。
@貝殼與知了
感覺前面幾位大佬講的已經講地很細致深入了,這里稍微闡述一下我自己的理解 _
1. 先說兩者的本質
參考wiki上面對compiler的定義:
In computing, a compiler is a computer program that translates computer code written in one programming language (the source language) into another language (the target language).
https://en.wikipedia.org/wiki/Compileren.wikipedia.org/wiki/Compiler
從這一點上來看,AI編譯器和傳統編譯器的本質是一樣的,都是一類能夠將不同的編程語言所表達code進行轉換的program。我想這也是AI編譯器之所以被稱之為“編譯器”的原因。
2.再說兩者的聯系
因為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,然后依賴傳統編譯器去做最終的機器碼生成。
3.最后說兩者的區別
我認為兩者最根本的區別是應用場景的區別:
- 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等,因為深度學習模型對計算精度不那么敏感。但傳統編譯器一般不會做這種優化。
@工藤福爾摩
針對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有工作做貪心搜索,能把調優時間降到一分鐘以內,最終性能也不算差 (MLSys 2021)。