TVM 运行时系统#

TVM 支持编译器堆栈开发和部署的多种编程语言。在这篇笔记中，解释了 TVM 运行时的关键元素。

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

需要满足一些有趣的要求：

Deployment：从 python/javascript/c++ 语言中调用编译过的函数。
Debug：用 python 定义函数，然后从编译过的函数中调用它。
Link：编写驱动代码来调用设备专用代码（CUDA），并从编译的 host 函数中调用它。
Prototype：从 python 定义 IR pass，并从 c++ 后端调用它。
Expose：用 c++ 开发的编译器堆栈到前端（即 python）
Experiment：将编译过的函数传递到嵌入式设备中，直接运行。

希望能够从任何语言定义函数，然后从另一种语言调用它。还希望运行时的核心最小化，以便部署到嵌入式设备上。

PackedFunc#

PackedFunc 是简单且优雅的解决方案，可以解决列出的挑战。单独的 PackedFunc 对象表示函数调用，它的调用者和被调用者可能使用不同的语言。

下面的代码块提供了 C++ 示例

#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 调用。

// register a global packed function in c++
TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);

import tvm

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

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

int、float 和 string
PackedFunc 自身
用于编译模块的 Module
对于张量对象交换的 DLTensor*
TVM Object 表示 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)

# 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 Object support 与 deployment runtime 隔离开来。根据包含多少运行时驱动模块（如 CUDA），最终的运行时间大约需要 200K - 600K。

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

Module#

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

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

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

远程开发#

PackedFunc 和 Module 系统也使得直接将函数传送到远程设备变得很容易。在底层，有 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 Object 和编译器堆栈#

正如我们之前提到的，在 PackedFunc 运行时系统之上构建编译器堆栈 API。为了研究的需要，面临着编译器 API 的不断变化。每当想要测试新的原语时，都需要新的语言对象或 IR 节点。然而，并不希望每次都改变 API。除此之外，还想

能够序列化任何语言对象和 IRs
能够在前端语言中探索、打印和操作 IR 对象来进行快速原型制作。

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

由于通常一个 Object 可以在语言的多个地方被引用，可以使用 shared_ptr 来跟踪引用。使用 ObjectRef 类表示对 Object 的引用。可以粗略地将 ObjectRef 类视为 shared_ptr 到 Object 容器。也可以定义子类 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:
  /*! \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 字段。

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 中的每个参数包含 union 值 TVMValue 和类型代码。这种设计允许动态类型语言直接转换为相应的类型，而静态类型语言在转换过程中进行运行时类型检查。

专用运行时信息#

Vulkan Runtime

TVM 文档

TVM 运行时系统

导航