TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

提出背景

? 現(xiàn)有的 DL 框架依賴于計(jì)算圖 IR 來(lái)實(shí)現(xiàn)優(yōu)化，比如自動(dòng)微分（Auto Differentiation）和動(dòng)態(tài)內(nèi)存管理。然而計(jì)算圖層級(jí)（Graph-level）的優(yōu)化對(duì)硬件后端特定算子層級(jí)（Operator-level）的變換來(lái)說(shuō)，往往視角太高了。這些框架大部分關(guān)注的都是服務(wù)器級(jí) GPU 設(shè)備中的一小撮，而把目標(biāo)機(jī)相關(guān)的優(yōu)化委派給高度工程化的、廠家特化的算子庫(kù)。這類算子庫(kù)需要大量的人工調(diào)優(yōu)，也因此過(guò)于特殊化（普適性差）和不透明，導(dǎo)致不易于在硬件平臺(tái)之間移植。即使是框架支持的后端，優(yōu)化計(jì)算圖時(shí)也需要在以下兩個(gè)選擇中權(quán)衡：

• 避免引入算子庫(kù)中未定義的新算子（擴(kuò)展性差）
• 新算子使用未優(yōu)化的實(shí)現(xiàn)代替（性能降級(jí)）

? 為了使得各種硬件后端的計(jì)算圖層級(jí)和算子層級(jí)優(yōu)化成為可能，TVM 作為一種全新意義上的端到端（End-to-end）方法被提出。TVM 編譯器從現(xiàn)有框架中取得 DL 程序的高層級(jí)表示，并產(chǎn)生多硬件平臺(tái)后端上低層級(jí)的優(yōu)化代碼。TVM 的目標(biāo)是展示與人工調(diào)優(yōu)的競(jìng)爭(zhēng)力，關(guān)鍵的挑戰(zhàn)是：

• 平衡特定硬件的特性和抽象。DL 加速器引入了張量計(jì)算原語(yǔ)（Tensor Compute Primitives），但 CPU 和 GPU 有各自的數(shù)據(jù)處理形式（標(biāo)量、向量），如何將多維度的數(shù)據(jù)布局通過(guò)變換，使之適合處理器和內(nèi)存層級(jí)（Memory Hierarchy），是一個(gè)巨大的挑戰(zhàn)。此外，加速器設(shè)計(jì)普遍偏好控制精簡(jiǎn)性的設(shè)計(jì)，而將任務(wù)調(diào)度的復(fù)雜性下放到編譯棧上。對(duì)于某些特定的加速器，編譯器甚至需要產(chǎn)生能夠顯式解決流水線數(shù)據(jù)依賴的代碼，來(lái)“隱藏”內(nèi)存訪問(wèn)的延遲（Memory Access Latency Hiding）。
• 優(yōu)化存在海量的搜索空間。安排內(nèi)存訪問(wèn)、線程模式和新硬件原語(yǔ)（Hardware Primitives）等元素是排列組合級(jí)別的復(fù)雜度，如果實(shí)現(xiàn)一個(gè)黑箱來(lái)自動(dòng)調(diào)優(yōu)會(huì)帶來(lái)巨大的搜索開(kāi)銷。也許我們可以預(yù)定義一個(gè)開(kāi)銷模型來(lái)指導(dǎo)搜索，但是現(xiàn)代硬件的復(fù)雜程度是不斷增長(zhǎng)的（而且速度很快），精確建模的難度非常大，更不用說(shuō)還得給每一種硬件類型都獨(dú)立建模了。

主要工作

? TVM 提出了三個(gè)重要模塊：

