[[409589]]

最近的十幾年深度學習發(fā)展十分迅速，業(yè)界出現(xiàn)了很多深度學習算法開發(fā)框架。同時，由于深度學習具有廣泛應(yīng)用場景和對算力的巨大需求，我們需要將深度學習算法運行在各種通用和專用的硬件上，比如各種類型的CPU，GPU，TPU，NPU等。那么這就出現(xiàn)了框架和硬件之間的組合爆炸，如圖 1所示。比如說TensorFlow要支持GPU計算，就要把tensorflow里面的所有算子開發(fā)一個GPU版本，如果又要支持D芯片，又需要把每個算子開發(fā)一個D芯片的版本。這個過程無疑非常耗時耗力。

圖 1

于此同時，我們現(xiàn)在有非常多的算法網(wǎng)絡(luò)，比如說YOLO, BERT, GPT等等。而這些算法網(wǎng)絡(luò)是是由不同類型、不同shape，不同連接關(guān)系的算子組成的。最終它們又運行在不同種類和型號的硬件上面。這就導(dǎo)致人工去為每個場景開發(fā)和實現(xiàn)最優(yōu)算子成本很高。這里舉了兩個例子，如圖 2所示，算子融合是一個常見的性能優(yōu)化方法，在融合之前，每個算子計算前后都需要把數(shù)據(jù)從內(nèi)存讀到緩存，再從緩存寫回到內(nèi)存。而融合之后，可以避免算子之間內(nèi)存讀寫從而提高性能。傳統(tǒng)的做法就是人工去根據(jù)算子連接關(guān)系開發(fā)融合算子，但是不同網(wǎng)絡(luò)不同類別算子連接關(guān)系幾乎不可能完全枚舉。另一個例子就是算子調(diào)優(yōu)，算子實現(xiàn)過程有很多參數(shù)會影響性能，但是傳統(tǒng)人工算子開發(fā)方式很難去表達和維護這些參數(shù)，并且對這些參數(shù)進行調(diào)優(yōu)從而實現(xiàn)不同shape和硬件的最優(yōu)性能。

圖 2

深度學習編譯器正是為了解決上面一系列問題而誕生的，它可以作為框架和硬件之間的公共組件和橋梁，最終希望實現(xiàn)的目標是我們只用開發(fā)一次，就能夠為自動為任何設(shè)備生成最優(yōu)代碼。比如為CPU開發(fā)的算子可以幾乎原封不同的用于GPU和D芯片，從而顯著降低成本。

這里簡單介紹一下深度學習編譯器的組成部分和功能，如圖 3所示。首先它的前端是從不同的框架拿到計算圖，并且使用這個High level IR的數(shù)據(jù)結(jié)構(gòu)來表示，然后在這個階段進行一系列圖優(yōu)化，比如常量折疊，算子融合，等價替換等。這里展示了一個等價替換的例子，原來計算圖是這樣的，我們給它換一個計算方式，結(jié)果不變，但是性能可能更優(yōu)。接著，對于計算圖里面的每一個算子，采用DSL一種領(lǐng)域特定的語言來描述算子的計算過程和對算子進行優(yōu)化。比如對算子進行tiling,多核，double-buffer等優(yōu)化。由于算子的計算過程通常是用多重循環(huán)來實現(xiàn)的，比如說矩陣乘法是一個三重的循環(huán)。深度學習編譯器可以很方便的對循環(huán)進行各種變換，并且對這些變換的參數(shù)進行調(diào)優(yōu)，從而得到不同shape和硬件的最佳算子實現(xiàn)。最后，基于low level IR為不同硬件生成具體的代碼。

圖 3

最后介紹下業(yè)界已有的編譯器項目。目前生態(tài)最完善，開源的，框架不依賴的項目首推TVM，已經(jīng)被很多公司所采用。TVM流程如如圖 3a所示，TVM可以導(dǎo)入各個框架的模型，例如TensorFlow pb，onnx，TorchScript等模型，統(tǒng)一用TVM稱為Relay的High level IR進行表示。IR中每個算子采用了Tensor expression的DSL來進行計算描述和調(diào)度。這個DSL采用Einstein’s notation的方式進行算子的compute描述，算子compute一般體現(xiàn)為多重for循環(huán)。然后基于Halide思想使用schedule對這個多重for循環(huán)進行各種變換，例如循環(huán)合并，split，順序變換等等。最后，lower到low-level IR生成具體的device代碼并進行推理。

這里再簡單介紹下TVM具體如何生成最優(yōu)的算子代碼。上面介紹了算子需要進行compute描述，然后需要對compute對應(yīng)的多重for循環(huán)進行調(diào)度變換，即schedule。TVM的算子生成和調(diào)優(yōu)經(jīng)歷了3代發(fā)展。第一代TVM/AutoTVM，這一代需要用戶編寫算子的compute和算子的schedule，AutoTVM與TVM的區(qū)別在于可以在schedule定義一些可變的參數(shù)，然后采用例如遺傳算法進行參數(shù)調(diào)優(yōu)。例如把一個loop切分為2段，那么在哪里進行切分是可以進行優(yōu)化的。第二代AutoScheduler (Ansor)，這一代只需要用戶開發(fā)算子ompute，Ansor內(nèi)部自動根據(jù)一些規(guī)則進行調(diào)度變換。由于調(diào)度開發(fā)需要同時熟悉TVM的表達機制和底層硬件原理，schedule開發(fā)往往具有很高的難度，因此Ansor可以顯著降低開發(fā)人員工作量和開發(fā)難度，缺點就是Ansor調(diào)優(yōu)時間很長，往往需要1小時才能調(diào)優(yōu)1個算子。以卷積網(wǎng)絡(luò)為例，Ansor在部分場景能超過TensorFlow算子性能，距離TensorRT實現(xiàn)有一定差距。第三代Meta Schedule (AutoTensorIR)才處于起步階段，預(yù)期會對調(diào)優(yōu)速度和性能進行優(yōu)化，暫時還不可用，我們拭目以待。

TVM的落地包括華為D芯片TBE算子開發(fā)工具，在TVM的基礎(chǔ)上增加了D芯片的代碼生成支持。TVM采用了Halide計算+調(diào)度的路線，還有另外一種采用polyhedral算法路線的編譯器，比如Tensor Comprehensions，Tiramisu，華為自研的AKG等。這種方法跟Ansor一樣，也只需要用戶開發(fā)算子compute，無需開發(fā)schedule，因此對用戶也較為友好。其中AKG已經(jīng)用在了MindSpore的圖算融合里面。其他的深度學習編譯器還有TensorFlow的XLA、TensorRT等，大家可能已經(jīng)用過。

總之，深度學習編譯器具有很多優(yōu)勢。比如易于支持新硬件，避免重復(fù)開發(fā)，采用一系列自動優(yōu)化代替人工優(yōu)化，可以實現(xiàn)極致性價比等。目前深度學習編譯器也有一些不足，仍然出于一個快速發(fā)展的狀態(tài)。例如調(diào)優(yōu)時間長，對于復(fù)雜的算子無法有效生成，一個模型中深度學習編譯器生成的算子能超過庫調(diào)用的算子比例較低等，仍然需要大家持續(xù)投入和優(yōu)化。