tvm - Runtime System

发表于 | 分类于

本文翻译自官方文档：TVM Runtime System

TVM 运行时系统

TVM 支持多种编程语言进行编译器堆栈开发和部署。在本说明中，我们解释了 TVM 运行时的关键元素。

https://tvm.apache.org/images/release/tvm_flexible.png

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

  • 部署：从python/javascript/c++语言调用编译好的函数。
  • 调试：在 python 中定义一个函数并从编译的函数中调用它。
  • 链接：编写驱动程序代码以调用设备特定代码（CUDA）并从编译的主机函数中调用它。
  • 原型：从 python 定义一个 IR 通道并从 C++ 后端调用它。
  • 公开：用 c++ 开发的编译器堆栈到前端（即 python）。
  • 实验：将编译后的函数发送到嵌入式设备上直接运行。

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

PackedFunc

PackedFunc 是一个简单而优雅的解决方案，我们发现它可以解决列出的挑战。单个 PackedFunc 对象代表一个函数调用，其调用者和被调用者可能使用不同的语言。

以下代码块提供了一个 C++ 示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#include <tvm/runtime/packed_func.h>

void MyAdd(TVMArgs args, TVMRetValue* rv) {
  // automatically convert arguments to desired type.
  int a = args[0];
  int b = args[1];
  // automatically assign value return to rv
  *rv = a + b;
}

void CallPacked() {
  PackedFunc myadd = PackedFunc(MyAdd);
  // get back 3
  int c = myadd(1, 2);
}

在上面的代码块中，我们定义了一个 PackedFunc MyAdd。它有两个参数：args 代表输入参数，rv 代表返回值。该函数是类型擦除的，这意味着函数签名不限制传入或返回的输入类型。在后台，当我们调用 PackedFunc 时，它会将输入参数打包到堆栈上的 TVMArgs，并通过 TVMRetValue 获取结果。

由于 C++ 中的模板技巧，我们可以像调用普通函数一样调用 PackedFunc。由于其类型擦除的性质，我们可以从像 python 这样的动态语言调用 PackedFunc，而无需为每个创建的新类型函数添加额外的胶水代码。以下示例在 C++ 中注册 PackedFunc 并从 python 调用。

1
2
3
// register a global packed function in c++
TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);
1
2
3
4
import tvm
myadd = tvm.get_global_func("myadd")
# prints 3
print(myadd(1, 2))

PackedFunc 的大部分魔力在于TVMArgs 和TVMRetValue 结构。我们限制了可以传递的可能类型的列表。以下是常见的：

  • 整数、浮点数和字符串
  • PackedFunc 本身
  • 编译模块的模块
  • DLTensor* 用于张量对象交换
  • TVM 对象表示 IR 中的任何对象

该限制使实现简单而无需序列化。尽管 PackedFunc 是最小的，但对于深度学习部署的用例来说已经足够了，因为大多数函数只需要 DLTensor 或数字。

由于一个 PackedFunc 可以将另一个 PackedFunc 作为参数，因此我们可以将函数从 python（作为 PackedFunc）传递给 C++。

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

def callback(msg):
  print(msg)

# convert to PackedFunc
f = tvm.convert(callback)
callhello = tvm.get_global_func("callhello")
# prints hello world
callhello(f)

TVM 提供了一个最小 C API，它允许我们将 PackedFunc 嵌入到任何语言中。除了 python，到目前为止，我们还支持 java 和 javascript。这种嵌入式 API 的哲学很像 Lua，只是我们没有新的语言，而是使用 C++。

PackedFunc 的一个有趣事实是我们将它用于编译器和部署堆栈。

  • TVM 的所有编译器传递函数都作为 PackedFunc 暴露给前端
  • 编译后的模块也将编译后的函数返回为 PackedFunc

为了将运行时保持在最低限度，我们将 IR 对象支持与部署运行时隔离开来。 生成的运行时大约需要 200K - 600K，具体取决于包含多少运行时驱动程序模块（例如，CUDA）。

与普通函数相比，调用 PackedFunc 的开销很小，因为它只在堆栈上保存了几个值。 所以只要我们不包装小函数就可以了。 总之，PackedFunc 是 TVM 中的通用粘合剂，我们广泛使用它来支持我们的编译器和部署。