• 張量表達(dá)語(yǔ)言（Tensor Expression Language）和變換原語(yǔ)（Transformation Primitives）。這是對(duì) Halide 計(jì)算調(diào)度解耦理念的擴(kuò)展，把硬件本質(zhì)（Hardware Intrinsic）和變換原語(yǔ)也分離了，使得 TVM 可以支持新的加速器和對(duì)應(yīng)的硬件本質(zhì)。
• 自動(dòng)程序優(yōu)化框架（Automated Program Optimization Framework）。使用一個(gè)基于 ML 的開(kāi)銷模型來(lái)指導(dǎo)如何尋找最優(yōu)張量算子，當(dāng)收集到更多硬件后端的數(shù)據(jù)時(shí)，模型的適應(yīng)性和表現(xiàn)會(huì)不斷提升。
• 計(jì)算圖重寫（Graph Rewriter）。作用于自動(dòng)代碼生成器上層，統(tǒng)籌計(jì)算圖級(jí)和算子級(jí)的聯(lián)合優(yōu)化。

工作流程

? TVM 的工作流程：

讀取現(xiàn)有框架中的模型，產(chǎn)生計(jì)算圖表示

進(jìn)行高層級(jí)數(shù)據(jù)流重寫（High-level Dataflow Rewriting），產(chǎn)生優(yōu)化后的計(jì)算圖

進(jìn)行算子層級(jí)優(yōu)化，為計(jì)算圖中的融合算子（Fused Operator）產(chǎn)生高效代碼

算子由張量表達(dá)語(yǔ)言聲明，執(zhí)行細(xì)節(jié)不需要指定

利用 ML 開(kāi)銷模型從目標(biāo)機(jī)潛在的優(yōu)化集合中，搜索算子的最優(yōu)代碼

將生成的代碼包裝為可部署的模塊

優(yōu)化計(jì)算圖（Optimizing Computational Graphs）

? 計(jì)算圖是一種高級(jí)表示，提供了對(duì)于算子的全局視野，而不需要指明實(shí)現(xiàn)的細(xì)節(jié)。就像 LLVM IR，計(jì)算圖也可以轉(zhuǎn)換成功能等價(jià)的子圖，以適合各種優(yōu)化手段。TVM 對(duì)計(jì)算圖的優(yōu)化包括：算子融合（Operator Fusion）、常量折疊（Constant Folding）、靜態(tài)內(nèi)存規(guī)劃（Static Memory Planning）、數(shù)據(jù)布局變換（Data Layout Transformation）等。

算子融合

? 算子融合應(yīng)用了多個(gè)算子、一次計(jì)算的理念，避免保存中間結(jié)果時(shí)對(duì)內(nèi)存的訪問(wèn)，從而減少執(zhí)行時(shí)間。TVM 將算子分為以下幾類：

• 單射（Injective），一對(duì)一的運(yùn)算，例如元素和
• 規(guī)約（Reduction），多對(duì)一的運(yùn)算，例如求和
• 復(fù)雜運(yùn)算、輸出可融合（Complex-out-fusable），可以將元素映射融合到輸出，例如二維卷積
• 不透明（Opaque），不可融合，例如排序

? 針對(duì)這幾類算子，TVM 提出了泛用的融合規(guī)則：

• 多個(gè)單射算子可以融合為單個(gè)單射算子
• 一個(gè)規(guī)約算子可以和多個(gè)單射算子融合，例如縮放后求和
• 逐元素（Element-wise）類型的算子可以融合到復(fù)雜運(yùn)算、輸出可融合算子的輸出

? 算子融合通過(guò)減少訪存可以實(shí)現(xiàn) 1.2× 到 2× 的加速。

數(shù)據(jù)布局變換

? 最常見(jiàn)的數(shù)據(jù)布局是列優(yōu)先（Column Major）和行優(yōu)先（Row Major）的方式。事實(shí)上，DL 加速器往往會(huì)采用更復(fù)雜的數(shù)據(jù)布局，比如 4×4 的矩陣，為了能夠充分利用空間局部性，要求數(shù)據(jù)能夠平鋪成 4×4 的“小磚塊”。數(shù)據(jù)布局變換將計(jì)算圖轉(zhuǎn)換成可以更好地利用內(nèi)部數(shù)據(jù)布局的形式，首先需要為每個(gè)算子規(guī)定來(lái)自內(nèi)存層級(jí)的約束，如果數(shù)據(jù)布局不符合要求，就進(jìn)行變換，這里采用的是生產(chǎn)者消費(fèi)者模式。

