TVM运行时系统

TVM运行时系统#

TVM 支持多种编程语言用于编译器堆栈的开发和部署。在本说明中，将解释 TVM 运行时的关键要素。

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

需要满足一些非常有趣的要求。

  • 部署：从 Python/JavaScript/C++ 语言调用编译后的函数。

  • 调试：在 Python 中定义函数，并在编译后的函数中调用它。

  • 链接：编写驱动(driver)程序代码以调用特定于设备的代码（CUDA），并从编译后的主机(host)函数中调用它。

  • 原型(Prototype)：在 Python 中定义 IR（中间表示）pass，并从 C++ 后端调用它。

  • 暴露：用 C++ 开发的编译器栈，面向前端（即 Python）

  • 实验：将编译好的函数发送到嵌入式设备上直接运行。

希望能够从任何语言定义函数，并从另一种语言调用它。同时，还希望运行时核心(core)尽可能小，以便部署到嵌入式设备上。

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，而不需要为每种新类型的函数创建额外的粘合(glue)代码。以下示例展示了如何在 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 和字符串

  • 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 提供了 minimum C API，这使得能够将 PackedFunc 嵌入到任何语言中。除了 Python，到目前为止还支持 java 和 javascript。这种嵌入式 API 的理念与 Lua 非常相似，不同之处在于没有创造一种新的语言，而是使用 C++。

关于 PackedFunc 的有趣的事实是，既用它来处理编译器，也用它来处理部署堆栈。

  • TVM 中的所有编译器 pass 函数都被作为 PackedFunc 暴露给前端。

  • 编译后的模块同样会以 PackedFunc 的形式返回编译好的函数。

为了最小化运行时的开销，将 IR 对象支持从部署运行时中分离出来。这样得到的运行时大约需要 200K 到 600K的空间，具体取决于包含了多少运行时驱动模块（例如 CUDA）。

调用 PackedFunc 与普通函数的开销很小，因为它仅在堆栈上保存了几个值。因此，只要不包装小型函数，就没有问题。总之，PackedFunc 是 TVM 中广泛使用的万能粘合剂(glue)，用它来支持编译器和部署工作。

Module#

由于 TVM 支持多种类型的设备，需要为不同类型的驱动程序提供支持。必须使用驱动 API 来加载内核，以打包格式设置参数并执行内核启动。还需要对驱动 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、树莓派甚至浏览器上启动最小化的 TVM RPC服务器。服务器上的交叉编译和模块的测试发布可以在相同的脚本中完成。有关更多详细信息，请参阅：tutorial-cross-compilation-and-rpc。

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

TVM Object 与编译器堆栈#

正如之前提到的，在 PackedFunc 运行时系统之上构建了编译器堆栈 API。由于研究需求，面临编译器 API 的持续变化。每当希望测试新的基元(primitives)时，就需要新的语言对象或 IR 节点。然而，不希望 API 频繁变动。此外，我们还希望

  • 能够将任何语言对象和 IRs 序列化。

  • 能够在前端探索、打印和操作 IR 对象，使用前端语言进行快速原型设计。

为了解决这个问题，引入了基类 Object 。编译器堆栈中的所有语言对象都是 Object 的子类。每个对象都包含字符串类型的 type_key，这个键唯一标识了对象的类型。选择使用字符串而非整数作为类型键，这样新的 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:
  /*! \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)

在不改变前端运行时的情况下，可以在 C++ 中添加新的 Object，这使得对编译器堆栈的扩展变得容易。值得注意的是，这并不是将成员暴露给前端语言的最快捷方式，但可能是最简单的方法之一。还发现这种方法适合我们的目的，因为我们主要使用 Python 进行测试和原型设计，而仍然使用 C++ 来完成繁重的工作。

实现细节#

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

相关文件

为了支持扩展类型，采用了注册系统来登记(register)类型的相关信息，比如在 C++ 中的支持情况，详情请见 Extension types 部分。

运行时特定信息#