天天看點

TVM | 一種用于深度學習的端到端自動優化編譯器

  ​

​歡迎關注我的公衆号 [極智視界],擷取我的更多筆記分享​

  大家好,我是極智視界,本文解讀一下 一種用于深度學習的端到端自動優化編譯器 TVM。

TVM | 一種用于深度學習的端到端自動優化編譯器

  現在越來越需要将機器學習部署到各種硬體裝置,目前的架構依賴于供應商特定的算子,并隻是針對範圍較窄的伺服器級的 GPU 進行了優化。将工作負載部署到新平台:例如手機、嵌入式裝置和加速器 (例如 FPGA、ASIC 等),需要大量的手工工作。作者提出了 TVM,這是一種公開了 graph級别 和 算子級别的優化,在跨不同硬體後端的深度學習工作負載的情況下也能表現優秀的性能。TVM 解決了深度學習特有的優化挑戰,例如進階算子融合、映射到任意硬體原語 和 記憶體延遲隐藏。它還通過采用一種新穎的、基于學習的成本模組化方法來快速搜尋代碼優化,進而自動優化低級程式以适應硬體特性。實驗結果表明,TVM 提供的跨硬體後端性能可與用于低功耗 CPU、移動 GPU 和 伺服器級GPU 的最先進的手動調整庫相媲美。另外作者還展示了 TVM 針對新加速器後端的适應能力,例如基于 FPGA 的通用深度學習加速器。值得一提的是,TVM 是開源的,并已經在幾家大廠内用于生産。

  論文位址:​​https://www.usenix.org/system/files/osdi18-chen.pdf​​

  github位址:​​https://github.com/apache/tvm​​

文章目錄

  • ​​1. 簡介​​
  • ​​2. 總覽​​
  • ​​3. 優化計算圖​​
  • ​​4. 生成張量算子​​
  • ​​4.1 張量表達和排程空間​​
  • ​​4.2 嵌套并行與協作​​
  • ​​4.3 張量化​​
  • ​​4.4 顯式記憶體延遲隐藏​​
  • ​​5. 自動優化​​
  • ​​5.1 排程空間規範​​
  • ​​5.2 基于機器學習的成本模型​​
  • ​​5.3 排程搜尋​​
  • ​​5.4 分布式裝置池和 RPC​​
  • ​​6 實驗​​
  • ​​7 總結​​
  • ​​8 參考​​

1. 簡介

  深度學習模型 (DL) 現在可以用來識别圖像、處理自然語言,并在具有挑戰性的政策遊戲中擊敗人類。從雲伺服器到自動駕駛汽車和嵌入式裝置,将智能應用程式部署到各種裝置的需求在不斷增長。由于硬體特性的多樣性,包括嵌入式 CPU、GPU、FPGA 和 ASIC (例如 TPU),将 DL 工作負載映射到這些裝置變得複雜。這些硬體目标在記憶體組織、計算功能單元等方面存在差異,如圖1所示。