生成張量運(yùn)算（Generating Tensor Operations）

? 雖然計(jì)算圖優(yōu)化能極大地提高 DL 工作負(fù)載，但是它的效果與算子庫(kù)提供的算子水平有很大關(guān)系。現(xiàn)在支持算子融合的 DL 框架很少有要求算子庫(kù)也提供算子融合模式的實(shí)現(xiàn)，因?yàn)殡S著神經(jīng)網(wǎng)絡(luò)算子的不斷提出，融合算子的數(shù)量也經(jīng)歷了排列組合級(jí)別的增長(zhǎng)，再考慮到各種不同硬件后端的出現(xiàn)，這種方式明顯是不可持續(xù)的。出于同樣的原因，理想的、多樣的算子也不可能經(jīng)由手工調(diào)制，于是張量算子的自動(dòng)生成就成了迫切需要。

張量表達(dá)式（Tensor Expression）和調(diào)度空間（Schedule Space）

? 張量表達(dá)式由結(jié)果形狀和運(yùn)算規(guī)則兩部分組成，支持常見(jiàn)的算數(shù)運(yùn)算和 DL 算子。它無(wú)需指明循環(huán)結(jié)構(gòu)和其他執(zhí)行細(xì)節(jié)，提供給硬件后端優(yōu)化更大的靈活性。

? 在維持程序邏輯等價(jià)性的前提下，TVM 對(duì)張量表達(dá)式逐次使用基本變換（調(diào)度原語(yǔ)），并記錄下過(guò)程中的循環(huán)結(jié)構(gòu)等其他所需信息，這些信息用來(lái)幫助生成最終調(diào)度（Final Schedule）的低層級(jí)代碼。

協(xié)作嵌套并行（Nested Parallelism with Cooperation）

? 嵌套并行是 Fork-join 模型的一種形式，指的是每一個(gè)子任務(wù)都可以遞歸地被更進(jìn)一步劃分成子任務(wù)并行處理，從而深度利用目標(biāo)架構(gòu)的多級(jí)線程層級(jí)（Multi-level Thread Hierarchy），比如 GPU 的線程組（Thread Group）。如果在并行計(jì)算階段中，一個(gè)線程無(wú)需訪問(wèn)相鄰線程的數(shù)據(jù)，這種模型又被叫做無(wú)共享嵌套并行（Shared-nothing Nested Parallelism）。

? 代替無(wú)共享方式的一種選擇是協(xié)作獲取數(shù)據(jù)：一組線程共同獲取到一塊數(shù)據(jù)，然后各取所需。這樣做的好處是能夠充分利用 GPU 顯存層級(jí)，通過(guò)共享內(nèi)存也使得線程之間的數(shù)據(jù)重用成為可能。

? TVM 在調(diào)度空間種引入了內(nèi)存作用域（Memory Scope）的概念，計(jì)算階段（Compute Stage）可以被標(biāo)記為共享（Shared）。如果沒(méi)有顯式的內(nèi)存作用域，自動(dòng)作用域推導(dǎo)會(huì)把計(jì)算階段標(biāo)記為線程局部（Thread-local）。共享任務(wù)必須計(jì)算彼此之間的依賴關(guān)系，內(nèi)存同步屏障（Memory Synchronization Barrier）技術(shù)也需要用來(lái)保證數(shù)據(jù)對(duì)數(shù)據(jù)的消費(fèi)者是可見(jiàn)的。另外，內(nèi)存作用域還可以標(biāo)記特殊內(nèi)存緩存，這對(duì) GPU 來(lái)說(shuō)很有用；當(dāng)以 DL 加速器為目標(biāo)機(jī)時(shí)，內(nèi)存作用域還可以創(chuàng)建額外的代碼低層級(jí)化規(guī)則。

