摘要：深度學習編譯器能夠做爲框架和硬件之間的公共組件和橋樑，最終但願實現的目標是咱們只用開發一次，就可以爲自動爲任何設備生成最優代碼。
本文分享自華爲雲社區《深度學習編譯器簡介》，原文做者：luchangli 。前端
最近的十幾年深度學習發展十分迅速，業界出現了不少深度學習算法開發框架。同時，因爲深度學習具備普遍應用場景和對算力的巨大需求，咱們須要將深度學習算法運行在各類通用和專用的硬件上，好比各類類型的CPU，GPU，TPU，NPU等。那麼這就出現了框架和硬件之間的組合爆炸，如圖 1所示。好比說TensorFlow要支持GPU計算，就要把tensorflow裏面的全部算子開發一個GPU版本，若是又要支持D芯片，又須要把每一個算子開發一個D芯片的版本。這個過程無疑很是耗時耗力。算法
圖 1express
於此同時，咱們如今有很是多的算法網絡，好比說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代碼並進行推理。ide
這裏再簡單介紹下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等，你們可能已經用過。
總之，深度學習編譯器具備不少優點。好比易於支持新硬件，避免重複開發，採用一系列自動優化代替人工優化，能夠實現極致性價比等。目前深度學習編譯器也有一些不足，仍然出於一個快速發展的狀態。例如調優時間長，對於複雜的算子沒法有效生成，一個模型中深度學習編譯器生成的算子能超過庫調用的算子比例較低等，仍然須要你們持續投入和優化。