TVM | 一種用于深度學習的端到端自動優化編譯器

  目前的 DL 架構,如 TensorFlow、MXNet、Caffe 和 PyTorch,依靠計算圖中間表示來實作優化,例如自動微分 和 動态記憶體管理。然而,圖級優化通常過于進階,無法處理特定于硬體後端的算子級别的轉換。這些架構中的大多數都專注于一小類伺服器級 GPU 裝置,針對特定目标的優化依賴于高度工程化和特定于供應商的算子庫。這些算子級别的庫需要大量的手動調整,是以過于專業且不透明,無法輕松跨硬體裝置移植。目前,在各種 DL 架構中為各種硬體後端提供支援需要大量的工程努力。即使對于已經支援的後端,架構也必須在以下之間做出艱難的選擇:(1) 避免産生不在預定義算子庫中的新算子的圖優化;(2) 使用這些新算子的未優化的視線。為了對不同的硬體後端啟用圖級 和 算子級的優化,作者采用了一種完全不同的端到端的方法。作者建構了 TVM,這是一個編譯器,它從現有架構中擷取深度學習程式的進階規範,并為各種硬體後端生成低級優化代碼。為了吸引更多使用者,TVM 需要提供與跨不同硬體後端的大量手動優化的算子庫相比 具有競争力的性能。這一目标需要解決下面描述的關鍵挑戰。

  Levaraging specific Hardware Features and Abstractions DL 加速器引入了優化的張量加速原語,而 GPU 和 CPU 則不斷改進其處理元素。這對為給定的算子描述生成優化代碼提出了重大挑戰。硬體指令的輸入是多元的,具有固定或可變長度;他們決定了不同的資料布局;他們對記憶體層次結構有特殊要求。系統必須有效地利用這些複雜的原語才能從加速中受益。此外,加速器設計通常也支援更加精簡的控制,并将大多數排程複雜性下放到編碼器堆棧。對于專用加速器,系統需要生成顯式控制管道依賴關系的代碼,以隐藏記憶體通路延遲 - 這是硬體為 CPU 和 GPU 執行的一項工作。

  Large Search Space for Optimization 另一個挑戰是在不手動調整算子的情況下生成高效的代碼。記憶體通路、線程模式和新穎的硬體原語的組合選擇為生成的代碼 (例如 循環塊、排序、緩存 和 展開) 建立了巨大的配置空間,如果要實作黑盒自動調整,将需要大量的搜尋成本。可以采用預定義的成本模型來指導搜尋,但由于現代硬體的複雜性日益增加,建構準确的成本模型很困難。此外,這種方法需要我們為每種硬體類型建立單獨的成本模型。

  TVM 通過三個關鍵子產品解決了這些挑戰。(1) 引入張量表達式語言來建構算子并提供程式轉換原語,這些原語可以生成具有各種優化的不同版本的程式。這個 layer 擴充了 Halide 的計算 / 排程分離概念,還将目标硬體内在函數與轉換原語分離,進而支援新的加速器 及其 相應的新内在函數。此外,引入了新的轉換原語來解決與 GPU 相關的挑戰,并支援部署到專門的加速器。然後應用不同的程式轉換序列,為給定的算子聲明形成一個豐富的有效程式空間。(2) 引入了一個自動程式優化架構來尋找優化的張量算子。優化器由基于 ML 的成本模型指導,随着我們從硬體後端收集更多的資料,該模型會适應和改進。(3) 在自動代碼生成器之上,引入了一個 graph重寫器,它充分利用了進階 和 算子級别的優化。

  通過結合這三個子產品,TVM 可以從現有的深度學習架構中擷取模型描述,執行進階和低級聯合優化,并為後端生成特定于硬體的優化代碼,例如 CPU、GPU 和 基于FPGA的專用加速器。本文的貢獻如下:

  • 确定了為跨不同硬體後端的深度學習工作負載提供性能可移植性方面的主要優化挑戰;
  • 介紹了利用跨線程記憶體重用、新穎的硬體内在函數 和 延遲隐藏的新穎排程原語;
  • 提出并實作了一個基于機器學習的優化系統,以自動搜尋和搜尋優化的張量算子;
  • 建構了一個端到端的編譯和優化堆棧,允許将進階架構 (TensorFlow、MXNet、PyTorch、Kears、CNTK) 中指定的深度學習工作負載部署到不同的硬體後端 (包括 CPU、伺服器 GPU、移動 GPU 和 基于FPGA的加速器);
  • 開源 TVM 在幾家大公司内用于生産;

  使用伺服器級GPU、嵌入式GPU、嵌入式CPU 和 基于FPGA 的定制通用加速器上的真實工作負載評估了 TVM。實驗結果表明,TVM 提供了跨後端的可移植性能,并且比由手動優化庫支援的現有架構實作了 1.2 倍到 3.8 倍的加速。

2. 總覽

  本節通過一個示例來介紹 TVM 的元件。圖2總結了 TVM 中的執行步驟及其在論文中的相應部分。

TVM | 一種用于深度學習的端到端自動優化編譯器

  該系統首先将來自現有架構的模型作為輸入,并将其轉換為計算圖表示。然後它執行進階資料流重寫以生成優化圖。算子級優化子產品必須為該圖中的每個融合算子生成高效代碼。算子以聲明性張量表達式語言指定;執行細節未指定。TVM 為給定硬體目标的算子識别一組可能的代碼優化。可能的優化形成了很大的空間,是以使用基于 ML 的成本模型來尋找優化的算子。最後,系統将生成的代碼打包到一個可部署的子產品中。

  End-User Example 隻需要幾行代碼,使用者就可以從現有的深度學習架構中擷取模型并調用 TVM API 以擷取可部署的子產品:

