为了便于Python和C++混合编程,TVM使用了统一的PackedFunc机制。PackedFunc可以将C++中的各类函数打包成统一的函数接口,并自动导出到Python模块中进行调用,并且也支持从Python中注册一个函数,并伪装成PackedFunc在C++和Python中调用。
c++中需要调用的函数
#include
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中使用,先使用tvm.get_global_func方法获取c++中的myadd函数
import tvm
myadd = tvm.get_global_func("myadd") # python
print(myadd(1, 2)) # prints 3
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)
TVM_REGISTER_GLOBAL("callhello")
.set_body([](TVMArgs args, TVMRetValue* rv) {
PackedFunc f = args[0];
f("hello world");
});