[[409589]]

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

圖 1

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

圖 2

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

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

圖 3

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

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

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

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