用TVM在硬件平臺上部署深度學習工作負載的端到端 IR 堆棧

深度學習已變得無處不在，不可或缺。這場革命的一部分是由可擴展的深度學習系統推動的，如滕索弗洛、MXNet、咖啡和皮托奇。大多數現有系統針對范圍狹窄的服務器級 GPU 進行了優化，需要在其它平臺，如移動電話、物聯網設備和專用加速器（FPGA、ASIC）上部署大量精力。隨著深度學習框架和硬件后端數量的增加，建議建立一個統一的中間表示 （IR） 堆棧，以縮小以生產力為中心的深度學習框架與面向性能或效率的硬件后端之間的差距。

TVM 是一個新穎的框架，可以：
? 表示并優化 CPU、GPU 和其它專業硬件的常見深度學習計算工作負載
? 自動轉換計算圖以最大限度地減少內存利用、優化數據布局和融合計算模式
? 提供從現有前端框架到裸機硬件的端到端編譯，一直到瀏覽器可執行的 javascripts。
TVM可以輕松地在移動電話、嵌入式設備甚至瀏覽器上運行深度學習工作負載，不需要額外的支持。TVM 還為眾多硬件平臺（包括依賴新型計算原始的專用加速器）的深度學習工作量提供了統一的優化框架。

采用編譯器方案，提供兩個中間表示層，有效地將高級深度學習算法降低到多個硬件后端。
開源TVM包包含x86、ARM、OpenCL、Metal、CUDA和Javascript的原型優化。致力于專業硬件加速和 Nvidia 的 GEMM 優化Volta架構提供支持。
技術細節
TVM 堆棧的目標是提供可重復使用的工具鏈，以編譯從深度學習框架前端到低級機器代碼的高層神經網絡描述，用于多個硬件后端。以 Apache MXNet 為例，以下代碼片段演示，如何使用 TVM 將深度學習模型的高level描述，編譯為針對目標硬件定制的優化可執行模塊。

挑戰在于支持多個硬件后端，同時將計算、內存和能量消耗保持在最低水平。借用AICompiler的智慧，以彌合眾多深度學習框架和硬件后端之間的差距：
構建了由 NNVM 組成的兩級中間層、用于任務調度和內存管理的高級中級代表 （IR） 和用于優化計算內核的具有表現力的低level IR 的 TVM。
堆棧的第一個級別是基于計算圖形的表示。計算圖是指示式循環圖，表示計算作為節點，數據流依賴性表示邊緣。此表示非常強大：
將操作屬性bake到計算圖中，指定轉換規則，迭代優化計算圖形。這是大多數現有深度學習框架所采用的常見方法，包括 TVM 堆棧中的 NNVM 圖形表示、TensorFlow XLA 和英特爾的 ngraph。

圖形優化框架可以支持許多強大的優化。例如，提供了sublinear亞線性內存優化功能。可在單個 GPU 上訓練 1000 層ImageNet ResNet。

發現僅基于計算圖的 IR 不足以解決支持不同硬件后端的挑戰。原因是，對于每個硬件的后端，可能會以非常不同的方式映射和優化單個圖形算子，如卷積或矩陣乘法。這些針對硬件的優化在內存布局、并行線程模式、緩存訪問模式和硬件原始選擇方面可能存在巨大差異。希望能夠以通用表示方式明確表達這些優化模塊，有效引導優化空間。
構建一個低level級別的表示來解決這個問題。此表示基于索引公式，為復發計算提供額外支持。

低級 IR 采用現有圖像處理語言（如Halid或darkroom）的原則，定制具有表現力的深度學習 DSL。TVM 構建低level級別優化的理論來自循環轉換工具，如循環和多面體分析。從 MXNet, TensorFlow, Theano.等深度學習框架中使用的數據流描述語言中汲取靈感。然后在調度階段處理 TVM 中描述的算法，應用針對目標硬件后端量身定制的轉換。

