技術探索

深度學習編譯器的發展趨勢

工研院資通所 董明智 林國弘 黃俊堯

深度學習模型可以識別影像與圖片、處理自然語言、以及其他AI應用上。在技術發展的過程中,持續精進的硬體運算能力扮演了不可或缺的作用。除了硬體之外,前端深度學習軟體套件與框架(Framework)的發展,如:TensorFlow、MXNet、Caffe2、Keras、和PyTorch等,這些可擴展的深度學習系統提供軟體開發人員便利的開發工具,由於前端便利的軟體工具與後端硬體的持續演進,讓深度學習應用蓬勃發展。但深度學習的套件與框架大多數僅支援CPU與GPU硬體環境,如果要將深度學習程式應用到其他硬體上,需要付出相當大的工程成本進行編譯器開發,這讓深度學習系統的靈活部署遇到困難。因此,大量多元的深度學習套件或框架,與越來越多的硬體加速器之間需要一個深度學習編譯器(Neural Network Compiler)當作橋樑。國際大廠Facebook、Intel、Amazon與華盛頓大學等機構紛紛提出深度學習編譯器架構,旨在讓深度學習模型能夠在各種不同的硬體平台執行,從伺服器GPU、手機、嵌入式設備到低功耗專用加速晶片等不同的系統上,快速地部署深度學習應用,亦可讓廠商不需為不同硬體提供不同編譯環境,降低廠商的硬體編譯環境開發時程與成本。

精彩內容

1. 深度學習編譯架構
2. 深度學習編譯技術的競合
3. TVM深度學習編譯技術
4. AI應用與晶片產業快速發展的關鍵

深度學習編譯器架構

在AI晶片軟體編譯環境中,分為計算圖優化(Graph Level Optimization)與運算優化(Operator Level Optimization),計算圖優化透過優化記憶體配置與運算融合(Operator Fusion)等方式減少記憶體資料搬移並提升運算效率;運算優化則是依據硬體特性進行張量(Tensor)運算轉換與執行緒管理(Threads Cooperation)等以提升運算效能。 目前的深度學習框架依賴於計算圖優化來實現最佳化,如自動微分和動態記憶體管理。然而,計算圖優化通常無法針對硬體特性進行優化處理,無法有效處理硬體運算層級的轉換;另一方面,目前深度學習框架的運算程式庫通常過於僵化,難以移植到不同硬體設備上,為了解決這些問題,需要一個可以兼容計算圖與運算優化的編譯架構,如圖1的深度學習編譯環境,不同的硬體之間除了有計算特性差異,也有不同的主記憶體與快取記憶體組成結構,因此,藉由運算融合與資料架構的優化策略,將記憶體存取進行最佳化處理至關重要;在執行緒的管理上,通常硬體加速器的記憶體可以被多個計算核心共享,透過共享記憶體管理與執行緒管理的整合,可以提升記憶體的存取效率,更進一步提升運算效能。一般的數值運算適合CPU的硬體架構,GPU可快速執行陣列運算,FPGA或TPU可支援矩陣運算,因此將深度學習相關運算分解為張量運算的函數組合,可以根據硬體架構提供重組或分割張量運算;至於延遲隱藏(Latency Hiding)的處理,大都是針對CPU或GPU提供多執行緒與自動暫存管理的延遲隱藏機制[1],但是其他專用的FPGA或TPU等加速器,在設計上通常使用精簡的控制與存取分流,在編譯環境中仍需對這些專用硬體提供適當的延遲隱藏演算法,藉以提升運算效能。

圖1  深度學習譯環境架構圖1 深度學習譯環境架構

深度學習編譯技術的競合

深度學習編譯器可以針對硬體特性進行運算優化,讓不同深度學習框架所開發的程式可以在不同硬體平台上達到更佳的執行效能,國際上許多知名軟硬體廠商及學術機構都投入深度學習的應用與發展,針對硬體的差異性,深度學習編譯器的開發成為運行深度學習於各硬體效能的關鍵重點,較著名的有Tensorflow XLA、nGraph、Glow、以及TVM(Tensor Virtual Machine)。Tensorflow XLA[2]是Google提出的深度學習編譯器,主要還是針對Google的Tensorflow框架進行優化。Intel則提出nGraph[3][4] 編譯器,擴展了深度學習模型的適用性和可移植性,nGraph編譯器能夠將深度學習模型轉換為可執行的優化程式,根據不同硬體的特性進行優化藉以提升執行效能。Glow是Facebook開發的機器學習編譯器[5][6],為了擴展深度學習框架能支援不同硬體加速器,需要提供能針對不同硬體特性優化程式的編譯器,善用硬體特徵藉以提升深度學習程式執行效能。TVM是華盛頓大學所研發的開源深度學習編譯環境[7],近來廣為各大指標性公司所採用,參與的機構除了Amazon之外,華為、Xilinx、ARM、Microsoft、Qualcomm、甚至自行提出編譯器的Facebook與Intel都發表了TVM之應用[8]。TVM近兩年急速發展,陸續擴增功能與支援,從最早的編譯器發展到包含深度學習程式庫的完整架構並持續研究更新,在深度學習編譯器領域提出了很多新概念。圖2整理了nGraph、Glow、TVM之深度學習編譯架構差異,他們都支援了計算圖優化與運算優化兩個層級,提供完整的深度學習編譯功能。依據目觀察,以開源(Open Source)的深度學習編譯器發展上,TVM會是深度學習編譯器開源技術的趨勢,而工研院也以TVM為基礎進行深度學習編譯器的開發,提供各種深度學習加速器完整的編譯環境,功能包含: 依據硬體組態進行運算排程自動生成與優化、支援在無作業系統的硬體上運行(Bare Metal)、已量化模型(Pre-Quantized Model)、編譯期間量化(Post-Quantization)、異質多核心推論、以及支援學習/推論混合多模型等。