import tvm as t
# Use keras framework as example, import model
graph, params = t.frontend.from_keras(keras_model)
target = t.target.cuda()
graph, lib, params = t.compiler.build(graph, target, params)      

  這個編譯後的運作時子產品包含三個元件:最終優化的計算圖 (graph)、生成的算子 (lib) 和 子產品參數 (params)。然後可以使用這些元件将模型部署到目标後端:

import tvm.runtime as t

module = runtime.create(graph, lib, t.cuda(0))
module.set_input(**params)
module.run(data=data_array)
output = tvm.nd.empty(out_shape, ctx=t.cuda(0))
module.get_output(0, output)      

  TVM 支援多種語言的部署後端,例如 C++、Java 和 Python。後面的其餘部分描述了 TVM 的架構以及系統程式員如何擴充它以支援新的後端。

3. 優化計算圖

  計算圖是在 DL 架構中表示程式的常用方法。圖3展示了一個兩層卷積網絡的示例計算圖表示。

TVM | 一種用于深度學習的端到端自動優化編譯器

  這種進階表示與低級編譯器中間表示 (IR) (如 LLVM) 之間的主要差別在于:中間資料項是大型的多元張量。計算圖提供了算子的全局視圖,但它們避免指定每個算子必須如何實作。與 LLVM IR 一樣,計算圖可以轉換為功能等效的圖以應用優化。作者還利用常見 DL 工作負載中的形狀特異性來優化一組固定的輸入形狀。

  TVM 利用計算圖表示來應用進階優化:節點表示對張量或程式輸入的操作,邊表示操作之間的資料依賴關系。它實作了許多圖級優化,包括:算子融合,将多個小操作融合在一起;常量折疊,預先計算可以靜态确定 graph 部分,節省執行成本;一個靜态記憶體規劃 pass,它預先配置設定記憶體來儲存每個中間張量;資料排布轉換,将内部資料排布轉換為後端友好的形式。接下來讨論算子融合和資料排布轉換。

  Operator Fusion 算子融合将多個算子組合到一個核心中,而不會将中間結果儲存在記憶體中。這種優化可以大大減少執行時間,尤其是在 GPU 和專用加速器中的執行時間。具體來說,有四類 graph 算子:(1) injective (一對一映射,例如 add);(2) reduction (例如 sum);(3) complex-out-fusable (可以将元素映射融合到輸出,例如 conv2d);(4) opaque (不能融合的,例如 sort)。這裡提供了融合這些算子的通用規則,如下所示。多個 injective 算子可以融合成另一個 injective 算子。reduction 算子可以與輸入 injective 算子融合 (例如 融合 scale 和 sum)。諸如 conv2d 之類的算子是 complex-out-fusable 的,咱們可以将元素算子融合到它的輸出中。可以應用這些規則将計算圖轉換為融合版本。圖4展示了這種優化對不同工作負載的影響。作者發現融合算子通過減少記憶體通路會産生高達 1.2 到 2 倍的加速。

TVM | 一種用于深度學習的端到端自動優化編譯器

  Data Layout Transformation 有多種方法可以在計算圖中存儲給定的張量。最常見的資料排布選擇是 列主序 和 行主序。在實踐中,咱們可能會更加喜歡使用更加複雜的資料排布。例如,DL 加速器可能會采用 4X4 矩陣運算,需要将資料排布平鋪成 4X4 塊以優化通路局部性。

  資料排布優化将計算圖轉換為 可以使用更加好的内部資料排布 以在目标硬體上執行計算圖。它首先在給定記憶體層次結構規定的限制條件下為每個算子指定首選資料排布。如果首選資料布局不比對,就在生産者和消費者之間執行适當的布局轉換。

  雖然進階圖優化可以極大地提高 DL 工作負載的效率,但他們也隻是與算子庫差不多有效。目前,支援算子融合的少數 DL 架構需要算子庫來提供融合模式的實作。随着定期引入更多的網絡算子,可能的融合核心數量會急劇增加。當需要面對越來越多的硬體後端時,這種方法不再可持續,因為所需要的融合模式實作數量 與 必須支援的資料排布、資料類型和加速器内在函數的數量相結合。為程式所需的各種操作和每個後端手工設計算子核心是不可行的。為此,接下來提出了一種代碼生成方法,可以為給定模型的算子生成各種可能的實作。