張量化（Tensorization）

? 類比向量化（Vectorization）之于 SIMD。

顯式內(nèi)存延遲隱藏（Explicit Memory Latency Hiding）

? CPU 隱藏內(nèi)存延遲的方式是多線程，GPU 隱藏內(nèi)存延遲的方式是線程組的快速上下文切換，但是特定 DL 加速器（比如 TPU）偏好控制精簡(jiǎn)型的解耦訪問(wèn)執(zhí)行（Decoupled Access Execute）架構(gòu)，會(huì)將細(xì)粒度的同步控制下放到軟件處理。

? DAE 架構(gòu)流水線需要保證正確的依賴關(guān)系，可以通過(guò)使用細(xì)粒度的依賴隊(duì)列實(shí)現(xiàn)。直接在低層級(jí)上實(shí)現(xiàn) DAE 加速器的同步控制是相對(duì)困難的，TVM 引入了虛擬線程調(diào)度原語(yǔ)（Virtual Threading Scheduling Primitive），開(kāi)發(fā)者可以假裝指定的硬件后端擁有多線程支持，TVM 來(lái)負(fù)責(zé)插入確保執(zhí)行順序所必須的低層級(jí)同步操作，并自動(dòng)生成單指令流。

自動(dòng)優(yōu)化（Automating Optimization）

? TVM 為 DL 模型的每一層產(chǎn)生針對(duì)輸入形狀和布局優(yōu)化過(guò)的算子，從而帶來(lái)巨大的性能增益，但如何選擇調(diào)度優(yōu)化（比如改變循環(huán)順序、平鋪大小、展開(kāi)因子等）卻是排列組合級(jí)別的復(fù)雜度。為此，TVM 提出了自動(dòng)調(diào)度優(yōu)化器（Automated Schedule Optimizer），它包含兩個(gè)主要組件：

• 調(diào)度探索器（Schedule Explorer），用來(lái)提出潛在的、有前途的優(yōu)化配置
• 基于機(jī)器學(xué)習(xí)的開(kāi)銷模型（ML-based Cost Model），用來(lái)預(yù)測(cè)和評(píng)估給定配置的表現(xiàn)

? TVM 提出了調(diào)度模板規(guī)范（Schedule Template Specification）API，使開(kāi)發(fā)者可以在調(diào)度空間中定義錨點(diǎn)，包括一些額外的特定領(lǐng)域的背景知識(shí)。TVM 為各種硬件后端創(chuàng)建了通用主模板（Generic Master Template），用來(lái)從張量表達(dá)語(yǔ)言中自動(dòng)提取可能的錨點(diǎn)。

基于機(jī)器學(xué)習(xí)的開(kāi)銷模型

? 比較而言，黑箱自動(dòng)調(diào)優(yōu)（Blackbox Auto-tuning）通常用來(lái)調(diào)優(yōu)高性能計(jì)算的運(yùn)行庫(kù)，但是為了取得較好的結(jié)果需要大量的實(shí)驗(yàn)。另一種方法是預(yù)定義開(kāi)銷模型（Predefined Cost Model），理想情況下，它應(yīng)該能夠綜合考慮內(nèi)存訪問(wèn)模式、數(shù)據(jù)重用、流水線依賴、線程模式等各種因素，然而不幸的是，為當(dāng)今愈加復(fù)雜的硬件架構(gòu)創(chuàng)建預(yù)定義開(kāi)銷模型困難重重。

