TVM运行时系统

[翻译]TVM运行时系统— tvm 0.6.0文档 - 图1

TVM支持多种编程语言用于编译器堆栈的开发和部署。在本说明中，我们解释了TVM运行时的关键元素。
[翻译]TVM运行时系统— tvm 0.6.0文档 - 图2

我们需要满足很多有趣的要求：

  • 部署：从python/javascript/c++语言调用已编译的函数。
  • 调试：在python中定义一个函数，然后从别的已编译函数中调用该函数。
  • 链接：编写驱动程序代码以调用设备特定代码（CUDA），然后从已编译的主机函数中调用它。
  • 原型：从python定义IR pass，并从c++后端调用它。
  • 公开：将使用c++开发的编译器堆栈对前端（例如python）公开
  • 实验：将已编译的函数传送到嵌入式设备以在其中直接运行它。

我们希望能够从任何一种语言定义一个函数并从另一种语言调用它，我们还希望将运行时核心最小化以部署到嵌入式设备。

PackedFunc

PackedFunc是解决前面所列挑战的一个简单但优雅的解决方案。以下代码块提供了一个示例（ c++）

  1. #include <tvm/runtime/packed_func.h>
  2. void MyAdd(TVMArgs args, TVMRetValue *rv)
  3. {
  4.     // automatically convert arguments to desired type.
  5.     int a = args[0];
  6.     int b = args[1];
  7.     // automatically assign value return to rv
  8.     *rv = a + b;
  9. }
  10. void CallPacked()
  11. {
  12.     PackedFunc myadd = PackedFunc(MyAdd);
  13.     // get back 3
  14.     int c = myadd(1, 2);
  15. }

在上面的代码块中，我们定义了PackedFunc MyAdd。它有两个参数：args代表输入参数，rv代表返回值。该函数是类型擦除的，这意味着函数签名不限制传入的输入类型或返回的类型。在内部，当我们调用PackedFunc时，它将输入参数打包到堆栈中的TVMArgs，并通过TVMRetValue返回结果。
多亏了c++中的模板技巧，我们可以像普通函数一样调用PackedFunc。由于它具有类型擦除的特性，因此我们可以从动态语言（如python）调用PackedFunc，而无需为创建的每个新类型函数添加额外的胶水代码。以下示例在c++中注册PackedFunc并从python进行调用。
// register a global packed function in c++
TVM_REGISTER_GLOBAL(“myadd”)
.set_body(MyAdd);
import tvm

myadd = tvm.getglobal_func(“myadd”)
# prints 3_
print(myadd(1, 2))
大多数PackedFunc的奇妙之处在于TVMArgs和TVMRetValue结构。我们限制了可以传递的类型，下面这些是常见的类型：

  • int，float和string
  • PackedFunc本身
  • 用于编译Module的Module
  • DLTensor *用于张量对象交换
  • TVM Node，它表示IR中的任何对象

该限制使实现变得简单而无需序列化。虽然很小，但PackedFunc对于深度学习部署的用例已足够，因为大多数函数仅采用DLTensor或数字作为参数。
由于一个PackedFunc可以将另一个PackedFunc用作参数，因此我们可以将函数从python（作为PackedFunc）传递给c++

  1. TVM_REGISTER_GLOBAL("callhello")
  2. .set_body([](TVMArgs args, TVMRetValue *rv) {
  3.     PackedFunc f = args[0];
  4.     f("hello world");
  5. });

**

  1. import tvm
  2. def callback(msg):
  3.   print(msg)
  4. # convert to PackedFunc
  5. f = tvm.convert(callback)
  6. callhello = tvm.get_global_func("callhello")
  7. # prints hello world
  8. callhello(f)

TVM提供了尽可能少的C API，使我们可以将PackedFunc嵌入任何语言。到目前为止，除python外，我们还支持 java和javascript。嵌入式API的思想非常类似于Lua，除了我们没有使用一种新语言而是使用c++。
关于PackedFunc的一个有趣的事实是，我们将其用于编译器和部署堆栈。

  • 所有TVM的编译器传递函数都以PackedFunc的形式公开给前端，请参见此处
  • 编译后的 Module还会将编译后的函数以PackedFunc的形式返回。

为了使运行时最少，我们将IR Node支持与部署运行时隔离开来。生成的运行时大约需要200K-600K，具体取决于包含的运行时驱动程序 Module（例如CUDA）数量。
与正常函数相比，调用PackedFunc的开销很小，因为它只在堆栈中保存了一些值。因此，只要我们不包装小的函数就可以。总而言之，PackedFunc是TVM中的通用粘合剂，我们在TVM中广泛使用它来支持我们的编译器和部署。

Module

由于TVM支持多种类型的设备，因此我们需要支持不同类型的驱动程序。我们必须使用驱动程序API来加载内核，以打包格式设置参数并执行内核启动。我们还需要修补驱动程序API，以使公开的函数具有线程安全性。因此，我们经常需要在c++中实现这些驱动程序粘合并将其公开给用户。当然，我们不能为每种功能都这样做，所以PackedFunc还是我们的答案。
TVM将已编译对象定义为Module。用户可以从 Module中以PackedFunc的形式获取已编译的函数。生成的编译代码可以在运行时从Module动态获取函数。它在第一次调用中缓存该函数句柄，并在后续调用中重用。我们使用它来链接设备代码并从生成的代码回调到任何PackedFunc（例如python）中。
ModuleNode是一个抽象类，它由具体的各种类型的设备实现。到目前为止，我们支持CUDA，Metal，OpenCL和加载动态共享库的 Module。这种抽象使新设备的引入变得容易，并且我们不需要为每种类型的设备重做主机代码生成。

