TVM中PackedFun机制学习笔记

为了便于Python和C++混合编程，TVM使用了统一的PackedFunc机制。PackedFunc可以将C++中的各类函数打包成统一的函数接口，并自动导出到Python模块中进行调用，并且也支持从Python中注册一个函数，并伪装成PackedFunc在C++和Python中调用。

python调用c++

c++中

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;
}

在c++中注册全局PackedFunc，把c++中定义的函数的名称MyAdd在python的调用函数名称改为myadd，

TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);

python中

则在python中使用，先使用tvm.get_global_func方法获取c++中的myadd函数

import tvm

myadd = tvm.get_global_func("myadd") # python
print(myadd(1, 2)) # prints 3

c++调用python

python中

import tvm

def callback(msg):
  print(msg)

f = tvm.convert(callback) # convert to PackedFunc
callhello = tvm.get_global_func("callhello")
# prints hello world
callhello(f) 

c++中

TVM_REGISTER_GLOBAL("callhello")
.set_body([](TVMArgs args, TVMRetValue* rv) {
  PackedFunc f = args[0];
  f("hello world");
});