TVM设计与构架构建

TVM设计与构架构建

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

• 实例编译流程Example Compilation Flow，描述TVM把一个模型的高级描述到可部署模块的步骤。

• “逻辑体系结构组件” Logical Architecture Components部分描述了逻辑组件。针对每个逻辑组件的特定内容，按组件名称组织。

• 开发人员操作手册，以获取有用的开发技巧。

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

编译流程示例

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

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

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

• 目标翻译：编译器将IRModule转换（代码生成）为目标指定的可执行格式。目标翻译结果，封装为runtime.Module，可以导出，加载和执行，对目标运行时runtime环境。

• 运行时runtime执行：用户load了runtime.Module，在支持的运行环境和运行编译功能。

 

 

关键数据结构

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

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

• relay :: Function是高级功能代码表示。relay.Function通常对应于端到端模型。可以将relay作为计算图，额外支持控制流，递归和复杂的数据结构。

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

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

转换

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

• 优化：将代码转换为等效的，可能更优化的版本。

• 降维：将代码转换为更接近目标的，较低层表示。

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

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

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

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

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

有几种不同的方法（在Relay/后端），处理对整个执行问题的调用。对于具有已知形状，无控制流的简单模型，可以降低到将执行结构，存储在图中的图运行时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"

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".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：Relay虚拟机

• 模块序列化简介

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，用于配置传递算子，提供了通用的复合pass，执行pass管道。以下代码段给出了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

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

• Relay IR简介

• Relay算子策略

• 转换布局通道

• 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前端