陳天奇等人提出TVM：深度學習自動優化程式碼生成器

TVM 是由華盛頓大學在讀博士陳天奇等人提出的深度學習自動程式碼生成方法，去年 8 月機器之心曾對其進行過簡要介紹。該技術能自動為大多數計算硬體生成可部署優化程式碼，其效能可與當前最優的供應商提供的優化計算庫相比，且可以適應新型專用加速器後端。近日，這項研究的論文《TVM: End-to-End Optimization Stack for Deep Learning》終於完成，內容包含新方法的介紹與討論，以及 TVM 在英偉達、AMD 的 GPU、樹莓派及一些 FPGA 上的效能評估。

專案連結：https://github.com/dmlc/tvm

深度學習模型可以識別影像、處理自然語言，以及在部分具有挑戰性的策略遊戲中擊敗人類。在其技術發展的過程中，現代硬體穩步推進的計算能力扮演了不可或缺的作用。很多目前最為流行的深度學習框架，如 TensorFlow、MXNet、Caffe 和 PyTorch，支援在有限型別的伺服器級 GPU 裝置上獲得加速，這種支援依賴於高度特化、供應商特定的 GPU 庫。然而，專用深度學習加速器的種類越來越多，這意味著現代編譯器與框架越來越難以覆蓋所有的硬體。

顯而易見，以現有的點到點方式實現不同深度學習框架對所有種類的硬體進行後端支援是不現實的。我們的最終目標是讓深度學習負載可以輕鬆部署到所有硬體種類中，其中不僅包括 GPU、FPGA 和 ASIC（如谷歌 TPU），也包括嵌入式裝置，這些硬體的記憶體組織與計算能力存在著顯著的差異（如圖 1 所示）。考慮到這種需求的複雜性，開發一種能夠將深度學習高階程式降低為適應任何硬體後端的低階優化程式碼的優化框架是最好的方法。

目前的深度學習框架依賴於計算圖的中間表示來實現優化，如自動微分和動態記憶體管理 [3,7,4]。然而，圖級別的優化通常過於高階，無法有效處理硬體後端運算元級別的轉換。另一方面，目前深度學習框架的運算元級別庫通常過於僵化，難以輕鬆移植到不同硬體裝置上。為了解決這些問題，我們需要一個可實現從計算圖到運算元級別的優化，為各種硬體後端帶來強大效能的編譯器框架。

圖 1：CPU、GPU 與類 TPU 加速器需要不同的片上儲存架構和計算基元。在生成優化程式碼時我們必須考慮這一問題。

圖 2：TVM 堆疊圖。目前的堆疊支援多種深度學習框架以及主流 CPU、GPU 以及專用深度學習加速器。

優化的四大基本挑戰

深度學習的優化編譯器需要同時展示高階別與低階別的優化，在論文中，研究人員總結了在計算圖級別與張量運算元級別上的四大基本挑戰：

1. 高階資料流複寫：不同的硬體裝置可能具有截然不同的記憶體層次結構，因此，融合運算元與優化資料佈局的策略對於優化記憶體訪問至關重要。
2. 跨執行緒記憶體複用：現代 GPU 與專用加速器的記憶體可被多個計算核心共享，傳統的無共享巢狀並行模式已不再是最優方法。為優化核心，在共享記憶體負載上的執行緒合作很有必要。
3. 張量計算內部函式：最新的硬體帶來了超越向量運算的新指令集，如 TPU 中的 GEMM 運算元和英偉達 Volta 架構中的 Tensor Core。因此在排程過程中，我們必須將計算分解為張量算術內部函式，而非標量或向量程式碼。
4. 延遲隱藏（Latency Hiding）：儘管在現代 CPU 與 GPU 上，同時擁有多執行緒和自動快取管理的傳統架構隱藏了延遲問題，但專用的加速器設計通常使用精簡控制與分流，這為編譯器堆疊的排程帶來了複雜性。所以，排程仍需仔細，以隱藏記憶體訪問延遲。

TVM：一個端到端優化堆疊（見圖 2），該端到端優化編譯器堆疊可降低和調整深度學習工作負載，以適應多種硬體後端。TVM 的設計目的是分離演算法描述、排程和硬體介面。該原則受到 Halide [22] 的計算／排程分離思想的啟發，而且通過將排程與目標硬體內部函式分開而進行了擴充套件。這一額外分離使支援新型專用加速器及其對應新型內部函式成為可能。TVM 具備兩個優化層：計算圖優化層，用於解決第一個排程挑戰；具備新型排程基元的張量優化層，以解決剩餘的三個挑戰。通過結合這兩種優化層，TVM 從大部分深度學習框架中獲取模型描述，執行高階和低階優化，生成特定硬體的後端優化程式碼，如樹莓派、GPU 和基於 FPGA 的專用加速器。該論文做出了以下貢獻：

• 我們構建了一個端到端的編譯優化堆疊，允許將高階框架（如 Caffe、MXNet、PyTorch、Caffe2、CNTK）專用的深度學習工作負載部署到多種硬體後端上（包括 CPU、GPU 和基於 FPGA 的加速器）。
• 我們發現了提供深度學習工作負載在不同硬體後端中的效能可移植性的主要優化挑戰，並引入新型排程基元（schedule primitive）以利用跨執行緒記憶體重用、新型硬體內部函式和延遲隱藏。
• 我們在基於 FPGA 的通用加速器上對 TVM 進行評估，以提供關於如何最優適應專用加速器的具體案例。

我們的編譯器可生成可部署程式碼，其效能可與當前最優的特定供應商庫相比，且可適應新型專用加速器後端。

圖 3：兩層卷積神經網路的計算圖示例。圖中每個節點表示一次運算，它消耗一或多個張量，並生成一或多個張量。張量運算可以通過屬性進行引數化，以配置其行為（如 padding 或 stride）。

論文：TVM: End-to-End Optimization Stack for Deep Learning

論文連結：https://arxiv.org/abs/1802.04799

摘要：可擴充套件框架，如 TensorFlow、MXNet、Caffe 和 PyTorch 是目前深度學習領域中最流行和易用的框架。但是，這些框架只對窄範圍的伺服器級 GPU 進行優化，要把工作負載部署到其他平臺，如手機、嵌入式裝置和專用加速器（如 FPGA、ASIC），則需要大量手動工作。我們提出了 TVM，一個端到端的優化堆疊，具備圖形級和運算元級的優化，以為多種硬體後端提供深度學習工作負載的效能可移植性。我們討論了 TVM 所解決的深度學習優化挑戰：高階運算元融合（operator fusion）、多執行緒低階記憶體重用、任意硬體基元的對映，以及記憶體延遲隱藏。實驗結果證明 TVM 在多個硬體後端中的效能可與適應低功耗 CPU 和伺服器級 GPU 的當前最優庫相比。我們還通過針對基於 FPGA 的通用深度學習加速器的實驗，展示了 TVM 對新型硬體加速器的適應能力。該編譯器基礎架構已開源。

原文地址：http://www.aibbt.com/a/14034.html