microTVM 设计文档

背景

TVM 是一个模型部署框架，已经在传统操作系统上展示出良好的性能，适用于各种模型。鉴于 TVM 的分层编译方法，将其扩展到目标裸机设备是一种自然的选择。虽然大部分编译流程在这样的设备上不需要改变，但运行时不能依赖于：

• 虚拟内存 (Virtual Memory)，以及任何系统提供的 malloc （动态内存分配）系统。此外，裸机设备通常具有非常有限的内存（以 KB 为单位）。因此，针对此类平台设计的库通常需要更谨慎地使用内存，并在不使用时释放内存。

• 传统的操作系统抽象，例如 文件、库 和 内核函数。一些项目实现了对它们的支持，但它们并不是标准。

• 支持除 C 以外的编程语言。

这样的变化需要与传统操作系统上通常使用的 TVM C++ 运行时不同的方法。

典型的使用

本节讨论了我们对“典型”微型 TVM 使用情况的愿景。为了达到这种典型的使用情况，使用的每个组件都旨在设计灵活，但是这个统一的愿景有助于激发设计的每个部分的加入。

这个过程的各个部分如下所述：

1. 模型导入。用户导入现有模型或描述新模型到 TVM，生成 Relay 模块。

2. 模型变换。用户可以对模型应用量化等变换。每个变换后，用户应该仍然有 Relay 模块。

3. 编译 （调度和代码生成）。TVM 通过为每个 Relay 算子分配调度和调度配置来将每个算子实现为 Tensor IR。然后，为每个算子生成代码（C 源代码或已编译的对象）。

4. 集成。生成的代码与 TVM C 运行时库一起集成到用户提供的二进制项目中。在某些情况下（例如当该项目在多个 SoC/开发板上标准化时），此过程会自动处理。

5. 部署。构建项目并将剩余的固件(firmware)二进制文件刷写到设备上。模型推理由 TVM 使用设备上的 RPC 服务器驱动，或者在设备上使用设备上的 Graph 执行器。

设计目标

microTVM 旨在实现以下设计目标：

1. 可移植代码。microTVM 可以将任何 Relay 模型翻译为 C 代码，只需使用 C 标准库即可编译。

2. 最小化开销。microTVM 生成面向目标的高度优化的代码。应尽可能消除来自运行时的开销。

3. 易于访问的代码。microTVM 将 C 源代码视为一流的输出机制，以便固件工程师更容易理解和调整。

概述

microTVM 需要对 TVM 编译器堆栈的所有 level 进行更改。以下子部分在高层次上列出了这些更改，后续部分详细讨论了具体内容。

建模目标平台

TVM 的基于搜索的优化方法允许其在实验结果方面避免大部分系统级目标建模。然而，一些建模是必要的，以确保 TVM 比较 apples-to-apples 的搜索结果，并避免在搜索期间尝试为目标编译无效代码而浪费时间。

microTVM 对目标的这些部分进行建模：

• 通过 -mcpu 和 -march 目标标志使用的 CPU。

• 通过目标设备的组件来判断加速器的有无（目前只能表示缺少加速器，但是这种机制应该会得到很好的扩展）。

microTVM 旨在未来模拟目标的这些部分：

• 内存，被建模为一组不相交的内存空间，每个空间具有标签（label）、大小（size）和预取/刷新（prefetch/flush）行为。有些内存可能与加速器共享。

• 目标运行时配置（即时钟树配置（clock tree configuration）、时钟速度（clock speed）等）。这只是为了提供 AutoTVM 调度键，而不是用于任何其他用途。

目前，TVM 不打算建模：

• 缓存的大小、类型或关系，除了预取(prefetching)和缓存刷新(cache flushing)。

TVM microTVM 的目标设备

在编译过程中，一个核心的数据结构是 tvm::target::Target 类。TVM 使用 Target 来决定启用哪些 TIR 调度并配置代码生成器。Target 类还应该唯一地标识特定算子的生成代码，因为自动调优日志使用它来排名性能（但请参见未来工作）。

目标当前被表示为结构类似于命令行参数的字符串。下面是目标的示例：

c -keys=arm_cpu -mcpu=cortex-m7 -model=stm32f746xx

对于 microTVM 相关的部分有：

• Code generator (llvm or c)

• -mcpu=cortex-m7：由 TOPI 使用以启用 Cortex-M 调度，并且当选择 C 源代码生成器时，会作为注释包含在输出中，以帮助识别代码并配置下游 C 编译器。

microTVM 的运行时和执行器配置