模块

由于 TVM 支持多种类型的设备，我们需要支持不同类型的驱动程序。我们必须使用驱动程序 API 来加载内核，以打包格式设置参数并执行内核启动。我们还需要修补驱动程序 API，以便公开的函数是线程安全的。所以我们经常需要在 C++ 中实现这些驱动粘合并将它们暴露给用户。我们当然不能对每种类型的函数都这样做，所以 PackedFunc 也是我们的答案。

TVM 将编译后的对象定义为 Module。用户可以从 Module 中获取编译后的函数为 PackedFunc。生成的编译代码可以在运行时动态地从 Module 获取函数。它在第一次调用中缓存函数句柄并在后续调用中重用。我们使用它来将设备代码和回调链接到生成的代码中的任何 PackedFunc（例如 python）。

ModuleNode 是一个抽象类，可以由每种类型的设备实现。到目前为止，我们支持 CUDA、Metal、OpenCL 和加载动态共享库的模块。这种抽象使得新设备的引入变得容易，我们不需要为每种类型的设备重新生成主机代码。

远程部署

PackedFunc 和模块系统还可以轻松地将函数直接发送到远程设备。在底层，我们有一个 RPCModule 序列化参数以进行数据移动并在远程启动计算。

https://tvm.apache.org/images/release/tvm_rpc.png

RPC 服务器本身是最小的，可以捆绑到运行时中。我们可以在 iPhone/android/raspberry pi 甚至浏览器上启动一个最小的 TVM RPC 服务器。服务器上的交叉编译和测试模块的交付可以在同一个脚本中完成。查看 交叉编译和 RPC 了解更多详情。

这种即时反馈给我们带来了很多优势。例如，在 iPhone 上测试生成代码的正确性，我们不再需要从头开始在 swift/objective-c 中编写测试用例——我们可以使用 RPC 在 iPhone 上执行，将结果复制回来并在主机上进行验证通过 numpy。我们也可以使用相同的脚本进行分析。

TVM 对象和编译器堆栈

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

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

我们引入了一个名为 Object 的基类来解决这个问题。编译器堆栈中的所有语言对象都是 Object 的子类。每个对象都包含一个字符串 type_key，用于唯一标识对象的类型。我们选择 string 而不是 int 作为类型键，因此可以以分散的方式添加新的 Object 类，而无需将代码添加回中央仓库。为了加快调度速度，我们在运行时为每个 type_key 分配一个整数 type_index。

由于通常可以在语言中的多个位置引用一个“对象”，因此我们使用 shared_ptr 来跟踪引用。我们使用ObjectRef 类来表示对Object 的引用。我们可以粗略地将ObjectRef 类视为Object 容器的shared_ptr。我们还可以定义子类ObjectRef来保存Object的每个子类型。 Object 的每个子类都需要定义 VisitAttr 函数。

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

class BaseAttrsNode : public Object {
public:
  virtual void VisitAttrs(AttrVisitor* v) {}
  // ...
};

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

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

  void VisitAttrs(AttrVisitor* v) final {
    v->Visit("shape", &shape);
    v->Visit("dtype", &dtype);
    v->Visit("op", &op);
    v->Visit("value_index", &value_index);
  }
};

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

1
2
3
4
5
6
import tvm
from tvm import te

x = te.placeholder((3,4), name="x")
# access the op field of TensorNode
print(x.op.name)

可以在不更改前端运行时的情况下将新的 Object 添加到 C++，从而轻松扩展编译器堆栈。请注意，这不是将成员暴露给前端语言的最快方法，但可能是最简单的方法之一。我们还发现它符合我们的目的，因为我们主要使用 python 进行测试和原型设计，并且仍然使用 c++ 来完成繁重的工作。

实现细节

PackedFunc 中的每个参数都包含一个联合值 TVMValue 和一个类型代码。这种设计允许动态类型语言直接转换为相应的类型，而静态类型语言在转换过程中进行运行时类型检查。

相关文件是

为了支持扩展类型，我们使用了注册系统来注册类型相关信息，例如 C++ 中对任何类型的支持，请参阅扩展类型更多细节。