TVM 架构设计

TVM 架构设计

本文面向希望了解TVM体系结构和/或，积极参与项目开发的开发人员。

主要内容如下：

示例编译流程，概述了TVM将模型的高级概念，转换为可部署模块的步骤。

逻辑架构组件部分，描述逻辑组件。针对每个逻辑组件，按组件的名称进行组织。

可以随时查看，开发人员如何指导有用的开发技巧。

提供了架构的一些补充视图。检查一个单一的端到端编译流程，讨论关键的数据结构和转换。这个基于runtime的视图，主要关注运行编译器时，每个组件之间的交互。将检查代码库的逻辑模块及关系。设计了静态总体视图。

Example Compilation Flow

研究编译器中的一个示例编译流。下图显示了流程，在高层，包含几个步骤：

导入：前端组件将模型摄取到IRModule中，IRModule包含内部表示模型的函数集合。

转换：编译器将一个IRModule，转换成另一个功能上等价或近似等价的IRModule（如在量化的情况下）。许多转换都是独立于目标（后端）的。还允许target影响转换管道的配置。

目标转换：编译器将IRModule（codegen），转换为目标指定的可执行格式。目标转换结果封装为runtime.Module，可以在目标runtime环境中导出、加载和执行。

Runtime执行：用户将runtime.Module，在支持的runtime环境中，运行编译的函数。

Key data structures

设计和理解复杂系统的最佳方法之一，识别关键数据结构和操作（转换）这些数据结构的API。一旦确定了关键的数据结构，可以将系统分解成逻辑组件，这些组件要么定义关键数据结构的集合，要么定义数据结构之间的转换。

IRModule是整个堆栈中，使用的主要数据结构。IRModule（中间表示模块）包含一组函数。支持函数的两个主要变体。

relay::Function函数是高级函数程序表示。一个relay功能通常对应于端到端模型。可以查看relay功能作为一个计算图，对控制流、递归和复杂的数据结构，有额外的支持。

tir::PrimFunc是一个低级程序表示，包括循环嵌套选择、多维加载/存储、线程和向量/张量指令在内的元素。通常用来表示执行模型中（可能是融合）层的算子程序。

在编译过程中，一个relay函数，可以降为多个tir::PrimFunc函数，与一个调用这些tir::PrimFunc函数的顶层函数。

Transformations

已经介绍了关键的数据结构，谈谈转换。每种转换，都可以达到以下目的之一：

优化：将一个程序转换成一个等效的，可能更优化的版本。

降低：将程序转换为更接近目标的低级表示。

Relay/转换包含优化模型的过程集合。这些优化包括常见的程序优化，如常量折叠和死代码消除，及张量计算特定的过程，如布局转换和缩放因子折叠。

将端到端部的Relay（例如，MobileNet），端到端部的优化（融合操作），称这些函数段。这个过程帮助将原始问题，分为两个子问题：

每个子功能的编译和优化。

整体执行结构：需要对生成的子函数，执行一系列调用，执行整个模型。

使用低级tir，编译和优化每个子函数。对于特定的目标，可以直接进入目标转换阶段，使用外部代码生成器。

有几种不同的方法（在Relay/后端），处理对整个执行问题的调用。对于具有已知形状，没有控制流的简单模型，可以降低到将执行结构，存储在图形中的图形runtime。支持用于动态执行的虚拟机后端。调度支持提前编译，将高级执行结构，编译成可执行的和生成的原始函数。所有这些执行模式，都被一个统一的runtime.Module接口，将在后面部分讨论。

tir/转换包含tir级功能的转换过程。许多tir pass的目的是降低。有一些过程，可以将多维访问，变为一维指针访问，将内部函数扩展为特定于目标的内部函数，及修饰函数入口，满足runtime调用约定。有一些优化过程，如访问索引简化和死代码消除。

许多低级优化，可以