4. 生成張量算子

  TVM 通過在每個硬體後端生成許多有效的實作并選擇優化的實作,為每個算子生成高效的代碼。這個過程建立在 Halide 将描述和計算規則 (或排程優化) 解耦的想法的基礎上,并将其擴充為支援新的優化 (嵌套并行、張量和延遲隐藏) 和 廣泛的硬體後端。

4.1 張量表達和排程空間

  引入了張量表達式語言來支援自動代碼生成,與進階計算圖表示不同,張量算子的實作是不透明的,每個算子都用索引公式表達語言來描述。以下代碼展示了計算轉置矩陣乘法的示例張量表達式。

TVM | 一種用于深度學習的端到端自動優化編譯器

  每個計算操作都指定輸出張量的形狀 和 描述如何計算它的每個元素的表達式。張量表達式語言支援常見的算術和數學運算,并涵蓋常見的 DL 算子形式。該語言沒有指定循環結構和許多其他執行細節,它提供了為各種後端新增硬體感覺優化的靈活性。采用來自 Halide 的解耦 計算/排程原則,使用排程來表示從張量表達式到低級代碼的特定映射。

  通過增量應用保持程式邏輯等價的基本轉換 (排程原語) 來建構排程。圖5展示了在專用加速器上排程矩陣乘法的示例。在内部,TVM 應用排程轉換,使用資料結構來跟蹤循環結構和其他資訊。然後此資訊可以幫助按給定的最終排程生成低級代碼。

TVM | 一種用于深度學習的端到端自動優化編譯器

  作者的張量表達式借鑒了 Halide、Darkroom 和 TACO。為了在許多後端實作高性能,必須要支援足夠多的排程原語來涵蓋不同硬體後端的各種優化。圖6總結了 TVM 支援的算子代碼生成過程和排程原語。作者重用了有用的原語和來自 Halide 的低級循環程式 AST,并且引入了新的原語來優化 GPU 和 加速器的性能。新的原語是實作最佳 GPU 性能所必需的,也是加速器所必需的。CPU、GPU 和 類TPU加速器是三種重要的深度學習硬體。

TVM | 一種用于深度學習的端到端自動優化編譯器

4.2 嵌套并行與協作

  并行性是提高 DL 工作負載中計算密集型核心效率的關鍵。現代 GPU 提供了大規模的并行性,要求咱們将并行模式融合到排程轉換中。大多數現有的解決方案都采用一種稱為嵌套并行的模型,這是一種 fork-join 的形式。該模型需要并行排程原語來并行化資料并行任務;每個任務可以進一步遞歸地細分為子任務,以利用目标架構的多級線程層次結構 (例如 GPU 中的線程組)。将此模型稱為無共享嵌套并行,因為一個工作線程無法在同一并行計算階段檢視其兄弟的資料。

  無共享方法的替代方法是協作擷取資料。具體來說,線程組可以協作擷取它們都需要的資料并将其放入共享記憶體空間中。這種優化可以利用 GPU 記憶體層次結構,并通過共享記憶體區域實作跨線程的資料重寫。TVM 使用排程原語來支援這種衆所周知的 GPU 優化,以實作最佳性能。以下 GPU 代碼示例對矩陣乘法進行了優化。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖7展示了這種優化的影響,将記憶體範圍的概念引入排程空間,以便可以将計算階段 (代碼中的 AS 和 BS) 标記為共享。

TVM | 一種用于深度學習的端到端自動優化編譯器

  如果沒有顯式記憶體範圍,自動範圍推理會将計算階段标記為 thread-local。共享任務必須計算組中所有工作線程的依賴關系,此外還必須正确插入記憶體同步屏障,以保證共享加載的資料對消費者可見。最後,除了對 GPU 有用之外,記憶體範圍還讓咱們可以标記特殊的記憶體緩沖區,并在針對專門的 DL 加速器時建立特殊的降低規則。

