TVM：简介

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）部署。