将VM放入TVM：中继虚拟机

最新推荐文章于 2024-03-17 23:15:10 发布

牛牛存

最新推荐文章于 2024-03-17 23:15:10 发布

阅读量1.2k

点赞数 1

分类专栏： tvm TVM Relay

本文链接：https://blog.csdn.net/weixin_42164269/article/details/104291655

版权

TVM Relay 同时被 2 个专栏收录

25 篇文章 21 订阅

订阅专栏

tvm

14 篇文章 6 订阅

订阅专栏

将VM放入TVM：Relay虚拟机

Relay是一种新的程序表示形式，它可以表示和优化大量的机器学习程序。不幸的是，通过支持一组更具表现力的程序，我们引入了一些新的执行挑战。

Relay的解释器可以执行完整的语言，但是有明显的限制，使其不适合生产部署。它被构造为通过AST遍历以执行程序的低效率解释器。这种方法在概念上很简单，但是效率很低，因为AST遍历很大程度上依赖于间接。

编译动态代码还面临其他挑战，例如动态调度和分配，完全动态张量形状以及控制流。解释器为此提供了简单的解决方案，但没有一个具有足够的说服力和被优化。

第二种执行机制是现有的图运行时。为了使Relay程序以此为目标，我们将它们的一小部分编译为旧的图形格式，并在运行时执行它们。图运行时提供了快速的执行体验，但仅适用于Relay程序的非常有限的子集。

另一种可选的但非标准的方法是Relay的提前编译器，该编译器将Relay程序编译到包含提前实现的共享库中。提前提供的编译器提供了令人信服的性能，但是难以扩展和检测，只能通过修改代码生成和优化机制来完成。

Relay虚拟机旨在成为平衡这些竞争方法的框架，从而提供一个动态执行环境，该环境可以通过灵活的扩展机制与其他方法（如提前编译）进行扩展，检测和集成。

虚拟机旨在在部署和执行Relay程序时在性能和灵活性之间取得平衡，而不会放弃TVM的优势。

虚拟机（VM）设计是编程语言和系统中经过广泛研究的领域，并且已经针对成熟和嵌入式编程语言进行了各种虚拟机设计。先前的语言VM设计已针对传统程序的执行配置进行了大规模定制。传统程序处理小的标量值，并且包含大量的低级指令。大量指令要求指令的执行和分发非常高效。在机器学习的上下文中，我们主要使用（相对）少量的高级指令来操纵张量值。机器学习程序的成本中心是在大量输入上进行昂贵的操作符调用，例如GEMM或卷积。由于ML程序显示的执行配置文件，标量VM中存在的微优化的重要性大大降低。

TVM为视觉模型提供了强大的支持，但我们希望发展以支持更广泛的模型。图运行时能够利用输入图的完全静态性质来执行积极的优化，例如完全静态分配和最佳内存重用。当我们介绍利用控制流，递归，动态形状和动态分配的模型时，我们必须改变执行的方式。用于Relay的虚拟机是很自然的选择。

本文档的其余部分概述了Relay虚拟机设计及其指令集。

设计

VM的设计着眼于简单性而不牺牲性能。为了实现这一目标，我们专注于设计张量VM而不是标量VM。

在张量VM设置中，我们优化了对象的廉价“分配”（通过尝试避免实际分配），静态片段的重用以及进行动态成形的能力（即锯齿张量）。

指令集

指令集和指令表示的选择是VM最关键的设计决策。指令的当前表示形式是带有操作码和数据有效载荷的带标记的联合。一个重要的设计决策是指令的抽象级别（RISC与CISC）以及它们如何获取数据（固定宽度指令编码与可变长度编码）。当前版本更接近CISC，具有诸如AllocTensor之类的复杂指令，并且由于形状是指令的一部分，因此其长度是可变的。当前指令集是非常高级的，并且大致对应于Relay中的高级操作。

RET

参数：

RegName dst
RegName result

