TVM运行系统

最新推荐文章于 2023-11-28 09:59:23 发布

牛牛存

最新推荐文章于 2023-11-28 09:59:23 发布

阅读量1.3k

点赞数 1

分类专栏： tvm TVM Relay

本文链接：https://blog.csdn.net/weixin_42164269/article/details/104291845

版权

TVM Relay 同时被 2 个专栏收录

25 篇文章 21 订阅

订阅专栏

tvm

14 篇文章 6 订阅

订阅专栏

TVM运行系统

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

我们需要满足很多有趣的要求：

部署：从python / javascript / c ++语言调用已编译的函数。
调试：在python中定义一个函数，然后从已编译函数调用该函数。
链接：编写驱动程序代码以调用设备专用代码（CUDA），然后从已编译的主机函数中调用它。
原型：从python定义IR转换，并从C ++后端调用它。
公开：使用C ++开发的编译器堆栈到前端（例如python）
实验：将已编译的函数运送到嵌入式设备以直接在其中运行。

我们希望能够从任何一种语言定义一个函数并从另一种语言调用。我们还希望最小化运行时核以将其部署到嵌入式设备。

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】结构。我们限制了可以传递的可能类型的列表。这是常见的：

整数，浮点数和字符串
PackedFunc本身
Module用于编译模块
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)

# 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支持与部署运行时隔离开来。生成的运行时大约需要200K-600K，具体取决于包含的运行时驱动程序模块（例如CUDA）数量。

与正常函数相比，调用PackedFunc的开销很小，因为它只在堆栈中保存了一些值。因此，只要我们不包装小的函数就可以。总而言之，PackedFunc是TVM中的通用粘合剂，我们在其中广泛使用它来支持我们的编译器和部署。

模块

由于TVM支持多种类型的设备，因此我们需要支持不同类型的驱动程序。我们必须使用驱动程序API来加载内核，以打包格式设置参数并执行内核启动。我们还需要修补驱动程序API，以使公开的函数具有线程安全性。因此，我们经常需要在C ++中实现这些驱动程序粘合并将其公开给用户。当然，我们不能为每种功能都做到这一点，所以PackedFunc还是我们的答案。

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

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

远程部署

PackedFunc和Module系统还使将函数直接移植到远程设备中变得容易。在后台，我们有一个RPCModule，它对参数进行序列化以进行数据移动并在远程上启动计算。

RPC服务器本身是最小的，可以捆绑到运行时中。我们可以在iPhone / android / raspberry pi甚至浏览器上启动最小的TVM RPC服务器。服务器上的交叉编译和测试模块的交付可以在同一脚本中完成。有关更多详细信息，请参阅Checkout Cross编译和RPC教程。

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

TVM对象和编译器堆栈

如前所述，我们在PackedFunc运行时系统之上构建编译器堆栈API。为了研究的需要，我们面临着不断变化的编译器API。每当我们要测试新的原语时，我们都需要一个新的语言对象或IR节点。但是，我们不想不时更改API。除此之外，我们也想

能够序列化任何语言对象和IR
能够以前端语言浏览，打印和操作IR对象以进行快速原型制作。

我们引入了一个称为Object的基类来解决此问题。编译器堆栈中的所有语言对象都是【Object】的子类。每个对象都包含一个字符串type_key，它唯一地标识对象的类型。我们选择字符串而不是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

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

新的【Object】可以在不更改前端运行时的情况下被添加到C ++，从而轻松扩展编译器堆栈。请注意，这不是向成员公开前端语言的最快方法，而是可能的最简单方法之一。我们还发现它符合我们的目的，因为我们主要使用python进行测试和原型设计，仍然使用c ++进行繁重的工作。

实现细节

PackedFunc中的每个参数都包含联合值TVMValue 和类型代码。这种设计允许动态类型化的语言直接转换为相应的类型，而静态类型化的语言可以在转换期间进行运行时类型检查。