在使用 microTVM 时，重要的是要使用 C 运行时（Runtime('crt')），这是在微型设备上表现最佳的运行时，而不是更动态的 C++ 运行时。除此之外，还有两个执行器可与 C 运行时结合使用：

• Executor("aot") - 预先编译（AOT）执行器会将网络预编译为可运行的函数，您可以直接将其添加到您的微应用程序中。

• Executor("graph", {"link-params": True}) - Graph 执行器提供了网络的 JSON 表示形式，并需要生成 C 运行时的系统库以在函数注册表中查找函数（Runtime("crt", {"system-lib": True})）。{"link-params":True} 允许将参数链接到生成的文件中，而不是从外部提供。

这些是在构建运行时模块时指定的：relay.build(..., runtime=..., executor=...)。

编写 microTVM 的调度策略

对于在 CPU 上调度的运算，microTVM 最初计划利用专用指令和外部（即手动优化的）函数来实现良好的性能。在 TVM 中，这种方法通常是通过张量化来实现的，其中 TVM 将计算分解成小块，TIR 外部函数加速每个小块。

目前，TVM 使用 tir.call_extern 来容纳这两种方法。首先，在调度中附加 pragma，以定义便携式 C 中的 extern 函数。

sched[output].pragma(n, "import_c", "void call_asm(int32_t* a, int32_t* b) { /* ... */ }")

接下来，使用 tensorize 来分解计算。

sched[output].tensorize(owi, gemm)

这种方法有几个注意事项，所有这些注意事项都可以通过将生成的代码链接到外部库来解决：

• 内联汇编(Inline assembly)是与编译器相关的。虽然 Clang 和 GCC 已经统一了一种语法，但这种语法可能无法适用于其他编译器。SDK 通过根据使用的编译器条件性地包含头文件来解决这个问题。然而，采用这种方法意味着生成的代码需要额外的编译器标志（即 -Isystempath/to/header）。

• 从生成的代码中引用帮助函数可能会有所帮助（例如，内联常见的手动优化汇编序列）。

• 最后，调用的 extern 函数可能完全由外部库编写。如果这些函数可以完全内联，则此警告与之前相同。如果不是，则需要编译额外的 C 代码，并将其链接到算子。

目前，microTVM 假定所有合适的调度都可以被编译。这意味着用户提供的项目（请参见下一节）必须包括由生成的代码使用的所有库。当不使用自动调优时，TVM 会随机选择备用调度，因此需要支持所有库。当使用自动调优时，TVM 选择最佳性能的调度，因此只需要该库。目前没有办法强制 TVM 在自动调优日志之外选择特定的调度，但这将是很好的补充。

最后，当使用 llvm 后端时，该过程类似，只是 LLVM bitcode 包含在生成的代码中（使用 import_llvm pragma）。LLVM bitcode 提供了一种调用内联汇编的便携式方法。然而，从 LLVM bitcode 中调用外部 C 函数可能更加复杂，而辅助函数当然不易从 LLVM bitcode 中使用。

执行模型

传统上，TVM 编译器输出以下三部分：

1. 如上所述，模型算子的实现；

2. 编码为 JSON 的模型执行 graph；以及

3. 简化的参数。

为了正确地执行模型，graph 执行器需要在内存中重构计算图，加载参数，然后按正确的顺序调用算子实现。

microTVM 支持两种方法来实现这一点：

1. 主机驱动。Graph Executor 可以在主机上运行，并通过使用具有类 UART 传输的 RPC 链接向设备发出命令来执行。

2. 独立模式。提供了 C Graph Executor，可以在设备上编译，但它的内存效率并不是特别高。这种方式实现了独立执行而无需连接到任何主机。

主机驱动模式旨在在设备上进行模型实验，类似于 AutoTVM，使用 RPC 服务器驱动设备上的计算。独立模式则适用于部署。

主机驱动模式执行。

在主机驱动模式执行中，固件二进制文件如下：

1. 从 TVM 生成的算子实现。

2. TVM C 运行时。

3. 特定于 SoC 的初始化。

4. TVM RPC 服务器。

5. （可选）简化参数。

该固件镜像被烧录到设备上，并在主机上创建了 GraphExecutor 实例。GraphExecutor 通过 UART 发送 RPC 命令来驱动执行：

独立执行

在独立执行中，GraphExecutor 被实例化在设备上：

microTVM 固件

现在可以讨论 microTVM 固件应该如何表现了。对于模型执行策略来说，重要的任务是配置 SoC 以匹配其在生产中的表现方式。microTVM 认为这是与项目和 SoC 相关的任务。无论是针对 AutoTVM、主机驱动的模型推理还是独立部署，用户都应该提供一个项目，其 main() 函数应该执行以下操作：

