全文翻譯（二）: TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
3．優(yōu)化計算圖
計算圖是在DL框架中表示程序的常用方法。圖3顯示了兩層卷積神經(jīng)網(wǎng)絡(luò)的計算圖表示示例。這種高級表示與低級編譯器中間表示（IR，如LLVM）間的主要區(qū)別，中間數(shù)據(jù)項是大型多維張量。計算圖提供了算子的全局視圖，但避免指定每個算子必須如何實現(xiàn)。與LLVM IRs一樣，計算圖可以轉(zhuǎn)換為功能等效圖優(yōu)化。利用常見DL工作負(fù)載中的形狀特異性，優(yōu)化一組固定的輸入shape。

Figure 3: Example computational graph of a two-layer convolutional neural network. Each node in the graph represents an operation that consumes one or more tensors and produces one or more tensors. Tensor operations can be parameterized by attributes to confifigure their behavior (e.g., padding or strides).

TVM利用計算圖表示應(yīng)用高級優(yōu)化：節(jié)點表示對張量或程序輸入的算子，邊表示算子間的數(shù)據(jù)依賴關(guān)系。實現(xiàn)多級圖形優(yōu)化，包括：算子fusion，將多個小算子融合在一起；常量折疊，預(yù)先計算可靜態(tài)確定的圖形部分，節(jié)省執(zhí)行成本；靜態(tài)內(nèi)存調(diào)度過程，預(yù)先分配內(nèi)存，保存每個中間張量；數(shù)據(jù)布局轉(zhuǎn)換，將內(nèi)部數(shù)據(jù)布局轉(zhuǎn)換為后端友好的形式。
現(xiàn)在討論算子融合和數(shù)據(jù)布局轉(zhuǎn)換。
算子融合
算子融合將多個算子組合到單個內(nèi)核中，無需將中間結(jié)果保存在內(nèi)存中。這種優(yōu)化可以大大減少執(zhí)行時間，特別是在GPU和專用加速器中。具體來說，認(rèn)識到四類圖算子：
（1）內(nèi)射（一對一映射，如add）
（2）歸約（如sum）
（3）復(fù)雜輸出可融合（將元素映射融合到輸出，如conv2d）
（4）不透明（不能融合，如sort）。
提供了融合這些算子的通用規(guī)則，如下所示。
多個內(nèi)射算子可以融合成另一個內(nèi)射算子。歸約算子可以與輸入內(nèi)射算子融合（例如，融合scale和sum）。如conv2d類的算子是復(fù)雜的可外fusion的，可以將元素算子fusion到輸出中。可以應(yīng)用這些規(guī)則，將計算圖轉(zhuǎn)換為融合版本。

Figure 4: Performance comparison between fused and non-fused operations. TVM generates both operations. Tested on NVIDIA Titan X.
圖4演示了對不同工作負(fù)載的影響的優(yōu)化。通過減少內(nèi)存訪問，融合算子可以產(chǎn)生高達(dá)1.2倍到2倍的加速比。
數(shù)據(jù)布局轉(zhuǎn)換。在計算圖中存儲給定張量有多種方法。最常見的數(shù)據(jù)布局選擇是列主視圖和行主視圖。實際上，可能喜歡使用復(fù)雜的數(shù)據(jù)布局。例如，DL accelerator可能利用4×4矩陣運算，需要將數(shù)據(jù)平鋪到4×4塊中，優(yōu)化局部訪問。

Figure 5: Example schedule transformations that optimize a matrix multiplication on a specialized accelerator.
數(shù)據(jù)布局優(yōu)化，可以使用更好的內(nèi)部數(shù)據(jù)布局，轉(zhuǎn)換計算圖形，在目標(biāo)硬件上執(zhí)行圖形。先指定每個算子的首選數(shù)據(jù)布局，給定內(nèi)存層次結(jié)構(gòu)指定的約束。然后，如果生產(chǎn)者和消費者的首選數(shù)據(jù)布局不匹配，將在兩者間執(zhí)行適當(dāng)?shù)牟季洲D(zhuǎn)換。
雖然高級圖優(yōu)化可以極大地提高DL工作負(fù)載的效率，但效率僅與operator library提供的一樣高。目前，支持算子融合的少數(shù)DL框架，要求算子庫提供融合模式的實現(xiàn)。隨著更多的網(wǎng)絡(luò)算子定期推出，可能的融合內(nèi)核數(shù)量可能會大幅增加。當(dāng)針對越來越多的硬件后端時，這種方法不再具有可持續(xù)性，因為所需的融合模式實現(xiàn)數(shù)量，隨著必須支持的數(shù)據(jù)布局，數(shù)據(jù)類型和加速器內(nèi)部結(jié)構(gòu)的數(shù)量，組合增長。接下來提出一種代碼生成方法，可以為給定模型的算子生成各種可能的實現(xiàn)。
4．生成張量運算
TVM通過在每個硬件后端，生成許多有效的實現(xiàn)，選擇優(yōu)化的實現(xiàn)，為每個算子生成高效的代碼。
該過程基于Halide，將描述與計算規(guī)則（或調(diào)度優(yōu)化）分離，將支持新的優(yōu)化（嵌套并行，張量化和延遲隱藏）和大量硬件后端。重點介紹TVM的特定功能。

