TVM：簡介概述

Apache TVM 是一個用于 CPU、GPU 和機器學習加速器的開源機器學習編譯器框架。它旨在使機器學習工程師能夠在任何硬件后端上高效地優化和運行計算。本教程的目的是通過定義和演示關鍵概念，引導您了解 TVM 的所有主要功能。新用戶應該能夠從頭到尾完成本教程的工作，并能夠操作 TVM 進行自動模型優化，同時對 TVM 架構及其工作原理有基本的了解。

內容

1. 介紹
2. 安裝 TVM
3. 使用命令行界面編譯和優化模型
4. 使用 Python 接口編譯和優化模型
5. 使用張量表達式使用運算符
6. 使用模板和 AutoTVM 優化 Operator
7. 使用無模板的 AutoScheduler 優化 Operator
8. 交叉編譯和遠程過程調用 (RPC)
9. 為 GPU 編譯深度學習模型

TVM 和模型優化概述

下圖說明了一個模型在使用 TVM 優化編譯器框架進行轉換時所采取的步驟。

1. 從Tensorflow、Pytorch或Onnx等框架導入模型。導入器（importer）層是 TVM 可以從其他框架（如 Tensorflow、PyTorch 或 ONNX）中攝取模型的地方。隨著我們不斷改進開源項目，TVM 為每個前端提供的支持水平各不相同。如果在將模型導入 TVM 時遇到問題，可以嘗試先將其轉換為 ONNX。

2. 轉換為Relay，Relay 是TVM 的高層模型語言。已導入 TVM 的模型由 Relay 來表示。Relay 是一種用于神經網絡的函數式語言和中間表示 (IR)。它支持：

   • 傳統的數據流樣式表示
   • functional-style scoping，let-binding，這使的它成為功能齊全的可微分語言
   • 允許用戶混合兩種編程風格

   Relay 進行圖（graph）級別的優化傳遞來優化模型。

3. 更低層的 張量表達式 (TE，tensor expression) 表示。降級（lowing）是指將較高層的表示轉換為較低層的表示。經過高層（high-level）優化后，Relay 通過 FuseOps pass 將模型劃分為許多小子圖，并將子圖降級為 TE 表示。張量表達式 (TE) 是一種用于描述張量計算的特定領域語言。TE 還提供了多個schedule 原語來指定低級循環優化，例如平鋪 (tiling)、矢量化 (vectorization)、并行化 (parallelization)、展開 (unrolling) 和融合 (fusion) 。為了幫助將 Relay 表示轉換為 TE 表示的過程，TVM 還包括一個Tensor Operator Inventory (TOPI)，它具有常見張量算子（例如，二維卷積 (conv2d)、轉置 (transpose) ）的預定義模板。

4. 使用自動調整模塊 AutoTVM 或 AutoScheduler 搜索最佳的 schedule。schedule 為 TE 中定義的算子或子圖指定低層循環優化。Auto-tuning 搜索最佳 schedule，并將其與成本模型和端側測度進行比較。TVM 中有兩個 Auto-tuning 模塊。

   • AutoTVM：基于模板（template-based）的 auto-tuning 模塊。它通過某種搜索算法在用戶定義的模板中找到可調節參數的最佳值。對于常見的算子，TOPI中已經提供了他們的模板。
   • AutoScheduler（又名 Ansor）：一個無模板（template-free）的 auto tuning 模塊。它不需要預定義的 schedule 模板。而是通過分析計算的定義來自動生成搜索空間。然后在生成的搜索空間中搜索最佳 schedule。

5. 選擇模型編譯的最佳配置。調優后，auto-tuning 模塊生成JSON格式的調優記錄。此步驟為每個子圖選擇最佳 schedule。

6. 更低層的 Tensor Intermediate Representation (TIR)，它是 TVM 的低層中間表示。根據 tuning 一步選擇最佳配置后，每個 TE 子圖被降級到 TIR 并進行低層優化。接下來，優化過的 TIR 被降級到硬件平臺的目標編譯器。這是生成可部署到生產中的優化模型的最終代碼生成階段。TVM 支持多種不同的編譯器后端，包括：

   • LLVM，可以針對任意微處理器架構，包括標準 x86 和 ARM 處理器、AMDGPU 和 NVPTX 代碼生成，以及 LLVM 支持的任何其他平臺。
   • 專用編譯器，例如NVIDIA 的編譯器 NVCC。
   • 嵌入式和專用 target，通過 TVM 的BYOC（Bring Your Own Codegen）框架實現。

7. 編譯成機器碼。在最后，針對特定編譯器的生成代碼可以降低為機器碼。

   TVM 可以將模型編譯為可鏈接的對象模塊，然后可以由輕量級 TVM Runtime 運行，該 Runtime 提供 C API 來動態加載模型，以及其他語言（如 Python 和 Rust）的入口點。TVM 還可以將 Runtime 與單個包中的模型相結合，從而構建捆綁（bundled）部署。