远程部署

PackedFunc和Module系统还使将函数直接传送到远程设备变得容易。在后台，我们有一个RPCModule，它将参数序列化以进行数据移动并在远程上启动计算。
[翻译]TVM运行时系统— tvm 0.6.0文档 - 图3
RPC服务端本身是极小的，可以捆绑到运行时中。我们可以在iPhone/android/raspberry pi甚至浏览器上启动极小的TVM RPC服务端。服务端的交叉编译和测试 Module的交付可以在同一脚本中完成。详情请参阅 Cross编译和RPC教程。
这种即时反馈有很多优势。例如，要测试iPhone上生成的代码的正确性，我们不再需要从头开始在swift/objective-c中编写测试用例，我们可以使用RPC在iPhone上执行，将结果复制回来并在主机上通过numpy进行验证。我们也可以使用相同的脚本进行分析。

TVM Node和编译器堆栈

如前所述，我们在PackedFunc运行时系统之上构建编译器堆栈API。为了研究的需要，我们面临着不断变化的编译器API。每当我们要测试新的原语(primitives)时，我们都需要一个新的语言对象或IR node。但是，我们不想频繁的更改API。除此之外，我们也想：

  • 能够序列化任何语言对象和IR
  • 能够以前端语言浏览、打印和操作IR对象以进行快速原型制作

我们引入了一个称为Node的基类来解决此问题。编译器堆栈中的所有语言对象都是Node的子类。每个Node都包含一个字符串type_key，它唯一地标识对象的类型。我们选择字符串而不是int作为类型键，这样就可以以分散方式添加新的Node类，而无需将代码添加回中央存储库。为了降低分发速度，我们在运行时为每个type_key分配了一个整数type_index。
由于通常可以用该语言在多个位置引用同一个Node对象，因此我们使用shared_ptr来跟踪引用。我们使用NodeRef类来表示对Node的引用。我们可以将NodeRef类粗略地视为Node容器的shared_ptr。我们还可以定义NodeRef子类来保存Node的每个子类型。每个Node类都需要定义VisitAttrs函数。

  1. class AttrVisitor
  2. {
  3. public:
  4.     virtual void Visit(const char *key, double *value) = 0;
  5.     virtual void Visit(const char *key, int64_t *value) = 0;
  6.     virtual void Visit(const char *key, uint64_t *value) = 0;
  7.     virtual void Visit(const char *key, int *value) = 0;
  8.     virtual void Visit(const char *key, bool *value) = 0;
  9.     virtual void Visit(const char *key, std::string *value) = 0;
  10.     virtual void Visit(const char *key, void **value) = 0;
  11.     virtual void Visit(const char *key, Type *value) = 0;
  12.     virtual void Visit(const char *key, NodeRef *value) = 0;
  13.     // ...
  14. };
  15. class Node
  16. {
  17. public:
  18.     virtual void VisitAttrs(AttrVisitor *visitor) {}
  19.     // ...
  20. };

每个Node子类都将覆盖它以访问其成员。下面是TensorNode的实现示例。

  1. class TensorNode : public Node
  2. {
  3. public:
  4.     /*! \brief The shape of the tensor */
  5.     Array<Expr> shape;
  6.     /*! \brief data type in the content of the tensor */
  7.     Type dtype;
  8.     /*! \brief the source operation, can be None */
  9.     Operation op;
  10.     /*! \brief the output index from source operation */
  11.     int value_index{0};
  12.     /*! \brief constructor */
  13.     TensorNode() {}
  14.     void VisitAttrs(AttrVisitor *v) final
  15.     {
  16.         v->Visit("shape", &shape);
  17.         v->Visit("dtype", &dtype);
  18.         v->Visit("op", &op);
  19.         v->Visit("value_index", &value_index);
  20.     }
  21. };

在上面的示例中，Operation和Array<Expr>均为NodeRef。VisitAttrs为我们提供了一个反射API来访问对象的每个成员。我们可以使用此函数访问Node并递归序列化任何语言对象。它还使我们能够轻松地以前端语言来获取对象的成员。例如，以下代码显示如何访问了TensorNode的op字段。

  1. import tvm
  2. from tvm import te
  3. x = te.placeholder((3,4), name="x")
  4. # access the op field of TensorNode
  5. print(x.op.name)

可以在不更改前端运行时的情况下将新对象添加到C++中，这使得对编译器堆栈进行扩展变得很容易。 请注意，这不是让成员接触前端语言的最快方法，但可能是最简单的方法之一。 我们还发现它符合我们的目的，因为我们主要使用python进行测试和原型开发，而仍然使用c++来完成繁重的提升工作。

实施细则

PackedFunc中的每个参数都包含联合值TVMValue 和类型代码。这种设计允许动态类型化的语言直接转换为相应的类型，而静态类型化的语言可以在转换期间进行运行时类型检查。
相关文件：

为了支持扩展类型，我们使用了注册表系统来注册与类型相关的信息，例如在c++中对任何类型的支持，请参阅扩展类型以获取更多详细信息。