4.1張量表達(dá)式和表空間

Figure 6: TVM schedule lowering and code generation process. The table lists existing Halide and novel TVM scheduling primitives being used to optimize schedules for CPUs, GPUs and accelerator back-ends. Tensorization is essential for accelerators, but it can also be used for CPUs and GPUs. Special memory-scope enables memory reuse in GPUs and explicit management of onchip memory in accelerators. Latency hiding is specifific to TPU-like accelerators.

引入了一種支持自動代碼生成的張量表達(dá)式語言。與高級計算圖表示不同，張量算子的實現(xiàn)是不透明的，每個算子都用索引公式表達(dá)式語言描述。以下代碼顯示了用于計算轉(zhuǎn)置矩陣乘法的張量表達(dá)式示例：

如何計算每個元素的表達(dá)式，每個計算算子都指定輸出張量的形狀和描述。張量表達(dá)式語言支持常見的算術(shù)和數(shù)學(xué)運算，涵蓋常見的DL算子模式。該語言沒有指定循環(huán)結(jié)構(gòu)和許多其它執(zhí)行細(xì)節(jié)，為各種后端添加硬件感知優(yōu)化。采用Halide的解耦計算/調(diào)度原理，從張量表達(dá)式到低級代碼的特定映射，使用調(diào)度表示。
許多可能的調(diào)度，都可以執(zhí)行此功能。
通過增量應(yīng)用基本轉(zhuǎn)換（調(diào)度原語）構(gòu)建調(diào)度，持程序的邏輯等價性。圖5顯示了在專用加速器上，調(diào)度矩陣乘法的示例。在內(nèi)部，TVM使用一個數(shù)據(jù)結(jié)構(gòu)，跟蹤循環(huán)結(jié)構(gòu)和其它信息，因為應(yīng)用了調(diào)度轉(zhuǎn)換。然后，這些信息可以為給定的最終調(diào)度生成低級代碼。

Figure 7: Performance comparison between TVM with and without cooperative shared memory fetching on matrix multiplication workloads. Tested on an NVIDIA Titan X.
張量表達(dá)式從Halide，Darkroom，TACO中獲取線索。主要增強(qiáng)功能，包括支持下面討論的新調(diào)度優(yōu)化。為了在許多后端上實現(xiàn)高性能，必須支持足夠多的調(diào)度原語，覆蓋不同硬件后端上的各種優(yōu)化。圖6總結(jié)了TVM支持的算子代碼生成過程和調(diào)度原語。重用有用的原語和來自Halide的低級循環(huán)程序AST，引入新的原語，優(yōu)化GPU和加速器性能。新原語是實現(xiàn)最佳GPU性能所必需的，也是加速器所必需的。CPU、GPU、TPU類加速器是深度學(xué)習(xí)的三種重要硬件類型。
本節(jié)介紹了CPU、GPU和TPU類加速器的新優(yōu)化原語，第5節(jié)說明了如何自動導(dǎo)出有效的調(diào)度。
4.2嵌套并行與協(xié)作
并行性是提高DL工作負(fù)載中，計算密集型內(nèi)核效率的關(guān)鍵。現(xiàn)代GPU提供了巨大的并行性，要求將所有并行模式bake到調(diào)度轉(zhuǎn)換中。大多數(shù)現(xiàn)有的解決方案，都采用了一種稱為嵌套并行（一種fork-join形式）的模型。該模型需要一個并行schedule原語，并行化數(shù)據(jù)并行任務(wù)；每個任務(wù)可以進(jìn)一步遞歸細(xì)分為子任務(wù)，利用目標(biāo)體系結(jié)構(gòu)的多級線程層次結(jié)構(gòu)（例如，GPU中的線程組）。稱此模型為無共享嵌套并行，因為一個工作線程，無法在同一并行計算階段，查看同級的數(shù)據(jù)。
無共享方法的另一種替代方法是協(xié)同獲取數(shù)據(jù)。具體來說，線程組可以協(xié)作獲取所有需要的數(shù)據(jù)，放入共享內(nèi)存空間。此優(yōu)化可以利用GPU內(nèi)存層次結(jié)構(gòu)，通過共享內(nèi)存區(qū)域跨線程重用數(shù)據(jù)。TVM支持這種眾所周知的GPU優(yōu)化，使用調(diào)度原語實現(xiàn)最佳性能。下面的GPU代碼示例優(yōu)化了矩陣乘法。

