TVM设计与构架构建

TVM设计与构架构建

本文档适用于希望了解TVM体系结构和/或在项目上进行积极开发的开发人员。该页面的组织如下：

• 实例编译流程Example Compilation Flow描述TVM把一个模型的高级描述到可部署模块的步骤。
• “逻辑体系结构组件” Logical Architecture Components部分描述了逻辑组件。针对每个逻辑组件的特定内容，按组件名称组织。
• 开发人员操作手册，以获取有用的开发技巧。

本文提供了一些体系结构的补充视图。首先，回顾一个端到端的编译流程，并讨论关键的数据结构和转换。这个基于运行时runtime的视图着重于运行编译器时每个组件的交互。然后，将回顾代码库的逻辑模块及其关系。这部分提供了设计的静态总体视图。

编译流程示例

本节将研究编译器中的示例编译流程。下图显示了流程。从高层次上讲，包含几个步骤：

• 导入：前端组件将模型提取到IRModule中，该模块包含内部表示模型的函数的集合。
• 转换：编译器将IRModule转换为另一个功能上等效或近似等效的（例如，在量化的情况下）IRModule。许多转换都是独立于目标（后端）的。还允许目标影响转换管道的配置。
• 目标翻译：编译器将IRModule转换（代码生成）为目标指定的可执行格式。目标翻译结果被封装为runtime.Module，可以导出，加载和执行对目标运行时runtime环境。
• 运行时runtime执行：用户负载背了runtime.Module在支持的运行环境和运行编译功能。

关键数据结构

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

IRModule是整个堆栈中使用的主要数据结构。IRModule（中间表示模块）包含函数的集合。当前，支持两种主要的功能变体。

• relay :: Function是高级功能代码表示。relay.Function通常对应于端到端模型。可以将relay作为计算图来查看，并额外支持控制流，递归和复杂的数据结构。
• tir :: PrimFunc是一种低级代码表示，其中包含一些元素，包括循环嵌套选择，多维加载/存储，线程和向量/张量指令。通常用于表示执行模型中（可能是融合的）层的算子代码。

在编译期间，可以将relay函数降低为多个tir :: PrimFunc函数和一个调用这些tir :: PrimFunc函数的顶级函数。

转换

已经涵盖了关键数据结构，来谈谈转换。每个转换可以满足以下目的之一：

• 优化：将代码转换为等效的，可能更优化的版本。
• 降维：将代码转换为更接近目标的较低层表示。

中继/转换包含优化模型的通道集合。这些优化包括常见的代码优化（例如恒定折叠和死代码消除）以及张量计算特定的遍历（例如布局转换和缩放因子折叠）。

在中继优化管道的末端附近，运行pass（FuseOps）将端到端功能（例如MobileNet）划分为子功能（例如conv2d-relu）段。称这些功能为段。此过程可将原始问题分为两个子问题：

• 每个子功能的编译和优化。
• 总体执行结构：需要对所生成的子函数进行一系列调用，以执行整个模型。

使用低级的Tir波段来编译和优化每个子功能。对于特定目标，也可以直接进入目标平移并使用外部代码生成器。

有几种不同的方法（在中继/后端）来处理对整个执行问题的调用。对于具有已知形状且无控制流的简单模型，可以降低到将执行结构存储在图中的图运行时runtime。还支持虚拟机后端进行动态执行。最后，计划支持提前编译，该编译将高级执行结构编译为可执行文件和生成的原始函数。所有这些执行模式都由统一的runtime.Module 接口封装 ，将在本文的后面部分中进行讨论。

tir / transform包含用于TIR级别功能的转换过程。许多Tir通道的目的是降维。例如，可以通过多种途径将多维接入到一维指针访问，将内在函数扩展为特定于目标的内在函数，并修饰函数条目，以满足运行时runtime调用约定。当然，也有一些优化过程，例如简化访问索引和消除无效代码。

LLVM，CUDA C和其它目标编译器，可以在目标处理许多低级优化。结果，将低级优化（例如寄存器分配）留给了下游编译器，只专注于未涵盖的优化。

搜索空间和基于学习的转换

到目前为止，描述的转换过程是确定性的且基于规则的。TVM堆栈的一个设计目标是支持针对不同硬件平台的高性能代码优化。为此，将需要研究尽可能多的优化选择，包括但不限于多维张量访问，循环切片行为，特殊的加速器内存层次结构和线程化。

很难定义做出所有选择的试探法。相反，将采用基于搜索和学习的方法。首先定义可以用来转换代码操作的集合。示例操作包括循环转换，内联，向量化。称这些操作为调度原语。调度原语的集合定义了，可以对代码进行的可能优化的搜索空间。然后，系统搜索不同的可能调度序列，以选择最佳调度组合。搜索过程通常以机器学习算法为原则。

搜索完成后，可以为（可能是融合的）算子记录最佳调度顺序。然后，编译器可以仅查找最佳调度序列，并将其应用于代码。值得注意的是，此调度应用代码布局完全类似于基于规则的变换，能够与传统流程共享相同的接口约定。

