TVM运行时系统

本文翻译自TVM Runtime System — tvm 0.9.dev0 documentation

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

我们需要满足一些有趣的需求:

• 部署:从python/javascript/c++语言调用编译后的函数。
• 调试:在python中定义一个函数，并从编译后的函数中调用它。
• 链接:编写驱动程序代码调用设备特定代码(CUDA)，并从编译后的host侧函数调用它。
• Prototype:用python定义一个IR pass，并从c++后端调用它。
• 公开:用c++开发的编译器栈到前端(python)
• 实验:将编译好的函数发布到嵌入式设备上，并直接在设备上运行。

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

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本身
• 编译得到的模块
• 表示张量对象交换的DLTensor*
• 用于表示IR中的任何对象的TVM Object 

这个限定使实现变得简单，不需要对数据做序列化。尽管只支持少数几种类型，但PackedFunc对于深度学习部署的用例来说已经足够了，因为大多数函数只使用DLTensor或数字。

因为一个PackedFunc可以接受另一个PackedFunc作为实参，所以我们可以将函数从Python传递给C++ (PackedFunc)。

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的编译器pass函数都作为PackedFunc公开给前端
• 编译后的模块还将编译后的函数以PackedFunc返回

为了使运行时最小化，我们将IR对象支持与部署运行时隔离开来。最终的运行时模块大小大约需要200K - 600K，这取决于包含了多少运行时驱动模块(例如CUDA)。

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

 Module

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

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

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

远程部署 

PackedFunc和Module系统还可以很容易地将函数直接发送到远程设备。在幕后，我们有一个RPCModule，它序列化参数来进行数据移动，并远程启动计算。

RPC服务器本身很小，可以绑定到运行时中。我们可以在iPhone/android/树莓派甚至浏览器上启动一个最小的TVM RPC服务器。服务器上的交叉编译和测试模块的交付可以在同一个脚本中完成。更多细节可以参阅。

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

如前所述，我们在PackedFunc运行时系统之上构建编译器栈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:
  /*! \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和一个类型编码。这种设计允许动态类型语言直接转换为相应的类型，而静态类型语言可以在转换期间进行运行时类型检查。 

相关文件如下：

• packed_func.h ： C++ API

• c_runtime_api.cc ： C API和如何提供回调。

为了支持扩展类型，我们使用了一个注册系统来注册与类型相关的信息，类似C++中对任何类型的支持，参见扩展类型了解更多细节。 

具体运行时信息 

• Vulkan Runtime