摘要：深度學習編譯器可以作為框架和硬體之間的公共元件和橋樑，最終希望實現的目標是我們只用開發一次，就能夠為自動為任何裝置生成最優程式碼。
本文分享自華為雲社群《深度學習編譯器簡介》，原文作者：luchangli 。
最近的十幾年深度學習發展十分迅速，業界出現了很多深度學習演算法開發框架。同時，由於深度學習具有廣泛應用場景和對算力的巨大需求，我們需要將深度學習演算法執行在各種通用和專用的硬體上，比如各種型別的CPU，GPU，TPU，NPU等。那麼這就出現了框架和硬體之間的組合爆炸，如圖 1所示。比如說TensorFlow要支援GPU計算，就要把tensorflow裡面的所有運算元開發一個GPU版本，如果又要支援D晶片，又需要把每個運算元開發一個D晶片的版本。這個過程無疑非常耗時耗力。
圖 1
於此同時，我們現在有非常多的演算法網路，比如說YOLO, BERT, GPT等等。而這些演算法網路是是由不同型別、不同shape，不同連線關係的運算元組成的。最終它們又執行在不同種類和型號的硬體上面。這就導致人工去為每個場景開發和實現最優運算元成本很高。這裡舉了兩個例子，如圖 2所示，運算元融合是一個常見的效能優化方法，在融合之前，每個運算元計算前後都需要把資料從記憶體讀到快取，再從快取寫回到記憶體。而融合之後，可以避免運算元之間記憶體讀寫從而提高效能。傳統的做法就是人工去根據運算元連線關係開發融合運算元，但是不同網路不同類別運算元連線關係幾乎不可能完全列舉。另一個例子就是運算元調優，運算元實現過程有很多引數會影響效能，但是傳統人工運算元開發方式很難去表達和維護這些引數，並且對這些引數進行調優從而實現不同shape和硬體的最優效能。
圖 2
深度學習編譯器正是為了解決上面一系列問題而誕生的，它可以作為框架和硬體之間的公共元件和橋樑，最終希望實現的目標是我們只用開發一次，就能夠為自動為任何裝置生成最優程式碼。比如為CPU開發的運算元可以幾乎原封不同的用於GPU和D晶片，從而顯著降低成本。
這裡簡單介紹一下深度學習編譯器的組成部分和功能，如圖 3所示。首先它的前端是從不同的框架拿到計算圖，並且使用這個High level IR的資料結構來表示，然後在這個階段進行一系列圖優化，比如常量摺疊，運算元融合，等價替換等。這裡展示了一個等價替換的例子，原來計算圖是這樣的，我們給它換一個計算方式，結果不變，但是效能可能更優。接著，對於計算圖裡面的每一個運算元，採用DSL一種領域特定的語言來描述運算元的計算過程和對運算元進行優化。比如對運算元進行tiling,多核，double-buffer等優化。由於運算元的計算過程通常是用多重迴圈來實現的，比如說矩陣乘法是一個三重的迴圈。深度學習編譯器可以很方便的對迴圈進行各種變換，並且對這些變換的引數進行調優，從而得到不同shape和硬體的最佳運算元實現。最後，基於low level IR為不同硬體生成具體的程式碼。
圖 3
最後介紹下業界已有的編譯器專案。目前生態最完善，開源的，框架不依賴的專案首推TVM，已經被很多公司所採用。TVM流程如如圖 3a所示，TVM可以匯入各個框架的模型，例如TensorFlow pb，onnx，TorchScript等模型，統一用TVM稱為Relay的High level IR進行表示。IR中每個運算元採用了Tensor expression的DSL來進行計算描述和排程。這個DSL採用Einstein’s notation的方式進行運算元的compute描述，運算元compute一般體現為多重for迴圈。然後基於Halide思想使用schedule對這個多重for迴圈進行各種變換，例如迴圈合併，split，順序變換等等。最後，lower到low-level IR生成具體的device程式碼並進行推理。
這裡再簡單介紹下TVM具體如何生成最優的運算元程式碼。上面介紹了運算元需要進行compute描述，然後需要對compute對應的多重for迴圈進行排程變換，即schedule。TVM的運算元生成和調優經歷了3代發展。第一代TVM/AutoTVM，這一代需要使用者編寫運算元的compute和運算元的schedule，AutoTVM與TVM的區別在於可以在schedule定義一些可變的引數，然後採用例如遺傳演算法進行引數調優。例如把一個loop切分為2段，那麼在哪裡進行切分是可以進行優化的。第二代AutoScheduler (Ansor)，這一代只需要使用者開發運算元ompute，Ansor內部自動根據一些規則進行排程變換。由於排程開發需要同時熟悉TVM的表達機制和底層硬體原理，schedule開發往往具有很高的難度，因此Ansor可以顯著降低開發人員工作量和開發難度，缺點就是Ansor調優時間很長，往往需要1小時才能調優1個運算元。以卷積網路為例，Ansor在部分場景能超過TensorFlow運算元效能，距離TensorRT實現有一定差距。第三代Meta Schedule (AutoTensorIR)才處於起步階段，預期會對調優速度和效能進行優化，暫時還不可用，我們拭目以待。
TVM的落地包括華為D晶片TBE運算元開發工具，在TVM的基礎上增加了D晶片的程式碼生成支援。TVM採用了Halide計算+排程的路線，還有另外一種採用polyhedral演算法路線的編譯器，比如Tensor Comprehensions，Tiramisu，華為自研的AKG等。這種方法跟Ansor一樣，也只需要使用者開發運算元compute，無需開發schedule，因此對使用者也較為友好。其中AKG已經用在了MindSpore的圖算融合裡面。其他的深度學習編譯器還有TensorFlow的XLA、TensorRT等，大家可能已經用過。
總之，深度學習編譯器具有很多優勢。比如易於支援新硬體，避免重複開發，採用一系列自動優化代替人工優化，可以實現極致價效比等。目前深度學習編譯器也有一些不足，仍然出於一個快速發展的狀態。例如調優時間長，對於複雜的運算元無法有效生成，一個模型中深度學習編譯器生成的運算元能超過庫呼叫的運算元比例較低等，仍然需要大家持續投入和優化。