圖2  nGraph、Glow、TVM之深度學習編譯架構圖2 nGraph、Glow、TVM之深度學習編譯架構

TVM深度學習編譯技術

TVM是華盛頓大學提出且Amazon支援開發[7][9]的開源程式,其架構如圖3所示,是一個端到端的優化編譯架構,具備計算圖優化和運算優化的分層,這種抽象化分層可為不同硬體提供深度學習工作負載的可攜性,TVM 探討了深度學習編譯器內進行優化處裡的挑戰:運算融合、多執行緒的記憶體管理與重用、自動調度(Auto-TVM)、延遲隱藏、以及異質(Heterogeneous)硬體間的程式執行。TVM的目標是讓深度學習負載可以輕鬆部署到所有硬體種類上執行,其中不僅包括 CPU、GPU、和FPGA,也包括嵌入式設備、Android設備、iOS設備、Raspberry Pi、無作業系統硬體(Bare Metal)設備等。

圖3  TVM編譯器之架構圖3 TVM編譯器之架構

TVM在運算優化的設計方式是將演算法描述(Description)、調度(Scheduling)和硬體介面分開到不同抽象層處理,此設計參考了Halide[10]的計算/調度分離之設計方式,並進一步將硬體內建程式庫也進行分離,這讓整體編譯環境在支援新硬體上較為便利。計算圖優化首先進行資料流程複寫(Graph-Level Dataflow Rewriting),將深度學習框架的運算轉譯成TVM的計算圖表示方式,再依據硬體的記憶體結構進行融合運算與資料佈局(Data Layout)轉換,藉以降低記憶體的存取次數,這對提升運算效能至關重要(記憶體存取較Buffer存取更耗時且耗能);運算優化抽象層主要進行記憶體重用(Memory Reuse)、張量運算分割(例如:根據硬體特性進行矩陣運算的Split & Tiling)、延遲隱藏等跟硬體相關的運算優化處理,最後再產生可在硬體執行的最佳化程式碼,提升深度學習程式在特定硬體上的執行效能。

AI應用與晶片產業快速發展的關鍵

Intel、Amazon和Facebook等積極提出深度學習編譯環境並將研發成果開源,為AI晶片廠商提供了編譯器框架的基礎,大幅降低了自行研發深度學習編譯器的門檻。從這些國際大廠陸續針對深度學習編譯器提出支援與開源專案,大幅降低開發者把深度學習模型部署到不同硬體設備的複雜度,這也形成了AI應用與晶片產業快速發展的關鍵。

參考文獻

[1] V. Volkov, "Understanding Latency Hiding on GPUs", Electrical Engineering and Computer Sciences, University of California at Berkeley, Technical Report No. UCB/EECS-2016-143, 2016.

[2] (2019) XLA Overview. [Online]. Available: https://www.tensorflow.org/xla/overview

[3] C. Scott, K. B. Arjun, B. Anahita, et al. (2018) Intel nGraph: An Intermediate Representation, Compiler, and Executor for Deep Learning [online]. Available: https://arxiv.org/abs/1801.08058

[4] (2019) nGraph: An end to end deep learning compiler for inference and training with extensive framework and hardware support. [Online]. Available: https://www.ngraph.ai/

[5] N. Rotem, J. Fix, S. Abdulrasool, et al. (2019) Glow: Graph Lowering Compiler Techniques for Neural Networks [online]. Available: https://arxiv.org/abs/1805.00907

[6] (2019) Glow: Compiler for Neural Network hardware accelerators. [Online]. Available: https://github.com/pytorch/glow

[7] T. Q. Chen, T. Moreau, Z. H. Jiang, et al. (2018) TVM: An Automated End-to-End Optimizing Compiler for Deep Learning [online]. Available: https://arxiv.org/abs/1802.04799

[8] (2019) TVM and Deep Learning Compilation Conference. [Online]. Available: https://sampl.cs.washington.edu/tvmconf/

[9] (2019) End to End Deep Learning Compiler Stack. [Online]. Available: https://tvm.ai/

[10] R. K. Jonathan, B. Connelly, A. Andrew, el al., “Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines,” In Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2013, P. 519–530.