TVM 包括 CPU 優化框架中常見的標準原型轉換。更重要的是，TVM 通過利用線程調度模式、數據布局轉換和強大的新原型計算，整合了針對 GPU 的新型優化。使用 TVM 與 NNVM 相結合，為優化整個軟件堆棧中的深度學習工作負載提供了豐富的機會，實現了聯合計算圖形級別和算子級別的優化。
多語言和平臺支持
TVM 的眾多優勢之一，在于對多個平臺和語言的豐富支持。框架的兩個組件：包含完整優化庫以生成優化機器代碼的編譯器堆棧，以及輕量級lightweight的runtime，以及在不同平臺上部署編譯模塊所需的便攜性。

TVM 當前支持嵌入式編譯器堆棧的python和C++接口。設計框架時，考慮到了最大的再利用，以便編譯器堆棧改進，可以在 Python 和C++組件之間互換應用。
提供輕量級的runtime，可以在包括 Android、iOS、raspberry pi和 Web 瀏覽器在內的平臺上，以javascript, java, python, C++等語言直接運行 TVM 編譯代碼。
遠程部署和執行

TVM 支持使用 TVM RPC 對嵌入式設備進行交叉編譯和測試，這是一個輕量級界面，可在遠程嵌入式設備上部署和執行 TVM 交叉編譯模塊。這為 TVM 用戶提供了熟悉的高級 Python 界面，以便在各種低級嵌入式設備上遠程編譯、優化和測試深度學習算法。
性能
TVM能夠靈活地探索各種深度學習內核的豐富優化空間，用于多個硬件平臺。例如，TVM為最關心的內核和平臺定制數據布局和融合模式要求。基線庫是針對更通用的問題創建的，而 TVM 的優化內核則針對通過自動調整過程評估的工作負載進行了大量調整。TVM 充當快速生成獲取此類專用內核的橋梁。
Raspberry Pi樹莓派
在結果的第一部分，將TVM CPU調度與樹莓派3B執行網格工作負載上的nnpack進行了比較。使用TVM實現直接卷積，而nnpack用于執行3x3內核的winograd conv。

可以發現，通過 TVM 的自動調諧內核，可以獲得類似于 nnpack 中hand-optimized kernels手動優化內核的性能，用于aspberry pi實驗。
GPU 結果
創建了一個端到端的編譯管道，可以將 MxNet 模型編譯到 TVM 執行圖形。通過自動將算子融合在一起，讓 TVM 生成融合內核，在圖形節點內和圖形節點之間應用優化。對mobilenet ImageNet的工作量進行了基準測試，并討論了以下結果：

TVM在速度方面可以優于baseline基線方法。更有趣的是，內核融合帶來了額外的加速。值得一提的是，TVM 能夠自行生成所有優化的 GPU 內核，無需依賴 CuDNN 等外部庫。
開源努力
TVM最初是華盛頓大學保羅·艾倫學校計算機科學與工程系的研究項目。TVM 堆棧旨在支持DLPack，這是多個主要深度學習框架對張力數據結構的共識。收到了來自UW, AWS, Qiho 360, Facebook, HKUST, TuSimple, UCDavis, SJTU以及DMLC開源社區和DLPack倡議成員的早期貢獻。今后，該項目將遵循Apache open-source model開源模式，創建社區維護項目。
? Halide：TVM使用HalideIR作為數據結構，用于arithematic簡單化和低level lowering。HalideIR來自Halide。還從Halide那里學習，在TVM中實施lowing管道。
? Loopy：使用整數集分析及其循環轉換原始。
? Theano：symbolic scan算子的再現設計靈感。
源代碼Source code
? Github page頁面可以在這里找到： https://github.com/dmlc/tvm
? TVM 兼容DLPack，因此可以輕松支持采用該標準的框架，如 MXNet, PyTorch, Caffe2和微小 dnn。

總結

以上是生活随笔為你收集整理的用TVM在硬件平台上部署深度学习工作负载的端到端 IR 堆栈的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。