使用 tvm._ffi._init_api() 管理 TVM 插件

使用 tvm._ffi._init_api() 管理 TVM 插件#

ref{ext-cpp} 调用细节可以借助 FFI 函数 tvm._ffi._init_api() 进行隐藏。

from tvm_book.tvm_ext.libinfo import load_lib

_LIB_EXT, _LIB_EXT_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/sym_add/outputs/libs"])

import tvm

tvm._ffi._init_api("tvm_ext", __name__)

下面便可以直接使用 tvm_ext 下的函数了:

sym_add
<tvm.runtime.packed_func.PackedFunc at 0x7f53d75f3210>