将寄存器【result】中的对象返回到调用的寄存器【dst】。

InvokePacked

参数：Index packed_index Index arity Index output_size RegName* packed_args

调用由【packed_index】表示的打包函数。【arity 】和【output_size 】用于确认VM所需要的输入和输出的数目。【packed_args 】存储了参数寄存器的列表。注意【Index 】是【int64_t】，也将用于其他的指令。

AllocTensor

参数：

RegName dst
RegName storage
uint32_t ndim
int64_t* shape
DLDataType dtype

使用常量【shape】（存储在【shape】中）和【dtype 】从给定的存储块【storage】中分配张量值。结果保存到寄存器【dst】。

AllocTensorReg

参数：

RegName dst
RegName storage
RegName shape_register
DLDataType dtype

从给定的存储块（存储在【storage】）中分配恰当形状（存储在【shape_register】）和【dtype】，结果保存到寄存器【dst】。

AllocStorage

参数：

RegName dst
RegName size
RegName alignment
DLDataType dtype_hint

分配给定【size】，【alignment】和数据类型，【dtype_hint】的存储块。分配的存储块存储在寄存器中dst。

AllocADT

参数：

RegName dst
Index tag
Index num_fields
RegName* datatype_fields

使用【num_fields】条目从【datatype_fields】寄存器分配一个标记为【tag】的数据类型datatype_fields，结果保存到寄存器【dst】。

AllocClosure

参数：

RegName dst
Index clo_index
Index num_freevar
RegName* free_vars;

使用【clo_index】位置的VMFunction 作为其代码分配一个闭包，并使用中的【num_freevar】寄存器项分配一个闭包【free_vars】，结果保存到寄存器【dst】。

getfield命令

参数：

RegName dst
RegName object
Index field_index

从【object】中获取带有索引【field_index】的字段值，并保存结果保存到寄存器【dst】。

If

参数：

RegName test
RegName target
Index true_offset
Index false_offset

检查寄存器【test】处的对象是否等于【target】。如果相等，则相对跳动【true_offset】，否则相对跳动【false_offset】。

GetTag

参数：

RegName object
RegName dst

获取【object】寄存器中的ADT对象的对象标签，并保存结果到寄存器【dst】。

Fatal

虚拟机执行失败。

GOTO

参数：

Index pc_offset

相对无条件跳跃【pc_offset】。

Invoke

参数：

Index func_index

调用【func_index】处的函数，会消耗VMFunction的arity字段中包含的参数数量。

InvokeClosure

参数：

RegName closure
Index num_closure_args
RegName* closure_args

调用【closure】，消耗在闭包的VMFunction中声明的参数数量。

LoadConst

参数：

RegName dst
Index const_index

从常量池中加载【const_index】位置的常量，结果保存到寄存器dst。

LoadConsti

参数：

Index val
RegName dst

将常量整数【val】加载到寄存器【dst】，结果是0秩张量。

对象表示

我们利用对象协议来表示VM使用的对象。

目前，三种类型的对象，NDArray，ADT，和Closure对象，分别用于表示张量，元组/列表和闭合数据。可以分别在include / tvm / runtime / ndarray.h， include / tvm / runtime / vm.h和include / tvm / runtime / container.h中找到有关它们的更多详细信息。

堆栈和状态

Relay VM维护一个堆栈框架，其中包含有关如何恢复上一个调用的信息。在连续空间（虚拟寄存器文件）中为每一个函数分配寄存器。

我们跟踪一组已调用的Relay函数，一个指向其字节码的指针，一个指向字节码的偏移量（称为程序计数器）。

struct VirtualMachine {
  ...
  std::vector<VMFrame> frames;
  ...
  // Current function.
  size_t func_index;
  // Pointer into the current function's instructions.
  const Instruction* code;
  // Current program counter relative to the code pointer.
  size_t pc;
  ...
};

派遣循环