使用基于搜索的优化来处理最初的Tir函数生成问题。此模块部分称为AutoTVM（auto_scheduler）。随着继续开发TVM堆栈，希望将基于学习的转换扩展到更多领域。

目标转换

目标转换阶段将IRModule转换为相应的目标可执行格式。对于x86和ARM等后端，使用LLVM IRBuilder来构建内存中的LLVM IR。还可以生成诸如CUDA C和OpenCL之类的源代码级语言。最后，支持通过外部代码生成器将Relay函数（子图）直接转换为特定目标。重要的是，最终代码生成阶段应尽可能轻巧。绝大部分的转换和降维都应在目标转换之前进行。

还提供了一个Target结构来指定编译目标。目标翻译阶段之前的转换也可能受到目标的影响-例如，目标的向量长度会改变向量化行为。

运行时runtime执行

TVM运行时runtime的主要目标是提供一个最小的API，以使用他们选择的语言（包括Python，C ++，Rust，Go，Java和JavaScript）加载和执行已编译的工件。下面的代码片段显示了Python中的这样一个示例：

import tvm

# Example runtime execution program in python, with type annotated

mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")

arr: tvm.runtime.NDArray = tvm.nd.array([1, 2, 3], ctx=tvm.gpu(0))

fun: tvm.runtime.PackedFunc = mod["addone"]

fun(a)

print(a.asnumpy())

tvm.runtime.Module封装编译结果。runtime.Module包含一个GetFunction方法，用于按名称获取PackedFuncs。

tvm.runtime.PackedFunc是两个生成的函数的类型清理的函数接口。runtime.PackedFunc可采用以下类型的参数并返回值：POD类型（int，float），字符串，runtime.PackedFunc，runtime.Module，runtime.NDArray以及runtime.Object的其它子类。

tvm.runtime.Module并且tvm.runtime.PackedFunc是将运行时runtime模块化的强大机制。例如，要在CUDA上获得上述addone函数，可以使用LLVM生成主机端代码以计算启动参数（例如线程组的大小），然后从CUDAModule调用另一个由PackedFunc支持的PackedFunc。 CUDA驱动代码API。相同的机制可用于OpenCL内核。

上面的示例仅处理简单的addone函数。下面的代码段给出了使用同一接口执行端到端模型的示例：

import tvm

# Example runtime execution program in python, with types annotated

factory: tvm.runtime.Module = tvm.runtime.load_module("resnet18.so")

# Create a stateful graph execution module for resnet18 on gpu(0)

gmod: tvm.runtime.Module = factory["resnet18"](tvm.gpu(0))

data: tvm.runtime.NDArray = get_input_data()

# set input

gmod["set_input"](0, data)

# execute the model

gmod["run"]()

# get the output

result = gmod["get_output"](0).asnumpy()

主要优点是，runtime.Module和runtime.PackedFunc足以封装算子级代码（例如addone）以及端到端模型。

总结与讨论

总之，编译流程中的关键数据结构为：

• IRModule：包含relay.Function和tir.PrimFunc
• runtime.Module：包含runtime.PackedFunc

编译的大部分内容是关键数据结构之间的转换。

• relay / transform和tir / transform是基于规则的确定性转换
• auto_scheduler和autotvm包含基于搜索的转换

最后，编译流程示例只是TVM堆栈的典型用例。将这些关键数据结构和转换开放给python和C ++ API。因此，除了感兴趣的数据结构从numpy.ndarray更改为tvm.IRModule之外，可以像使用numpy，一样使用TVM。以下是一些用例示例：

• 使用python API直接构造IRModule。
• 组成一组自定义的转换（例如，自定义量化）。
• 使用TVM的python API直接操作IR。

逻辑架构组件

TVM体系结构图

上图显示了项目中的主要逻辑组件。阅读以下各节，以获取有关组件及其关系的信息。

tvm/support

支持模块包含最常用的基础构建实用代码，例如通用竞技场分配器，套接字和日志记录。

tvm/runtime

运行时runtime是TVM堆栈的基础。提供了加载和执行已编译工件的机制。运行时runtime定义了一组稳定的标准C API，以与诸如Python和Rust的前端语言进行接口。

除了runtime :: PackedFunc之外，runtime :: Object是TVM运行时runtime中的主要数据结构之一。它是带有类型索引的引用计数基类，以支持运行时runtime类型检查和向下转换。目标系统允许开发人员向运行时runtime引入新的数据结构，例如数组，映射和新的IR数据结构。

除了部署用例之外，编译器本身还大量使用TVM的运行时runtime机制。所有的IR数据结构都是runtime :: Object的子类，可以从Python前端直接访问和操作它们。使用PackedFunc机制将各种API公开给前端。

在运行时runtime的子目录（例如runtime / opencl）中定义了对不同硬件后端的运行时runtime支持。这些特定于硬件的运行时runtime模块定义，用于设备内存分配和设备功能序列化的API。