1. 将 SoC 配置以匹配部署性能。

2. 初始化 TVM C 运行时。

当为主机驱动的推理或 AutoTVM 进行配置时，剩余的任务是明确定义的：

3. 初始化传输（如 UART），以便与 TVM RPC 服务器一起使用。

4. 启动 TVM RPC 服务器。

在配置为独立部署时，固件需要：

1. 通过调用 runtime.SystemLib PackedFunc 来实例化系统库。

2. 实例化 GraphExecutor，传递系统库模块。

3. 根据需要配置参数和输入。

4. 运行模型。

microTVM 二进制文件的组成部分。

总之，microTVM 固件二进制文件必须包含以下几个部分：

1. 由 TVM 生成的算子符实现。

2. 由 TVM 提供的 TVM C 运行时库，作为静态库供使用。

3. SoC 初始化，由用户提供。

对于基于主机的模型执行，固件还需要：

4. TVM RPC 服务器库。

对于独立模型执行，固件还需要：

4. 由 TVM 作为静态库提供的 TVM C GraphExecutor 库。

5. 剩余的编译器输出（简化的参数和 Graph JSON）。

自动化构建流程

一旦代码生成完成，tvm.relay.build 将返回 tvm.runtime.Module，用户可以将生成的 C 源代码或二进制对象保存到 .c 或 .o 文件中。从这一点上来说，TVM 理论上可以退后一步，用户可以将代码编译和运行分开进行。

然而，对于 AutoTVM，TVM 需要一些自动化的流程来处理以下任务：

1. 将算子实现、TVM C 运行库和 TVM RPC 服务器库集成到包含用户提供的 SoC 初始化的固件项目中。

2. 构建生成的项目。

3. 将构建好的固件烧录到（特定的）连接设备上。

4. 确定 TVM 用于驱动远程执行的串口或其他传输方式。

目前，TVM 希望用户提供 tvm.micro.Compiler、tvm.micro.Flasher 和 tvm.micro.Transport 接口的实现。TVM 随后：

1. 将每个部分分别构建为库。

2. 将这些库构建成二进制固件映像。

3. 将固件映像烧录到已连接设备上。

4. 打开串口作为 RPC 服务器传输。

选择这种设计是为了减少 microTVM 的构建时间（只需对候选算子实现构建一次共同的库）。实际上，这些项目非常小，编译相对较快。与 TVM 更紧密的构建集成相比，这种更紧密的构建集成增加的复杂性可能不值得。未来的设计将把构建任务合并到一个步骤中，并缩小接口以提供更好的集成。

测量算子性能。

TVM C 运行时依赖于用户提供的函数在设备上测量时间。用户应该实现 TVMPlatformTimerStart 和 TVMPlatformTimerStop。这些函数应该测量 clock 时间，因此在实现这些函数时有一些需要注意的地方：

1. 如果在计算过程中 CPU 可以停止或休眠（例如在加速器上进行计算），那么不应该使用循环计数器，因为这些计数器在 CPU 休眠时往往停止计数。

2. 这些函数的粒度可以根据需要进行放宽，以扩展计时器设备的范围。但是，如果粒度太粗，可能会使用次优调度。

3. 如果计时器溢出，应该引发异常。

4. 除非绝对必要，计时器不应中断计算。这样做可能会影响结果的准确性。

5. 将输出校准到 clock 是理想的，但可能太繁琐。未来的 PR 可以通过例如将内部振荡器与外部晶体等参考进行比较，实现对平台计时器的一些特征描述。

未来的工作

提前运行时（Ahead-of-Time Runtime）

Graph Executor 的局限性在于解析 JSON 需要大量的内存开销。当前的实现对 microTVM 的动态内存使用量做出了重大贡献，限制了其实用性。提前运行时可以避免对 Graph JSON 进行任何解析，并通过生成 C 代码直接调用生成的算子实现来提高推理速度，而不是依赖于 Graph Executor 的数据驱动方法。

Memory Planning

目前的内存规划器仅试图限制中间张量的 TVMBackendDeviceAlloc() 调用次数。由于临时缓冲区的大小可能会有很大差异，并且规划器会将内存分配合并在彼此的 16 倍范围内，因此这种策略通常会导致高峰值内存使用率。

异构执行

新的 Cortex-M SoC 可以包含多个 CPU 和内置的 ML 加速器。

自动调优目标设备

如前所述，