? TVM 采用了數(shù)據(jù)驅(qū)動(dòng)的方法，ML 模型將低層級(jí)的循環(huán)程序作為輸入，預(yù)測(cè)它在指定硬件后端上的運(yùn)行時(shí)長(zhǎng)。模型使用探索過(guò)程中的運(yùn)行時(shí)刻測(cè)量數(shù)據(jù)來(lái)訓(xùn)練，不需要用戶輸入硬件細(xì)節(jié)信息，模型的準(zhǔn)確率會(huì)隨著試驗(yàn)次數(shù)的增加而改進(jìn)，它是對(duì)前兩種方法的折衷。

? 模型的選擇上面，質(zhì)量和速度是關(guān)鍵考量。調(diào)度探索器會(huì)頻繁地對(duì)開(kāi)銷模型發(fā)起查詢請(qǐng)求，使調(diào)度優(yōu)化過(guò)程中引入了模型預(yù)測(cè)和模型更新的時(shí)間花費(fèi)，這類時(shí)間花費(fèi)在真實(shí)硬件上應(yīng)當(dāng)被嚴(yán)格限制。因此與傳統(tǒng)的超參數(shù)調(diào)優(yōu)不同，大模型在調(diào)度優(yōu)化里不是一種好的選擇。目標(biāo)函數(shù)或者損失函數(shù)可以采用預(yù)測(cè)運(yùn)行時(shí)長(zhǎng)與真實(shí)運(yùn)行時(shí)長(zhǎng)的偏差，由于調(diào)度探索器只會(huì)從最優(yōu)秀的候選項(xiàng)里選擇，因此事實(shí)上不必預(yù)測(cè)運(yùn)行時(shí)長(zhǎng)的絕對(duì)大小，TVM 讓目標(biāo)函數(shù)支持排名作為替代。

? TVM 采用了基于 XGBoost 的梯度樹(shù)狀提升模型（Gradient Tree Boosting Model），從循環(huán)程序提取的內(nèi)存訪問(wèn)計(jì)數(shù)、每級(jí)循環(huán)緩存重用率、循環(huán) One-hot 向量表示等特征預(yù)測(cè)運(yùn)行時(shí)長(zhǎng)；另一種基于神經(jīng)網(wǎng)絡(luò)的 TreeRNN 模型則是從循環(huán)程序的 AST 中提取特征，不需要自行構(gòu)造特征。兩者預(yù)測(cè)質(zhì)量接近，但前者訓(xùn)練和推導(dǎo)都更快。

調(diào)度探索（Schedule Exploration）

? 調(diào)度探索器最簡(jiǎn)單的策略就是讓每一個(gè)配置都跑一邊開(kāi)銷模型，然后選擇前幾個(gè)預(yù)測(cè)表現(xiàn)好的，問(wèn)題是搜索空間大起來(lái)時(shí)間花銷就不能接受了。TVM 采用的是模擬退火算法，從一個(gè)隨機(jī)配置開(kāi)始，每次在鄰近配置中隨機(jī)游走，如果開(kāi)銷降低了就接收下一個(gè)配置，否則以一定概率拒絕該配置。

分布式設(shè)備池（Distributed Device Pool）和遠(yuǎn)過(guò)程調(diào)用（Remote Process Call）

? 分布式設(shè)備池使得模型在硬件上的試驗(yàn)次數(shù)大大增加，同時(shí)多個(gè)優(yōu)化任務(wù)之間能夠進(jìn)行細(xì)粒度的資源共享。TVM 的分布式設(shè)備池基于 RPC 技術(shù)的、可定制化的，支持動(dòng)態(tài)裝載和運(yùn)行交叉編譯得到的模塊。這樣一套相同的基礎(chǔ)設(shè)施可以進(jìn)行單工作負(fù)載的優(yōu)化和端到端的圖推導(dǎo)任務(wù)。

參考資料

T. Chen, T. Moreau, Z. Jiang, et al. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI’18. https://arxiv.org/abs/1802.04799

總結(jié)

以上是生活随笔為你收集整理的TVM：一种自动端到端优化的深度学习编译器的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

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