自動生成硬體優化內核：陳天奇等人發布深度學習編譯器TVM

有了 TVM，業界與學界開發者們可以快速、輕鬆地在各個系統（包括手機、嵌入式設備與低功耗晶元）上部署深度學習應用程序，同時無須擔心資源與速度的限制。「TVM 作為神經網路和硬體後端之間的共同層，消除了為每類設備或伺服器優化各自基礎架構的需要。」TVM 項目負責人陳天奇表示，「我們的框架允許開發人員快速、輕鬆地部署和優化大量硬體設備上的深度學習系統。」

項目鏈接：dmlc/tvm

深度學習如今無處不在且必不可少。這次創新部分得益於可擴展的深度學習系統，比如 TensorFlow、MXNet、Caffe 和 PyTorch。大多數現有系統針對窄範圍的伺服器級 GPU 進行了優化，並且需要在其他平台（如手機、IoT 設備和專用加速器（FPGA、 ASIC））上部署大量工作。隨著深度學習框架和硬體後端數量不斷增加，我們提出了一個統一的中間表徵（IR）堆棧，可以彌補以生產力為中心的深度學習框架與性能或效率導向的硬體後端之間的差距。

我們很高興宣布 TVM 的發行將解決這一問題。TVM 是一個全新的框架，它可以：

• 優化 CPU、GPU 和其他專業化硬體的常規深度學習的計算量
• 自動轉換計算圖以最小化內存利用率，優化數據布局，融合計算模式
• 提供從現有的前端框架到裸機硬體的端到端編譯，一直到瀏覽器可執行的 Javascript

有了 TVM 的幫助，只需少量額外工作我們即可輕易地在手機端、嵌入式設備甚至瀏覽器上運行深度學習任務。TVM 還為多硬體平台上的深度學習工作負載提供了統一的優化框架，其中包括依賴全新計算原語的專用加速器。

通過採用編譯器社區的共同理念，並提供兩個中間表徵層，我們有效地將高級深度學習演算法降低到多硬體後端。在今天的版本中，我們開源了包含 x86、ARM、OpenCL、Metal、CUDA 和 Javascript 優化原語的 TVM 工具包。目前，我們正積極地為專用硬體加速和英偉達的 GEMM 優化的 Volta 架構提供支持。

技術細節

TVM 堆棧的目標在於提供一個可重複使用的工具鏈來將高級神經網路描述從深度學習框架前端向下編譯為多硬體後端的低級機器代碼。拿 Apache MXNet 作為前端的例子，以下代碼段演示了如何使用 TVM 將深層學習模型的高級描述編譯為針對目標硬體定製的優化可執行模塊。

挑戰在於支持多個硬體後端，同時將計算、內存和能源足跡（energy footprint）保持在最低水平。我們借鑒了編譯器社區的共同理念，以彌合大量深度學習框架和硬體後端之間的差距：我們構建了一個由 NNVM 組成的兩級中間層，一個可用於任務調度、內存管理以及 TVM（用於優化計算內核的富有表現力的低級 IR）的高級中間表徵（IR）。

堆棧的第一級是基於計算圖的表徵。計算圖是一個有向無環圖，其表徵作為節點的計算和作為邊緣的數據流依賴性。這一表徵非常強大：它允許我們將操作屬性融入到計算圖中，並指定轉換規則以迭代優化計算圖。這是一個被大多數現有深度學習框架的採用的慣常方法，包括 TVM 堆棧中的 NNVM 圖表徵、TensorFlow XLA 和英特爾的 ngraph。

很多強大的優化可由圖優化框架支持。例如，我們提供了一個子線性內存優化功能，允許用戶在單塊 GPU 上訓練 1000 層的 ImageNet ResNet。

然而，我們發現只基於 IR 的計算圖並不足以應對支持不同硬體後端的挑戰。原因是像卷積或矩陣乘法一樣的圖運算符可以非常不同的方式映射和優化每個硬體後端。這些特定硬體的優化在內存布局、並行線程模式、緩存訪問模式和硬體原語選擇方面的變數可能很大。我們希望能夠以通用表徵的方式明確地表達這些優化旋鈕（optimization knob），以有效地導航優化空間。

我們構建了一個低級表徵以解決這一問題。該表徵基於索引公式，並額外支持遞推計算。

低級 IR 採用現有的圖像處理語言（比如 Halide 或 darkroom）的原理以制定一個富有表現力的深度學習 DSL。TVM 構建了由循環轉換工具比如 loopy、多面體分析啟發的圖優化。我們也從深度學習框架比如 MXNet、TensorFlow、Theano 使用的數據流描述語言汲取靈感，然後在調度階段對 TVM 中描述的演算法進行處理，以應用針對目標硬體後端定製的轉換。

