TVM apps extension示例扩展库

TVM apps extension示例扩展库

此文件夹包含TVM的示例扩展库。演示了其它库如何在C++和Python API中扩展TVM。

该库扩展了TVM的功能。

python模块加载新的共享库，可以使用TVM的python API进行插值。

https://github.com/apache/tvm/tree/main/apps/extension

 

 

 test_ext.py修改了一些代码

 

 

  

 

  

 

  

 

 运行结果

 

 

 TVM Runtime System

TVM支持多种编程语言用于编译器堆栈的开发和部署。将解释TVMRuntime的关键元素。

 

  需要满足许多有趣的要求：

 部署：从python/javascript/c++语言调用编译后的函数。

 调试：在python中定义一个函数，从编译后的函数调用该函数。

 链接：编写驱动程序代码，调用设备特定代码（CUDA），从编译后的主机函数中调用。

 原型：从Python定义IR传递，从C++后端调用。

 暴露：C++开发的编译器栈到前端（即，Python）

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

从任何语言定义函数，从另一种语言调用函数。希望Runtime核心最小，部署到嵌入式设备。

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 and string
• · PackedFunc itself
• · Module for compiled modules
• · DLTensor* for tensor object exchange
• · TVM Object to represent any object in 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的所有编译器pass函数，都作为PackedFunc公开给前端

编译后的模块将编译后的函数作为PackedFunc返回

为了保持Runtime最小值，将IR对象支持与部署Runtime隔离。根据包含多少Runtime驱动程序模块（如CUDA），生成的Runtime大约需要200K-600K。

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

Module

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

TVM将编译对象定义为模块。用户可以从模块中以PackedFunc的形式获取编译后的函数。生成的编译代码可以在Runtime从模块中动态获取函数。在第一次调用中缓存函数句柄，在后续调用中重用。将设备代码和回调从生成的代码链接到任何PackedFunc（如python）中。

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

远程部署

PackedFunc和模块系统可以轻松地将功能直接发送到远程设备。在后端，有一个RPCModule，序列化参数以进行数据移动，在远程启动计算。

 

 

 RPC服务器本身是最小的，可以捆绑到Runtime中。可以在iPhone/android/raspberry pi甚至浏览器上，启动最小的TVM RPC服务器。服务器上的交叉编译和用于测试的模块的发布，可以在同一个脚本中完成。有关更多详细信息，请参考交叉编译和RPC。

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

TVM对象和编译器堆栈

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

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

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

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

 因为通常一个对象可以在语言中的多个位置引用，所以使用一个共享的ptr跟踪引用。使用ObjectRef类表示对对象的引用。可以粗略地将ObjectRef类视为对象容器的shared_ptr。可以定义子类ObjectRef，保存对象的每个子类型。对象的每个子类都需要定义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) {}

  // ...

};

每个对象子类都将覆盖此项访问成员。下面是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);

  }

};

在上面的示例中，操作和数组<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++，便于对编译器堆栈进行扩展。这不是向成员公开前端语言的最快方法，但可能是最简单的方法之一。主要使用Python进行测试和原型开发，仍然使用C++完成繁重的工作。

实施详情

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

有关档案如下：

C++ API的PACKEDFUNC.H

c_runtime_api.cc用于c API，如何提供回调。

为了支持扩展类型，使用注册表系统登记类型相关信息，如C++中的任何支持，参阅扩展类型获取更多细节：https://github.com/apache/tvm/tree/main/apps/extension。

 

参考链接：

https://github.com/apache/tvm/tree/main/apps/extension

https://tvm.apache.org/docs/arch/runtime.html