VM的关键是派遣循环。派遣循环通常会主导虚拟机的执行时间，但是我们通过实验发现Relay不是这种情况。我们刚刚实现了一个简单的switch/ goto派遣循环，该循环基于指令操作码进行派遣。

此循环由实现VirtualMachine::Run()。

VM编译器

该基础架构的重要组成部分是将Relay的完整IR编译为字节码序列的编译器。VM编译器将一个【tvm::relay::Module】转换为【tvm::relay::vm::Executable】。可执行文件包含一组编译函数，这些编译函数包含在【tvm::relay::vm::Function】中。这些函数包含有关该函数的元数据及其编译的字节码。然后，可以由【tvm::relay::vm::VirtualMachine】对象加载并运行发出的可执行对象。对于数据结构的完整定义，请参见 include/tvm/runtime/vm.h。

优化

VM编译器需要进行很多优化。它们中的每一个都被实现为由Relay过程管理器管理的过程。

标有TODO的优化尚未实现。

A范式
Lambda提升（请参阅src / relay / vm / lambda_lift.cc）
内联基元（请参阅src / relay / vm / inline_primitives.cc）
常量池布局（请参阅src / relay / backend / vm / compiler.cc）
尾部调用优化（TODO）
活力分析（TODO）

序列化

必须序列化和反序列化由Relay VM编译器生成的可执行文件，因为我们可能希望将模型保存到磁盘并在以后执行推理。以前，Relay以json文件的形式为图运行时生成了序列化表格。但是，相同的格式不适用于VM，因为它发出字节码而不是图样式的程序。可执行文件的序列化本质上需要处理特定于模型的数据（即权重和内核）以及与VM相关的数据（即字节码和全局函数名）。

对于内核，我们可以方便地利用现有的TVM基础框架来保存和加载已编译的库模块。在这里，我们仅关注以二进制格式序列化其他几个组件，该二进制格式按顺序按以下部分进行组织。

全局部分。这一部分包含虚拟机使用的全局变量（函数名称）。
常量部分。这一部分用于存储虚拟机的常量池（即模型的权重）。
原始名称部分。引入本节是为了容纳将由虚拟机调用的原始运算符名称的列表，即以【fused_】开头的名称。原语名称用作在已编译内核库中查找函数指针的符号。
代码部分。VM函数（包括字节码）位于本部分中。分派循环遍历此部分以获取要执行的指令。

因此，与包含权重（.params），json图（.json）和已编译内核库（.so）的图运行时工件不同，序列化的可执行工件由Relay对象文件（.ro）和已编译内核组成库（.so）。

【save】函数被实现为将可执行存储到磁盘，并将其序列化到上面的格式。同时，使用一个【load_exec】函数来加载序列化的内核二进制文件和与可执行文件相关的二进制代码，这些代码将再次用于实例化VM对象。有关更多示例，请参考test_vm_serialization.py文件。

尚未解决的问题

我们如何处理动态形状？

TODO

我们如何修改VM以支持某些代码路径的JIT编译？

在代码生成空间中，仍然需要进行许多折衷分析，并且VM的设计非常灵活，因此我们可以对其进行修改以供将来进行实验。

我们如何支持异构执行？

假设我们已经注释了适当的设备副本，则异类执行应该立即可用。为了正确执行此操作，我们需要运行设备注释和复制过程。

牛牛存

关注

1
点赞
踩
2

收藏

觉得还不错? 一键收藏
0
评论
将VM放入TVM：中继虚拟机

将VM放入TVM：中继虚拟机Relay是一种新的程序表示形式，它可以表示和优化大量的机器学习程序。不幸的是，通过支持一组更具表现力的程序，我们引入了一些新的执行挑战。Relay的解释器可以执行完整的语言，但是有明显的限制，使其不适合生产部署。它被构造为执行AST遍历以执行程序的低效率解释器。这种方法在概念上很简单，但是效率很低，因为AST遍历很大程度上依赖于间接。编译动态代码还面临其他...
复制链接

扫一扫

专栏目录