TVM 包括通常出現在 CPU 優化框架中的標準轉換原語。更重要的是，TVM 通過利用線程協作模式、數據布局轉換和強大的新計算原語，集成了針對 GPU 的全新優化原語。TVM 與 NNVM 的結合為優化軟體堆棧中深度學習工作負載提供了一個很好的機會，從而實現了圖級別與運算符級別優化的聯合計算。

多語言和平台支持

TVM 的眾多優勢之一在於它可以支持多種語言和平台。TVM 框架中有兩個組件：編譯器棧，其中包含完整的優化庫以產生優化的機器代碼；以及輕量級運行環境，提供在不同平台上部署編譯模塊所需的可移植性。

TVM 目前支持嵌入式編譯器堆棧的 Python 和 C++ 介面。我們在設計中最大限度地重用框架，以便編譯器堆棧的改進可以在 Python 和 C++ 組件之間互換使用。我們還提供了一個輕量級運行環境，可以在包括 Android、iOS、樹莓派和 Web 瀏覽器在內的平台上直接運行諸如 Javascript、Java、Python 和 C++ 等語言的 TVM 編譯代碼。

遠程部署和執行

TVM 支持使用 TVM PRC 進行交叉編譯，測試嵌入式設備，這是一種輕量級界面，用於在遠程嵌入式設備上部署和執行 TVM 交叉編譯模塊。我們 TVM 用戶提供了易用的高級 Python 界面，用於在各種低級嵌入式設備上進行遠程編譯、優化，並測試深度學習演算法。

性能

TVM 目前仍然處在早期發展階段，我們會在未來看到更多改進，但目前它已經展現出了強大的實力。

TVM 為我們提供了靈活性，讓我們可以探索不同硬體平台上，不同深度學習內核中豐富的優化空間。例如，TVM 允許我們裁剪數據布局，為內核與平台融合模式需求。注意，構建基線庫可用於更通用的問題，而 TVM 的優化內核通過自動調整過程對人們評估的工作負載進行了大量調整。TVM 可以作為快速生成專屬內核的絕佳工具。

本節列出的結果仍有提升空間。

樹莓派

在結果的第一部分，我們比較了 TVM CPU 計劃與 nnpack 在樹莓派 3B 上執行 ResNet 任務的表現。由於時間有限，我們使用 TVM 實現直接卷積，而 nnpack 對 3×3 內核執行 winograd 轉換。

我們可以看到，使用 TVM 自動調整內核，我們可以獲得類似於 nnpack 中的手動優化內核的性能（實驗在樹莓派上進行）。

GPU 結果

PS. 該測試結果由 Leyuan Wang（AWS/UCDavis）、Yuwei Hu（圖森未來）與 Weitang Liu（AWS/UCDavis）提供。

為了進行概念證明，我們構建了一種端到端編譯管道，可以將 MXNet 模型編譯成 TVM 執行圖。我們通過自動融合運算符並使 TVM 生成融合的內核，在圖形節點之間和之中應用優化。我們對 MobileNet 在 ImageNet 上的運行表現進行基準測試，結果如下：

我們可以看到，TVM 在速度上快過基準方法。更有意思的是，內核融合可以帶來更快的速度提升。值得一提的是，TVM 可以自行生成所有優化 GPU 內核，而不依賴於外部庫，如 CuDNN。

我們正在開展新的實驗，期待獲得更好的結果。

開源

TVM 始於華盛頓大學 Paul G. Allen 計算機科學與工程學院的研究項目。它被設計成支持 DLPack——一個被多種深度學習框架共享的張量數據結構。在開發過程中，我們得到了來自華盛頓大學、AWS、奇虎 360、Facebook、香港科技大學、圖森未來、UCDavis、上海交大以及 DMLC 開源社區和 DLPack 計劃成員們的支持。展望未來，該項目將繼續遵循 Apache 開源模式，力圖構建一個由社區維護的項目，歡迎大家的參與。

致謝

本項目的成功離不開以下所有人的通力合作：Yizhi Liu（奇虎 360）、Yuwei Hu（圖森未來）、Xingjian Shi（香港科技大學）、Leyuan Wang（AWS/UCDavis）、Nicolas Vasilache（Facebook）、Jian Weng（UCLA）、Weitang Liu（AWS/UCDavis）、Edward Z. Yang（Facebook）、Lianmin Zheng（上海交通大學）、Qiao Zhang（華盛頓大學）、William Moses（Facebook/MIT）與 Hu Shiwen。作者同時感謝 Xianyi Zhang 和他富有建設性的討論。

源代碼

• 本工具 GitHub 頁面：dmlc/tvm
• 下周將向 Apache MXNet 上游提供基於 TVM 堆棧和 MXNet 的圖形編譯工具鏈示例
• TVM 與 DLPack 兼容，這意味著它可以輕鬆支持 MXNet、PyTorch、Caffe2 和 tiny-dnn 等採用該標準的框架。

選自tvmlang 機器之心編譯