将VM放入TVM:Relay虚拟机

Relay是一种新的程序表示方法，它实现了大量机器学习程序的表示和优化。不幸的是，在引入支持更有表现力的程序集的同时，我们也引入了一些新的执行上的挑战。

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

在编译动态代码方面还有更多的挑战，比如动态调度和内存分配、完全动态张量形状和控制流。解释器为这些提供了简单的解决方案，但没有一个是足够引人注目或优化的。

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

另一种替代的但不是标准的方法是Relay的提前编译器，它将一个Relay程序编译到一个包含提前实现的共享库中。提前编译器提供了令人信服的性能，但是很难扩展和使用，这只能通过修改代码生成和优化机制来实现。

Relay虚拟机旨在成为一个平衡这些相互竞争的方法的框架，提供一个动态执行环境，可以通过灵活的扩展机制进行扩展、插装，并与其他方法(如提前编译)集成。

设计虚拟机是为了在部署和执行Relay程序时, 在性能和灵活性之间取得平衡，而不放弃TVM的好处。

在编程语言和系统中，虚拟机(VM)设计是一个得到充分研究的领域，已经有各种成熟的和嵌入式编程语言的虚拟机设计。以前的语言VM设计都是针对传统程序的执行特征进行了大量的定制。传统程序操作小标量值，并由大量低级指令组成。指令的数量要求指令的执行和调度非常高效。在机器学习中我们主要使用(相对)少量的高级指令操作张量值。ML程序的成本昂贵在大输入的算子调用上，例如GEMM或卷积。由于ML程序所展示的执行特征，在标量VM中进行微优化的重要性大大降低。

TVM为视觉模型提供了强有力的支持，但我们希望发展到支持更广泛的模型。图执行器能够利用输入图的完全静态特性来执行主动优化，例如完全静态分配和最佳内存重用。当我们引入使用控制流、递归、动态形状和动态分配的模型时，我们必须改变执行的工作方式。Relay的虚拟机是一个自然的选择。

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

设计 

VM的设计注重简单性的同时不牺牲性能。为了实现这一点，我们着重设计了一个张量VM，而不是标量VM。

在张量VM设置中，我们优化了对象的廉价“分配”(通过尝试避免真正的分配)、静态片段的重用和动态形状(即交错张量)的能力。

指令集

 指令集和指令表示的选择是虚拟机最关键的设计决策。指令的当前表示形式是一个带标记的联合，包含操作码和数据负载。一个重要的设计决策是指令的抽象级别(RISC vs. CISC)以及它们如何获取数据(固定宽度指令编码vs.可变长度指令编码)。当前版本更接近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

从给定的存储块storage中分配一个使用常量shape(存储在shape中)和dtype的张量值。结果保存到寄存器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

使用寄存器datatype_fields的num_fields条目分配带有标记tag的数据类型。保存结果到寄存器dst中。

AllocClosure

参数：

RegName dst
Index clo_index
Index num_freevar
RegName* free_vars;

 分配一个闭包，该闭包将clo_index处的VMFunction作为其编码，并从free_vars中的寄存器中分配num_freevar条目。结果保存到寄存器dst。

GetField

参数：

RegName dst
RegName object
Index field_index

用索引field_index从object获取字段值。并保存结果到寄存器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/vm.h和include/tvm/runtime/container.h中找到。

栈和状态 

Relay VM维护一个栈帧，其中包含关于如何恢复前一个调用的信息。在连续空间(虚拟寄存器文件)中为每个函数分配寄存器。

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

struct VirtualMachine {
  ...
  std::vector 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/executable.h和include/tvm/runtime/vm/vm.h。

优化

 VM编译器需要相当多的优化。它们中的每一个都被实现为一个由Relay pass管理器管理的pass。

用TODO标记的优化还没有实现。

• 一个正常的形式
• Lambda Lift(参见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)组成的。

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

未解决的问题 

我们如何处理动态shape

动态shape的支持是TVM正在进行的工作，因为我们升级Relay，TVM编译器。关于动态shape支持的最新更新，我们建议在TVM的讨论论坛(https://discuss.tvm.apache.org/)中进行更新。

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

在代码生成空间中，仍然有许多需要分析的权衡，VM被设计得非常灵活，所以我们可以为未来的实验修改它。

我们如何支持异构执行?

异构执行应该是开箱即用的，假设我们已经注释了适当的设备副本。为了正确地做到这一点，我们需要运行设备说明和复制pass。