版本：0.13.0

TVM Runtime 系统

TVM 支持多种编程语言进行编译器堆栈开发和部署。本文档将介绍 TVM runtime 的关键元素。

需要满足以下要求：

部署：从 Python/JavaScript/C++ 语言调用编译好的函数。
调试：在 Python 中定义一个函数，并从编译好的函数中调用它。
链接：编写驱动程序代码以调用特定设备代码（CUDA），并从编译的主机函数中调用它。
原型：从 Python 中定义一个 IR pass，并从 C++ 后端调用它。
公开：将 C++ 开发的编译器堆栈用到前端（即 Python）。
实验：将编译好的函数发送到嵌入式设备上直接运行。我们期望能够用任何语言定义一个函数，然后用另一种语言调用。还期望将 runtime 内核最小化，部署到嵌入式设备。

PackedFunc

PackedFunc 是一个简单而优雅的解决方案，它可以解决以上问题。单个 PackedFunc 对象代表一个函数调用，其调用者和被调用者可能使用不同的语言。

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

#include <tvm/runtime/packed_func.h>

void MyAdd(TVMArgs args, TVMRetValue* rv) {
  // 自动将参数转换为所需的类型。
  int a = args[0];
  int b = args[1];
  // 自动赋值返回给 rv
  *rv = a + b;
}

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

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

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

// 在 C++ 中注册一个全局打包函数
TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);

import tvm

myadd = tvm.get_global_func("myadd")
# 打印 3
print(myadd(1, 2))

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

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

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

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

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

import tvm

def callback(msg):
    print(msg)

# 转换成 PackedFunc
f = tvm.convert(callback)
callhello = tvm.get_global_func("callhello")
# 打印 hello world
callhello(f)

TVM 提供了一个最小的 C API，因此可将 PackedFunc 嵌入到任何语言中。除了 Python，目前还支持 java 和 javascript。这种嵌入式 API 的原理很像 Lua，除了它用的是 C++ 语言而非新的语言。

PackedFunc 用于编译器和部署堆栈：

TVM 的所有编译器 pass 函数都以 PackedFunc 的类型暴露给前端
编译好的模块还将编译好的函数作为 PackedFunc 类型返回

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

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

模块

由于 TVM 支持多种类型的设备，因此需要支持不同类型的驱动程序。必须用驱动程序 API 来加载内核，以打包格式设置参数，并执行内核启动。

还需要为驱动程序 API 打补丁，以便公开的函数是线程安全的。因此，经常要在 C++ 中实现这些驱动粘合，并将它们提供给用户。但不能对每种类型的函数都这样做，PackedFunc 又可用来辅助实现。

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

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

远程部署

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

RPC 服务器本身是最小的，可以捆绑到 runtime 中。可以在 iPhone/android/raspberry pi 甚至浏览器上启动一个最小的 TVM RPC 服务器。服务器上的交叉编译和测试模块的交付可以在同一个脚本中完成。查看交叉编译和 RPC 以获取更多详细信息。

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

TVM 对象和编译器堆栈

如前所述，在 PackedFunc runtime 系统之上构建编译器堆栈 API。因研究需要，编译器 API 在不断变化。要测试新的原语时，都需要一个新的语言对象或 IR 节点。但我们又不想经常更改 API。除此之外，还想

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

引入一个名为 Object 的基类来解决这个问题。编译器堆栈中的所有语言对象都是 Object 的子类，每个对象都包含一个字符串 type_key，用于唯一标识对象的类型。

之所以选择 string 而不是 int 作为类型键，是因为可以去中心化的方式来添加新的 Object 类，无需将代码添加回中心仓库。为了加快调度速度，在运行时为每个 type_key 分配一个整数 type_index。

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

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 的实现示例：

class TensorNode : public Object {
public:
  // 张量的形状
  Array<Expr> shape;
  // 张量内容中的简要数据类型
  Type dtype;
  // 简述源码操作，可以是None
  Operation op;
  // 简述源操作的输出索引
  int value_index{0};
  // 简要构造函数
  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 字段：

import tvm
from tvm import te

x = te.placeholder((3,4), name="x")
# 访问 TensorNode 的 op 字段
print(x.op.name)

可在不更改前端 runtime 的情况下，将新 Object 添加到 C++，从而轻松扩展编译器堆栈。

注意，这不是将成员提供给前端语言的最快方法，但可能是最简单的方法之一。并且它符合需求，因为我们主要用 Python 进行测试和原型设计，并仍用 C++ 来完成繁重的工作。

实现细节

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

Runtime 特定信息

Vulkan Runtime

TVM Runtime 系统

PackedFunc​

模块​

远程部署​

TVM 对象和编译器堆栈​

实现细节​

Runtime 特定信息

PackedFunc

模块

远程部署

TVM 对象和编译器堆栈

实现细节