4.3 張量化

  DL 工作負載具有很高的算術強度,通常可以分解為張量算子,如矩陣乘 或 一維卷積。這些自然分解導緻了最近添加的張量計算原語的趨勢。這些新的原語為基于排程的編譯創造了機遇和挑戰,使用它們可以提高性能,是以編譯架構必須無縫的內建它們。這種張量化類似于 SIMD 架構的矢量化,但又有着顯著差異。指令輸入是多元的,具有固定或可變的長度,并且具有不同的資料排布。更為重要的是,這種原語不能是固定的,因為新的加速器有可能會出現張量指令的變體,是以需要一個可擴充的解決方法。

  通過使用張量内在聲明機制将目标硬體内在特性 與 排程分離,進而使張量化變得可以擴充。作者使用相同的張量表達式語言來聲明每個新硬體内在行為 和 與之相關的降低規則。以下代碼展示了如何聲明一個 8x8 張量的硬體内在函數。

TVM | 一種用于深度學習的端到端自動優化編譯器

  此外,引入了張量排程原語,以用相應的内在函數替換計算單元。編譯器将計算模式與硬體聲明相比對,并将其降低到相應的硬體内在特性。張量化将排程與特定的硬體原語分離,進而 可以輕松擴充 TVM 以支援新的硬體架構。生成的張量排程代碼符合高性能計算特性:将複雜的算子分解為一系列 micro-kernel 調用。咱們還可以使用 tensorize 原語來利用手工設計的 micro-kernel,這在某些平台上可能是有益的。例如,通過利用 bit-serial 矩陣向量乘法 micro-kernel 為移動 CPU 實作超低精度算子,這些算子對一比特或兩比特寬的資料類型進行操作。這個 micro-kernel 将結果累積成越來越大的資料類型,以最大限度地減少記憶體占用。将 micro-kernel 表達為 TVM 固有的張量 與 非張量化版本相比 可産生高達1.5倍的加速。

4.4 顯式記憶體延遲隐藏

  延遲隐藏是指将記憶體操作與計算重疊以最大限度地利用記憶體和計算資源的過程。它需要不同的政策,具體取決于目标硬體後端。在 CPU 上,記憶體延遲隐藏是通過同步多線程 或 硬體預取 來隐式實作的。GPU 依賴于許多線程的快速上下文切換。相比之下,諸如 TPU 之類的專用 DL 加速器通常傾向于使用解耦通路執行 (DAE) 架構記性更加精簡的控制,并将細粒度同步的問題轉移給軟體。

  圖9顯示了減少運作時延遲的 DAE 硬體管道。與單片硬體設計相比,流水線可以隐藏大部分記憶體通路開銷,幾乎可以充分利用計算資源。為了實作更加高的使用率,指令流必須增加細粒度的同步操作。沒有它們,就無法強制執行依賴關系,進而導緻錯誤執行。是以,DAE 硬體流水線需要在流水線階段依賴于細粒度的入隊/出隊操作,以保證正确執行,如圖9的指令流所示。

TVM | 一種用于深度學習的端到端自動優化編譯器

  對需要顯式低級同步的 DAE 加速器進行程式設計很困難。為了減少程式設計負擔,引入了虛拟線程排程原語,讓程式員可以指定進階資料并行程式,就像他們指定支援多線程的硬體後端一樣。然後 TVM 會自動将程式降低為具有低級顯式同步的單個指令流,如圖8所示。該算法從進階多線程程式排程開始,然後插入必要的低級同步操作以保證每個線程的正确執行。接下來,它将所有虛拟線程的操作傳遞給單個指令流中。最後,硬體恢複指令流中低級同步所規定的可用流水線的并行度。

  Hardware Evaluation of Latency Hiding 展示了延遲隐藏在基于 FPGA 的定制加速器設計上的有效性。在加速器上運作 ResNet 的每一層,并使用 TVM 生成兩個排程:一個有延遲隐藏,一個沒有。具有延遲隐藏的排程将程式與虛拟線程并行化 以進行管道并行,進而隐藏了記憶體通路延遲。結果在圖10 中顯示為 roofline圖表;roofline性能圖可以深入了解給定系統在不同基準測試中使用計算和記憶體資源的情況。總體而言,延遲隐藏提高了所有 ResNet 層的性能。峰值計算使用率從沒有延遲隐藏的 70% 提高到了隐藏延遲的 88%。

TVM | 一種用于深度學習的端到端自動優化編譯器

