TVM：一种自动端到端优化的深度学习编译器
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
提出背景
? 現有的 DL 框架依賴于計算圖 IR 來實現優化，比如自動微分（Auto Differentiation）和動態內存管理。然而計算圖層級（Graph-level）的優化對硬件后端特定算子層級（Operator-level）的變換來說，往往視角太高了。這些框架大部分關注的都是服務器級 GPU 設備中的一小撮，而把目標機相關的優化委派給高度工程化的、廠家特化的算子庫。這類算子庫需要大量的人工調優，也因此過于特殊化（普適性差）和不透明，導致不易于在硬件平臺之間移植。即使是框架支持的后端，優化計算圖時也需要在以下兩個選擇中權衡：
- 避免引入算子庫中未定義的新算子（擴展性差）
- 新算子使用未優化的實現代替（性能降級）
? 為了使得各種硬件后端的計算圖層級和算子層級優化成為可能，TVM 作為一種全新意義上的端到端（End-to-end）方法被提出。TVM 編譯器從現有框架中取得 DL 程序的高層級表示，并產生多硬件平臺后端上低層級的優化代碼。TVM 的目標是展示與人工調優的競爭力，關鍵的挑戰是：
- 平衡特定硬件的特性和抽象。DL 加速器引入了張量計算原語（Tensor Compute Primitives），但 CPU 和 GPU 有各自的數據處理形式（標量、向量），如何將多維度的數據布局通過變換，使之適合處理器和內存層級（Memory Hierarchy），是一個巨大的挑戰。此外，加速器設計普遍偏好控制精簡性的設計，而將任務調度的復雜性下放到編譯棧上。對于某些特定的加速器，編譯器甚至需要產生能夠顯式解決流水線數據依賴的代碼，來“隱藏”內存訪問的延遲（Memory Access Latency Hiding）。
- 優化存在海量的搜索空間。安排內存訪問、線程模式和新硬件原語（Hardware Primitives）等元素是排列組合級別的復雜度，如果實現一個黑箱來自動調優會帶來巨大的搜索開銷。也許我們可以預定義一個開銷模型來指導搜索，但是現代硬件的復雜程度是不斷增長的（而且速度很快），精確建模的難度非常大，更不用說還得給每一種硬件類型都獨立建模了。
主要工作
? TVM 提出了三個重要模塊：
- 張量表達語言（Tensor Expression Language）和變換原語（Transformation Primitives）。這是對 Halide 計算調度解耦理念的擴展，把硬件本質（Hardware Intrinsic）和變換原語也分離了，使得 TVM 可以支持新的加速器和對應的硬件本質。
- 自動程序優化框架（Automated Program Optimization Framework）。使用一個基于 ML 的開銷模型來指導如何尋找最優張量算子，當收集到更多硬件后端的數據時，模型的適應性和表現會不斷提升。
- 計算圖重寫（Graph Rewriter）。作用于自動代碼生成器上層，統籌計算圖級和算子級的聯合優化。
工作流程
? TVM 的工作流程：
優化計算圖（Optimizing Computational Graphs）
? 計算圖是一種高級表示，提供了對于算子的全局視野，而不需要指明實現的細節。就像 LLVM IR，計算圖也可以轉換成功能等價的子圖，以適合各種優化手段。TVM 對計算圖的優化包括：算子融合（Operator Fusion）、常量折疊（Constant Folding）、靜態內存規劃（Static Memory Planning）、數據布局變換（Data Layout Transformation）等。
算子融合
? 算子融合應用了多個算子、一次計算的理念，避免保存中間結果時對內存的訪問，從而減少執行時間。TVM 將算子分為以下幾類：
- 單射（Injective），一對一的運算，例如元素和
- 規約（Reduction），多對一的運算，例如求和
- 復雜運算、輸出可融合（Complex-out-fusable），可以將元素映射融合到輸出，例如二維卷積
- 不透明（Opaque），不可融合，例如排序
? 針對這幾類算子，TVM 提出了泛用的融合規則：
- 多個單射算子可以融合為單個單射算子
- 一個規約算子可以和多個單射算子融合，例如縮放后求和
- 逐元素（Element-wise）類型的算子可以融合到復雜運算、輸出可融合算子的輸出
? 算子融合通過減少訪存可以實現 1.2× 到 2× 的加速。
數據布局變換
? 最常見的數據布局是列優先（Column Major）和行優先（Row Major）的方式。事實上，DL 加速器往往會采用更復雜的數據布局，比如 4×4 的矩陣，為了能夠充分利用空間局部性，要求數據能夠平鋪成 4×4 的“小磚塊”。數據布局變換將計算圖轉換成可以更好地利用內部數據布局的形式，首先需要為每個算子規定來自內存層級的約束，如果數據布局不符合要求，就進行變換，這里采用的是生產者消費者模式。
生成張量運算（Generating Tensor Operations）
? 雖然計算圖優化能極大地提高 DL 工作負載，但是它的效果與算子庫提供的算子水平有很大關系。現在支持算子融合的 DL 框架很少有要求算子庫也提供算子融合模式的實現，因為隨著神經網絡算子的不斷提出，融合算子的數量也經歷了排列組合級別的增長，再考慮到各種不同硬件后端的出現，這種方式明顯是不可持續的。出于同樣的原因，理想的、多樣的算子也不可能經由手工調制，于是張量算子的自動生成就成了迫切需要。
張量表達式（Tensor Expression）和調度空間（Schedule Space）
? 張量表達式由結果形狀和運算規則兩部分組成，支持常見的算數運算和 DL 算子。它無需指明循環結構和其他執行細節，提供給硬件后端優化更大的靈活性。
? 在維持程序邏輯等價性的前提下，TVM 對張量表達式逐次使用基本變換（調度原語），并記錄下過程中的循環結構等其他所需信息，這些信息用來幫助生成最終調度（Final Schedule）的低層級代碼。
協作嵌套并行（Nested Parallelism with Cooperation）
? 嵌套并行是 Fork-join 模型的一種形式，指的是每一個子任務都可以遞歸地被更進一步劃分成子任務并行處理，從而深度利用目標架構的多級線程層級（Multi-level Thread Hierarchy），比如 GPU 的線程組（Thread Group）。如果在并行計算階段中，一個線程無需訪問相鄰線程的數據，這種模型又被叫做無共享嵌套并行（Shared-nothing Nested Parallelism）。
? 代替無共享方式的一種選擇是協作獲取數據：一組線程共同獲取到一塊數據，然后各取所需。這樣做的好處是能夠充分利用 GPU 顯存層級，通過共享內存也使得線程之間的數據重用成為可能。
? TVM 在調度空間種引入了內存作用域（Memory Scope）的概念，計算階段（Compute Stage）可以被標記為共享（Shared）。如果沒有顯式的內存作用域，自動作用域推導會把計算階段標記為線程局部（Thread-local）。共享任務必須計算彼此之間的依賴關系，內存同步屏障（Memory Synchronization Barrier）技術也需要用來保證數據對數據的消費者是可見的。另外，內存作用域還可以標記特殊內存緩存，這對 GPU 來說很有用；當以 DL 加速器為目標機時，內存作用域還可以創建額外的代碼低層級化規則。
張量化（Tensorization）
? 類比向量化（Vectorization）之于 SIMD。
顯式內存延遲隱藏（Explicit Memory Latency Hiding）
? CPU 隱藏內存延遲的方式是多線程，GPU 隱藏內存延遲的方式是線程組的快速上下文切換，但是特定 DL 加速器（比如 TPU）偏好控制精簡型的解耦訪問執行（Decoupled Access Execute）架構，會將細粒度的同步控制下放到軟件處理。
? DAE 架構流水線需要保證正確的依賴關系，可以通過使用細粒度的依賴隊列實現。直接在低層級上實現 DAE 加速器的同步控制是相對困難的，TVM 引入了虛擬線程調度原語（Virtual Threading Scheduling Primitive），開發者可以假裝指定的硬件后端擁有多線程支持，TVM 來負責插入確保執行順序所必須的低層級同步操作，并自動生成單指令流。
自動優化（Automating Optimization）
? TVM 為 DL 模型的每一層產生針對輸入形狀和布局優化過的算子，從而帶來巨大的性能增益，但如何選擇調度優化（比如改變循環順序、平鋪大小、展開因子等）卻是排列組合級別的復雜度。為此，TVM 提出了自動調度優化器（Automated Schedule Optimizer），它包含兩個主要組件：
- 調度探索器（Schedule Explorer），用來提出潛在的、有前途的優化配置
- 基于機器學習的開銷模型（ML-based Cost Model），用來預測和評估給定配置的表現
? TVM 提出了調度模板規范（Schedule Template Specification）API，使開發者可以在調度空間中定義錨點，包括一些額外的特定領域的背景知識。TVM 為各種硬件后端創建了通用主模板（Generic Master Template），用來從張量表達語言中自動提取可能的錨點。
基于機器學習的開銷模型
? 比較而言，黑箱自動調優（Blackbox Auto-tuning）通常用來調優高性能計算的運行庫，但是為了取得較好的結果需要大量的實驗。另一種方法是預定義開銷模型（Predefined Cost Model），理想情況下，它應該能夠綜合考慮內存訪問模式、數據重用、流水線依賴、線程模式等各種因素，然而不幸的是，為當今愈加復雜的硬件架構創建預定義開銷模型困難重重。
? TVM 采用了數據驅動的方法，ML 模型將低層級的循環程序作為輸入，預測它在指定硬件后端上的運行時長。模型使用探索過程中的運行時刻測量數據來訓練，不需要用戶輸入硬件細節信息，模型的準確率會隨著試驗次數的增加而改進，它是對前兩種方法的折衷。
? 模型的選擇上面，質量和速度是關鍵考量。調度探索器會頻繁地對開銷模型發起查詢請求，使調度優化過程中引入了模型預測和模型更新的時間花費，這類時間花費在真實硬件上應當被嚴格限制。因此與傳統的超參數調優不同，大模型在調度優化里不是一種好的選擇。目標函數或者損失函數可以采用預測運行時長與真實運行時長的偏差，由于調度探索器只會從最優秀的候選項里選擇，因此事實上不必預測運行時長的絕對大小，TVM 讓目標函數支持排名作為替代。
? TVM 采用了基于 XGBoost 的梯度樹狀提升模型（Gradient Tree Boosting Model），從循環程序提取的內存訪問計數、每級循環緩存重用率、循環 One-hot 向量表示等特征預測運行時長；另一種基于神經網絡的 TreeRNN 模型則是從循環程序的 AST 中提取特征，不需要自行構造特征。兩者預測質量接近，但前者訓練和推導都更快。
調度探索（Schedule Exploration）
? 調度探索器最簡單的策略就是讓每一個配置都跑一邊開銷模型，然后選擇前幾個預測表現好的，問題是搜索空間大起來時間花銷就不能接受了。TVM 采用的是模擬退火算法，從一個隨機配置開始，每次在鄰近配置中隨機游走，如果開銷降低了就接收下一個配置，否則以一定概率拒絕該配置。
分布式設備池（Distributed Device Pool）和遠過程調用（Remote Process Call）
? 分布式設備池使得模型在硬件上的試驗次數大大增加，同時多個優化任務之間能夠進行細粒度的資源共享。TVM 的分布式設備池基于 RPC 技術的、可定制化的，支持動態裝載和運行交叉編譯得到的模塊。這樣一套相同的基礎設施可以進行單工作負載的優化和端到端的圖推導任務。
參考資料
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
總結
以上是生活随笔為你收集整理的TVM：一种自动端到端优化的深度学习编译器的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 无尘车间净化装修方案
- 下一篇： Python 数据库开发实战-Mac系统