This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new f255650 feat: Introduce `tvm_ffi.libinfo.load_lib_module` (#336)
f255650 is described below
commit f255650b8e5121452dd60805b206a0cb5bcb6525
Author: Junru Shao <[email protected]>
AuthorDate: Fri Dec 12 05:07:22 2025 -0800
feat: Introduce `tvm_ffi.libinfo.load_lib_module` (#336)
This PR makes downstream library loading process correct and convenient,
by bringing in a new API `tvm_ffi.libinfo.load_lib_module`.
To demonstrate the effect, the file
`examples/packaging/python/my_ffi_extension/base.py`, which degenerates
to only 2 lines after this PR:
```python
import tvm_ffi
_LIB = tvm_ffi.libinfo.load_lib_module("my-ffi-extension",
"my_ffi_extension")
```
It's worth mentioning that the new API works smoothly for both editable
and non-editable builds, which is more correct than the existing one.
---
examples/packaging/CMakeLists.txt | 3 --
examples/packaging/python/my_ffi_extension/base.py | 28 +----------------
python/tvm_ffi/libinfo.py | 35 +++++++++++++++++++++-
3 files changed, 35 insertions(+), 31 deletions(-)
diff --git a/examples/packaging/CMakeLists.txt
b/examples/packaging/CMakeLists.txt
index 0463a88..3648da6 100644
--- a/examples/packaging/CMakeLists.txt
+++ b/examples/packaging/CMakeLists.txt
@@ -60,9 +60,6 @@ add_library(my_ffi_extension SHARED src/extension.cc)
target_link_libraries(my_ffi_extension tvm_ffi_header)
target_link_libraries(my_ffi_extension tvm_ffi_shared)
-# show as my_ffi_extension.so
-set_target_properties(my_ffi_extension PROPERTIES PREFIX "")
-
if (TVM_FFI_EXT_SHIP_DEBUG_SYMBOLS)
# ship debugging symbols for backtrace on macos
tvm_ffi_add_prefix_map(my_ffi_extension ${CMAKE_CURRENT_SOURCE_DIR})
diff --git a/examples/packaging/python/my_ffi_extension/base.py
b/examples/packaging/python/my_ffi_extension/base.py
index 4165023..5f1005e 100644
--- a/examples/packaging/python/my_ffi_extension/base.py
+++ b/examples/packaging/python/my_ffi_extension/base.py
@@ -16,32 +16,6 @@
# Base logic to load library for extension package
"""Utilities to locate and load the example extension shared library."""
-import sys
-from pathlib import Path
-
import tvm_ffi
-
-def _load_lib() -> tvm_ffi.Module:
- # first look at the directory of the current file
- file_dir = Path(__file__).resolve().parent
-
- path_candidates = [
- file_dir,
- file_dir / ".." / ".." / "build",
- ]
-
- if sys.platform.startswith("win32"):
- lib_dll_name = "my_ffi_extension.dll"
- elif sys.platform.startswith("darwin"):
- lib_dll_name = "my_ffi_extension.dylib"
- else:
- lib_dll_name = "my_ffi_extension.so"
- for candidate in path_candidates:
- for path in Path(candidate).glob(lib_dll_name):
- return tvm_ffi.load_module(str(path))
-
- raise RuntimeError(f"Cannot find {lib_dll_name} in {path_candidates}")
-
-
-_LIB = _load_lib()
+_LIB = tvm_ffi.libinfo.load_lib_module("my-ffi-extension", "my_ffi_extension")
diff --git a/python/tvm_ffi/libinfo.py b/python/tvm_ffi/libinfo.py
index 32295ce..d1c1bcf 100644
--- a/python/tvm_ffi/libinfo.py
+++ b/python/tvm_ffi/libinfo.py
@@ -27,7 +27,10 @@ import importlib.metadata as im
import os
import sys
from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+ from .module import Module
def find_libtvm_ffi() -> str:
@@ -169,6 +172,36 @@ def load_lib_ctypes(package: str, target_name: str, mode:
str) -> ctypes.CDLL:
return ctypes.CDLL(str(lib_path), getattr(ctypes, mode))
+def load_lib_module(package: str, target_name: str, keep_module_alive: bool =
True) -> Module:
+ """Load the tvm_ffi shared library by searching likely paths.
+
+ Parameters
+ ----------
+ package
+ The package name where the library is expected to be found. For
example,
+ ``"apache-tvm-ffi"`` is the package name of `tvm-ffi`.
+
+ target_name
+ Name of the CMake target, e.g., ``"tvm_ffi"``. It is used to derive
the platform-specific
+ shared library name, e.g., ``"libtvm_ffi.so"`` on Linux,
``"tvm_ffi.dll"`` on Windows.
+
+ keep_module_alive
+ Whether to keep the loaded module alive to prevent it from being
unloaded.
+
+ Returns
+ -------
+ The loaded shared library.
+
+ """
+ from .module import load_module # noqa: PLC0415
+
+ lib_path: Path = _find_library_by_basename(package, target_name)
+ # The dll search path need to be added explicitly in windows
+ if sys.platform.startswith("win32"):
+ os.add_dll_directory(str(lib_path.parent))
+ return load_module(lib_path, keep_module_alive=keep_module_alive)
+
+
def _find_library_by_basename(package: str, target_name: str) -> Path: #
noqa: PLR0912
"""Find a shared library by target_name name across known directories.
