使用 tvm._ffi._init_api() 管理 TVM 插件#
拓展 TVM C++ 功能 调用细节可以借助 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.testing", __name__)
---------------------------------------------------------------------------
ModuleNotFoundError Traceback (most recent call last)
Cell In[1], line 1
----> 1 from tvm_book.tvm_ext.libinfo import load_lib
3 _LIB_EXT, _LIB_EXT_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/sym_add/outputs/libs"])
5 import tvm
File /media/pc/data/lxw/ai/tvm-book/src/tvm_book/tvm_ext/libinfo.py:4
2 import sys, os
3 import ctypes
----> 4 from tvm._ffi.libinfo import find_lib_path
6 def load_lib(name: str|list[str]|None=None, search_path: str|list[str]|None=None):
7 """通过搜索可能的路径加载库
8
9 Arg:
10 name: 需要导入的库名称（比如 ``libtvm_ext.so``），可为空（则为 ``libtvm.so``），也可为名称列表（暂未实现）
11 search_path: 搜索路径，可为空，也可为名称列表
12 """
ModuleNotFoundError: No module named 'tvm'
下面便可以直接使用 tvm_ext 下的函数了：
sym_add
<tvm.runtime.packed_func.PackedFunc at 0x7f992ebb3e80>