圖7展示了這種優(yōu)化的影響。將內(nèi)存作用域的概念引入調(diào)度空間，以便將計算階段（代碼中的AS和BS）標(biāo)記為共享。沒有顯式存儲范圍，自動范圍推斷，將計算階段標(biāo)記為線程本地。共享任務(wù)必須計算組中，所有工作線程的依賴關(guān)系。
必須正確插入內(nèi)存同步屏障，確保共享加載的數(shù)據(jù)對用戶可見。最后，除了對GPU有用之外，內(nèi)存作用域標(biāo)記特殊的內(nèi)存緩沖區(qū)，在針對專用DL加速器時，創(chuàng)建特殊的降低規(guī)則。
4.3 張量化
DL工作負(fù)載具有很高的運算強(qiáng)度，通常可以分解為張量算子，如矩陣乘法或一維卷積。這些自然分解導(dǎo)致了最近的添加張量計算原語。這些新原語為基于調(diào)度的編譯，帶來了機(jī)遇和挑戰(zhàn)；雖然這些使用可以提高性能，但編譯框架必須無縫集成。稱為張量化：類似于SIMD體系結(jié)構(gòu)的矢量化，但有顯著差異。指令輸入是多維的，具有固定或可變的長度，每個都有不同的數(shù)據(jù)布局。無法支持一組固定的原語，因為新的加速器用自定義的張量指令變體出現(xiàn)。
需要一個可擴(kuò)展的解決方案。
通過使用張量內(nèi)在聲明機(jī)制，將目標(biāo)硬件結(jié)構(gòu)從調(diào)度中分離出來，使張量化具有可擴(kuò)展性。使用相同的張量表達(dá)式語言，聲明每個新硬件內(nèi)在的行為，以及與之相關(guān)的降低規(guī)則。下面的代碼顯示了如何定義8×8張量硬件的內(nèi)在特性。

引入了一個tensorize調(diào)度原則，用相應(yīng)的內(nèi)部函數(shù)替換計算單元。編譯器使用硬件聲明匹配計算模式，降低到相應(yīng)的硬件本身。
Tensorization將調(diào)度與特定的硬件原語分離，易于擴(kuò)展TVM支持新的硬件體系結(jié)構(gòu)。Tensorization調(diào)度生成的代碼，符合高性能計算的實踐：將復(fù)雜算子分解為一系列微內(nèi)核調(diào)用。可以使用tensorize原語，進(jìn)行手工制作的微內(nèi)核的優(yōu)點，這在某些平臺形式中可能是有益的。例如，為移動CPU實現(xiàn)超低精度算子，這些CPU通過利用位串行數(shù)據(jù)算子1或2位寬的數(shù)據(jù)類型矩陣向量乘法微內(nèi)核。該微內(nèi)核將結(jié)果累加為越來越大的數(shù)據(jù)類型，以最小化內(nèi)存占用。將微內(nèi)核作為TVM固有的張量表示，可產(chǎn)生高達(dá)1.5倍于非張量化的加速比。

4.4顯式內(nèi)存延遲隱藏
延遲隱藏將內(nèi)存算子與計算重疊的過程，最大限度地利用內(nèi)存和計算資源。根據(jù)目標(biāo)硬件后端的不同，需要不同的策略。

Figure 8: TVM virtual thread lowering transforms a virtual thread-parallel program to a single instruction stream; the stream contains explicit low-level synchronizations that the hardware can interpret to recover the pipeline parallelism required to hide memory access latency.
在CPU上，內(nèi)存延遲隱藏通過同步多線程或硬件預(yù)取隱式實現(xiàn)的。GPU依賴于許多線程的快速上下文切換。相反，像TPU這樣的專用DL加速器，通常支持使用解耦訪問執(zhí)行（DAE）的精簡控制架構(gòu)和卸載細(xì)粒度的問題與軟件同步。

Figure 9: Decoupled Access-Execute in hardware hides most memory access latency by allowing memory and computation to overlap. Execution correctness is enforced by low-level synchronization in the form of dependence token enqueueing/dequeuing actions, which the compiler stack must insert in the instruction stream.
圖9顯示了一個減少運行時延遲的DAE硬件管道。與單片硬件設(shè)計相比，該管道可以將大部分內(nèi)存訪問隱藏在頭上，幾乎充分利用計算資源。要為了獲得更高的利用率，必須使用細(xì)粒度同步算子，擴(kuò)充指令流。如果沒有這些算子，依賴關(guān)系將無法強(qiáng)制執(zhí)行，導(dǎo)致錯誤的執(zhí)行。因此，DAE硬件管道，需要細(xì)粒度依賴排隊，在管道階段間均衡操作，保證正確執(zhí)行，如圖9的指令流所示。
編程需要顯式低級同步的DAE加速器很困難。為了減輕編程負(fù)擔(dān)，引入了虛擬線程調(diào)度原語，讓程序員指定一個高級數(shù)據(jù)并行程序，就像指定一個支持多線程的硬件后端一樣。TVM會自動降低program，轉(zhuǎn)換為具有低級顯式同步的單個指令流，如圖8所示。該算法從高級多線程程序調(diào)度開始，插入必要的低級同步算子，確保在每個線程內(nèi)正確執(zhí)行。接下來，將所有虛擬線程的算子交錯到single的指令流。

Figure 10: Rooflfline of an FPGA-based DL accelerator running ResNet inference. With latency hiding enabled by TVM, performance of the benchmarks is brought closer to the rooflfline, demonstrating higher compute and memory bandwidth effificiency.

參考文獻(xiàn)：
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

總結(jié)

以上是生活随笔為你收集整理的全文翻译（二）: TVM: An Automated End-to-End Optimizing Compiler for Deep Learning的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。