5. 自動優化

  鑒于豐富的排程原語集合,剩下的問題是為 DL 模型的每一層找到最佳的算子實作。在這裡,TVM 為與每一層關聯的特定輸入形狀和布局建立了一個專門的算子。這種專門的優化提供了顯著的性能優勢 (與針對較小的形狀和布局多樣性的手工代碼相比),但它也提出了自動化的挑戰。系統需要選擇排程優化:例如修改循環順序或優化記憶體層次結構,以及特定于排程的參數,例如切片大小和循環展開因子。這樣的組合選擇為每個硬體後端建立了算子實作的大的搜尋空間。為了應對這一挑戰,建構了一個具有兩個主要元件的自動排程優化器:一個新配置的排程搜尋器 以及 一個預測給定配置性能的機器學習成本模型。如圖11展示了這些元件和 TVM 的自動優化流程。

TVM | 一種用于深度學習的端到端自動優化編譯器

5.1 排程空間規範

  作者建構了一個排程模版規範 API,讓開發人員在排程空間中聲明旋鈕。在指定可能的排程時,模版規範允許在必要時結合開發人員的特定領域知識。作者還為每個硬體後端建立了一個通用主模闆,該模闆根據使用張量表達式語言表達的計算描述自動提取可能的旋鈕。在高層次上,咱們可能希望盡可能多的配置,并讓優化器管理選擇負擔。是以 優化器必須為咱們實驗中使用的真實世界 DL 工作負載搜尋數十億種可能的配置。

5.2 基于機器學習的成本模型

  從大型配置空間中找到最佳排程的一種方法是通過黑盒優化,即自動調整,該方法用于調優高性能計算庫。然而,自動調整需要許多實驗來确定一個好的配置。

  另一種方法是建構一個預定義的成本模型來指導對特定硬體後端的搜尋,而不是運作所有可能性并測量它們的性能。理想情況下,一個完美的成本模型會考慮影響性能的所有因素:記憶體通路模式、資料重用、管道依賴性和線程模式等。不幸的是,由于現代硬體日益複雜,這種方法很繁瑣。此外,每個新的硬體目标都需要一個新的 (預定義的) 成本模型。

  與此相反,作者采用統計方法來解決成本模組化問題。在這種方法中,排程搜尋器提出可以提高算子性能的配置。對于每個排程配置,使用一個 ML 模型,該模型将降低循環程式作為輸入,并預測其在給定硬體後端的運作時間。該模型使用搜尋期間收集的運作時測量資料進行訓練,不需要使用者輸入詳細的硬體資訊。當在優化過程中搜尋更多配置時,會定期更新模型,這也會提高其他相關工作負載的準确性。通過這種方式,ML 模型的品質會随着更多的試驗而提高。表1總結了自動化方法之間的主要差別。基于 ML 的成本模型在自動調整和預定義成本模組化之間取得平衡,并且可以從相關工作負載的曆史性能資料中受益。

TVM | 一種用于深度學習的端到端自動優化編譯器

  Machine Learning Model Design Choices 在排程搜尋器選擇将使用哪種 ML 模型時,咱們必須考慮兩個關鍵因素:品質和速度。排程搜尋器頻繁查詢成本模型,由于模型預測時間和模型調整時間而産生開銷。為了可用,這些開銷必須小于在真實硬體上測量性能所需的時間,根據具體的工作負載/硬體目标,這可能是幾秒鐘的數量級。這種速度要求将我們的問題與傳統的超參數調整問題區分開來,在傳統超參數調整問題中,執行測量的成本相對于模型開銷非常高,并且可以使用更加昂貴的模型。除了模型的選擇,咱們還需要選擇一個目标函數來訓練模型,例如配置預測運作時間的誤差。然而,排程搜尋器僅根據預測的相對順序 (A 運作速度比 B 快) 來選擇最佳候選,而不需要直接預測絕對執行時間。相反,作者采用 rank 目标來預測運作時成本的相對順序。

  作者在 ML 優化器中實作了幾種類型的模型。采用梯度樹提升模型 (基于 XGBoost),它根據從循環程式中提取的特征進行預測;這些功能包括每個循環級别的每個記憶體緩沖區的記憶體通路計數和重用率,以及循環注釋的 one-hot 編碼,例如 vectorize、unroll 和 parallel。作者還評估了一個神經網絡模型,該模型使用 TreeRNN 在沒有特征工程的情況下總結循環程式的 AST。圖13總結了成本模型的工作流程。

