深度学习编译中间件之NNVM(四)TVM设计理念与开发者指南

参考文档

http://docs.tvmlang.org/dev/index.html TVM Design and Developer Guide

本文档为官方指导手册的中文翻译版本，主要涉及到TVM的设计理念和开发者指南，适用于计划深入掌握TVM深度定制开发技术的开发者。

TVM运行时系统

TVM支持多种编程语言下的编译器堆栈开发和部署，针对本文档我们主要会介绍TVM运行时的关键组件。

这里写图片描述

我们需要满足相当多的软件需求：

Deployment(部署)：能够通过Python/Javascript/C++来调用被编译的函数
Debug(调试)：定义一个Python函数，被编译的函数能够调用这个Python函数
Link(链接)：设计设备相关代码(负责调用设备特定代码，例如CUDA)，并且这些代码能够被主机函数调用
Prototype(原型)：通过Python定义一个IR Pass1，此Pass能够被C++后端调用
Expose(暴露接口)：通过C++设计的编译器堆栈需要暴露接口给前端语言(例如Python)
Experiment(验证支持)：主要是针对嵌入式设备设计一套RPC接口(远程调用接口)从而加速验证过程

简而言之，我们需要确保通过一种语言定义的函数能够被另外的语言调用，另外还要针对嵌入式设备最小化运行时核心。

PackedFunc

对于上面列举的软件需求，PackedFunc是一个简单却优雅的解决方案。下面列举一个C++的PackedFunc示例：

#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类型函数。因为PackedFunc类型函数是无类型的，所以Python语言无需古怪的语法就可以调用PackedFunc函数。下面通过示例来展示：

// 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的良好设计。下面列举PackedFunc函数能够传递哪些类型的数据：

int float and string
PackedFunc类型自身
Module for compiled modules
DLTensor交换格式
TVM结点，表示IR

在不同语言间传递上上述类型的数据时不需要进行专门的序列化处理，而对于深度学习部署这种使用场景，PackedFunc能够满足部署需求，大部分函数只需要传递DLTensor和数字类型的数据。

因为PackedFunc可以传递PackedFunc类型的参数，所以我们可以把函数从Python传递到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，可以通过C语言API把PackedFunc嵌入到任何编程语言中。除了Python，我们还计划添加对Java和JavaScript的支持。嵌入API的设计哲学类似Lua。

关于PackedFunc有一个比较有意思的地方，就是它同时被编译器堆栈和部署堆栈使用了。

所有的TVM编译器Pass函数通过PackedFunc暴露接口给前端语言
已经被编译的模块也通过PackedFunc返回已编译的函数

为保证TVM Runtime的最小化，我们把运行时和IR Node隔离开。最终整个运行时的体积只有200K-600K，浮动的区间取决于包含了驱动支持(例如CUDA)。

因为只有非常少的参数在堆栈中，所以调用PackedFunc的负担和普通函数相比是小的。

模块

因为TVM需要支持多种类型的硬件设备，所以我们需要支持不同类型的驱动。我们必须使用驱动API来加载Kernel，设置Packed格式的参数和启动Kernel。我们也需要修补驱动API，以此让被暴露的函数是线程安全的。所以我们经常需要用C++来实现这些驱动胶水，并把这些暴露给用户。因为PackedFunc的原因，我们不需要为每一种类型的函数都做一个适配。

TVM定义Module作为已编译对象。用户可以通过PackedFunc从Module中获取已经编译的函数。在运行时可以动态地从Module中获取已经编译生成的代码。当代码被第一次调用之后会被缓存，以保证接下来相同代码的调用能重用已经缓存的代码。

ModuleNode是一个抽象类，被用来实现每个类型的设备驱动。到目前为止，我们已经支持了CUDA，Metal，Opencl模块。此处的抽象能够使新设备的添加变得容易，我们不需要为每个类型的设备重新设计host端代码生成逻辑。

远程部署

TVMNode和编译器堆栈

在文章的前面部分已经提到过，编译器堆栈API处于PackedFunc运行时系统之上。为了研究的需要，我们面对着一个编译器API经常需要变化的现实。我们需要一种新的IR语言，但是我们并不想大幅改变我们现有的API，我们总结我们对于编译器语言的需求：

能够序列化任何语言对象和IR
能够在前端语言中比较快捷地浏览、打印、操作IR对象

我们先介绍一个基类Node来满足上面的需求，在编译器堆栈中所有的语言对象都是Node类的子类。每个Node包含一个字符串type_key来惟一标识对象的类型。我们选择字符串作为type_key的类型是为了能够让新的Node类可以被添加分散管理的代码库中。为了缓解调度时的速度问题，在Runtime运行时我们也分配了一个int类型的type_index来标识对象的类型。

因为一般情况下一个Node对象可以在一种语言中的不同位置被引用，我们使用shared_ptr来记录引用。NodeRef类被用来标识Node的引用。我们也定义多个NodeRef的子类来处理Node的子类，每个Node类都需要定义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, NodeRef* value) = 0;
  // ...
};

class Node {
 public:
  virtual void VisitAttrs(AttrVisitor* visitor) {}
  // ...
};

每个Node的子类都会Override(重载)VisitAttrs来访问它的成员。在这里展示一个相应的示例：

class TensorNode : public Node {
 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>都是NodeRef。VisitAttrs提供了一个ReflectionAPI(反射API)来访问对象里面的每一个成员。我们也可以使用这个函数来访问Node节点和递归地序列化任何语言对象。它也允许我们在前端语言中容易地获取对象的成员。例如在下面的示例中，我们存取TensorNode的op成员：

import tvm

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

当添加一个新的Node类型到C++时，我们不需要改变前端运行时，这使得扩展编译器堆栈变得容易。

实现细节

PackedFunc的每一个参数都包含一个联合体类型的数据TVMValue和一个类型代码。这种设计允许动态类型语言能够直接转换到相应的类型，静态类型语言可以在运行时检查数据类型。

此术语为编译器领域专用，在LLVM的架构中，Pass的作用是优化LLVM IR。详见LLVM Cookbook中文版第4章 ↩