runtime / rpc为PackedFunc实现RPC支持。可以使用RPC机制将交叉编译的库发送到远程设备，并确定执行性能的基准。rpc基础架构支持从广泛的硬件后端收集数据，以进行基于学习的优化。

• TVM运行系统
• 调试器
• 将VM放入TVM：中继虚拟机
• 模块序列化简介

tvm/node

节点模块在runtime :: Object的基础上为IR数据结构添加了其它功能。主要包括反射，序列化，结构等效和散列。

使用了节点模块，可以通过在Python中的名称，直接访问TVM的IRNode的任何字段。

x = tvm.tir.Var("x", "int32")

y = tvm.tir.Add(x, x)

# a and b are fields of a tir.Add node

# we can directly use the field name to access the IR structures

assert y.a == x

可以将任意IR节点序列化为JSON格式，然后将其加载回。保存/存储和检查IR节点的能力，为使编译器更易于访问提供了基础。

tvm/ir

在TVM / IR文件夹包含跨所有IR功能变体的统一的数据结构和接口。tvm / ir中的组件由tvm / relay和tvm / tir共享，包括

• ir模块
• 类型
• PassContext和Pass
• 算子

功能的不同变体（例如relay.Function和tir.PrimFunc），可以共存于IRModule中。尽管这些变体可能不具有相同的内容表示，但是它们使用相同的数据结构来表示类型。因此，使用相同的数据结构来表示这些变量的功能（类型）名称。一旦明确定义了调用约定，统一类型系统就允许一个函数变换调用另一个函数。这为将来的跨功能变量优化打开了大门。

还提供了一个统一的PassContext用于配置传递操作，并提供了通用的复合传递来执行传递管道。以下代码段给出了PassContext配置的示例。

# configure the behavior of the tir.UnrollLoop pass

with tvm.transform.PassContext(config={"tir.UnrollLoop": { "auto_max_step": 10 }}):

# code affected by the pass context

Op是表示所有系统定义的原始算子/内部函数的通用类。开发人员可以向系统注册新的Op以及它们的其它属性（例如Op是否是元素化的）。

• Pass Infrastructure

tvm/target

目标模块包含将IRModule转换为目标runtime.Module的所有代码生成器。还提供了描述目标的通用Target类。

通过查询目标中的属性信息和注册到每个目标id（cuda，opencl）的内置信息，可以根据目标定制编译管道。

tvm/ir

TIR包含低级代码表示的定义。使用tir :: PrimFunc表示可以通过TIR传递转换的函数。除IR数据结构外，tir模块还通过公共Op注册表，以及tir / transform中的转换，传递定义了一组内置的内在函数及其属性。

tvm/arith

此模块与TIR紧密相关。低级代码生成中的关键问题之一是分析索引的算术属性-正则性，变量边界以及描述迭代空间的整数集。arith模块提供了一组进行（主要是整数）分析的工具。TIR pass可以使用这些分析来简化和优化代码。

tvm/te

te代表“张量表达式”。这是一个特定领域的语言模块，允许通过编写张量表达式来快速构建tir :: PrimFunc变体。重要的是，张量表达式本身不是可以存储到IRModule中的自包含函数。相反，它是IR的一个片段，可以将其拼接在一起以构建IRModule。

te / schedule提供了一组调度原语，以控制所生成的功能。将来，可能会将其中的一些调度组件引入tir :: PrimFunc本身。

• InferBound Pass
• 混合前端开发人员指南

• InferBound Pass
• Hybrid Frontend Developer Guide

tvm/topi

可以针对每个用例直接通过TIR或张量表达式（TE）构造算子。 topi（张量算子清单）提供了一组预定义的算子（在TE或TIR中），由numpy定义并在常见的深度学习工作负载中找到。提供了一组公共时间表模板，以在不同目标平台上获得高性能的实现。

tvm/relay

继电器是用于表示完整模型的高级功能性IR。在relay.transform中定义了各种优化。Relay编译器定义了多种语言，每种语言旨在支持特定的优化样式。值得注意的是QNN（用于导入预量化模型），VM（用于降级为动态虚拟机），内存（用于内存优化）。

• 中继IR简介
• 中继算子策略
• 转换布局通道

• Introduction to Relay IR
• Relay Operator Strategy
• Convert Layout Pass

tvm/autotvm

AutoTVM和AutoScheduler都是自动进行基于搜索的代码优化的组件。主要包括：

• Cost models和特征提取。
• 一种记录格式，用于存储计划基准结果以进行cost model构建。
• 一组有关代码转换的搜索策略。

• Cost models and feature extraction.
• A record format for storing program benchmark results for cost model construction.
• A set of search policies over program transformations.

自动化代码优化仍然是活跃的研究领域。试图对设计进行模块化，以便研究人员可以通过Python绑定，快速修改组件或自己的应用算法，自定义搜索并从Python绑定中插入其算法。

• 基准性能日志格式Benchmark Performance Log Format

前端

前端将来自不同框架的模型，存放在TVM堆栈中。 tvm.relay.frontend是模型提取API的命名空间。

• TensorFlow前端