TVM | 一種用于深度學習的端到端自動優化編譯器

  作者發現樹增強和 TreeRNN 具有相似的預測品質。然而,前者的預測速度要快兩倍,而且訓練時間要少得多。是以,在實驗中選擇梯度樹提升作為預設成本模型。盡管如此,作者還是相信這兩種方法都是有價值的,并期待在這個問題上進行更加多的研究。平均而言,樹增強模型能夠在 0.67 毫秒内進行預測,這比運作實際測量要快數千倍。圖12将基于 ML 的優化器與黑盒自動調整方法進行了比較;前者比後者更加快地找到了更好的配置。

TVM | 一種用于深度學習的端到端自動優化編譯器

5.3 排程搜尋

  一旦咱們選擇了一個成本模型,我們就可以使用它來選擇建議的配置,在這些配置上疊代地運作實際測量。在每次疊代中,搜尋器使用 ML 模型的預測來選擇一批候選者來運作測量。然後将收集的資料用作訓練資料以更新模型。如果不存在初始訓練資料,則搜尋器會選擇随機候選者進行測量。

  最簡單的搜尋算法通過成本模型枚舉并運作每個配置,選擇前 k 個預測的執行者。但是這種政策再搜尋空間很大時變得難以處理。相反,作者采用并行模拟退火算法。搜尋器從随機配置開始,并且在每一步中,随機走到附近的配置。如果成本如成本模型所預測的那樣降低,則這種轉變是成功的。如果目标配置成本較高,則可能會失敗 (拒絕)。這樣随機遊走在傾向于收斂于成本模型預測的成本較低的配置。搜尋狀态在成本模型更新中持續存在;咱們從這些更新後的最後一個配置繼續搜尋。

5.4 分布式裝置池和 RPC

  分布式裝置池可以擴充硬體試驗的運作,并在多個優化作業之間實作細粒度的資源共享。TVM 實作了一個定制的、基于 RPC 的分布式裝置池,使用戶端能夠在特定類型的裝置上運作程式。咱們可以使用此接口在主機編譯器上編譯程式,請求遠端裝置、遠端運作函數,并在主機上通路相同腳本中的結果。TVM 的 RPC 支援動态上傳并運作使用其運作時約定的交叉編譯子產品和函數。是以,相同的基礎架構可以執行單個工作負載優化和端到端 graph 推理。作者的方法可以跨多個裝置自動執行編譯、運作和配置步驟。這種基礎設施對于嵌入式裝置尤其重要,因為傳統上這些裝置需要繁瑣的手動操作來進行交叉編譯、代碼部署和測量。

6 實驗

  圖14展示了 TVM、MXNet、Tensorflow 和 Tensorflow XLA 在 NVIDIA Titan X 上的性能評估。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖15展示了 cuDNN、TensorComprehensions、MX Kernel、TVM、TVM PT 在 TITAN X 上對 ResNet18 中所有 conv2d 算子,和 MobileNet 中所有深度 conv2d 算子的相對加速情況。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖16展示了 TVM 和 TFLite 在 ARM A53 上端到端性能評估。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖17展示了 cuDNN、TensorComprehensions、MX Kernel、TVM、TVM PT 在 ARM A53 上對 ResNet18 中所有 conv2d 算子,和 MobileNet 中所有深度 conv2d 算子的相對加速情況。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖18展示了手工優化、TVM單線程 和 TVM多線程對 ResNet 中 conv2d 算子的相對加速。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖19展示了在 Mali-T860MP4 的端到端實驗結果。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖20展示了 VDLA 硬體設計總覽。

TVM | 一種用于深度學習的端到端自動優化編譯器

  圖21展示了将 ResNet 工作負載中的卷積解除安裝到基于FPGA的加速器上的加速效果。

TVM | 一種用于深度學習的端到端自動優化編譯器

7 總結

  作者提出了一個端到端的編譯堆棧,以解決跨各種硬體後端的深度學習的基本優化挑戰。提出的系統包括自動化的端到端的優化,這在以前是一項勞動密集型和高度專業化的任務。作者希望這項工作将鼓勵對端到端編譯方法的更多研究,并為 DL 系統軟硬體協同設計技術開辟新的機會。

8 參考

  [1] TVM: An Automated End-to-End Optimizing Compiler for Deep Learning.

繼續閱讀