This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ebc9348bd7 [REFACTOR][PYTHON] Slim tvm.libinfo to info-only helpers
(#19719)
ebc9348bd7 is described below
commit ebc9348bd727d8cbd132f7ab27e8d939d3c1c75a
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Jun 10 13:41:37 2026 -0400
[REFACTOR][PYTHON] Slim tvm.libinfo to info-only helpers (#19719)
This PR slims `tvm.libinfo` into a thin *info* layer that delegates path
discovery to the `tvm_ffi.libinfo` primitives and never loads libraries.
Loading responsibilities move to `tvm.base`, and the various ad-hoc
path-finding helpers are phased out in favor of the tvm-ffi resolvers.
## Changes
- **libinfo**: add `find_libtvm_runtime()` (resolves `libtvm_runtime`
via
`_find_library_by_basename` + `_resolve_and_validate`) and
`find_tvm_include_path()` (TVM's own `include/`). `find_include_path()`
now
returns `[find_tvm_include_path(), *tvm_ffi.libinfo.include_paths()]`,
folding
in the FFI + dlpack + python-helper include dirs. Remove
`find_lib_path`,
`get_dll_directories`, `use_runtime_lib`, `split_env_var`, and
`load_backend_libs`.
- **base**: receive `load_backend_libs` and the backend DSO list; the
runtime-only switch becomes a strict `TVM_USE_RUNTIME_LIB == "1"` check.
- **rpc**: `with_minrpc` uses `find_libtvm_runtime()` (the `runtime`
kwarg is
retained as an inert back-compat parameter); the rpc server
`load_library`
resolves the literal library name against the current working directory.
- **wasm**: move the `web/dist` asset search into `emcc.find_wasm_lib`,
used by
`emcc.create_tvmjs_wasm` and the tvmjs asset lookup.
- **hexagon**: fix a latent bug where `_get_hexagon_rpc_lib_dir` called
a
non-existent `tvm_ffi.libinfo.find_lib_path`; it now relies solely on
the
`HEXAGON_RPC_LIB_DIR` environment variable.
---
python/tvm/base.py | 29 +-
python/tvm/contrib/hexagon/build.py | 11 +-
python/tvm/contrib/tvmjs.py | 7 +-
python/tvm/libinfo.py | 311 +++------------------
python/tvm/rpc/minrpc.py | 14 +-
python/tvm/rpc/server.py | 11 +-
python/tvm/runtime/module.py | 10 +-
python/tvm/support/emcc.py | 46 ++-
.../python/relax/test_frontend_nn_extern_module.py | 9 +-
9 files changed, 143 insertions(+), 305 deletions(-)
diff --git a/python/tvm/base.py b/python/tvm/base.py
index 11601a9dcc..f3f42263b4 100644
--- a/python/tvm/base.py
+++ b/python/tvm/base.py
@@ -21,6 +21,7 @@
import os
import sys
+from pathlib import Path
from tvm_ffi.libinfo import load_lib_ctypes
@@ -37,6 +38,26 @@ if not (sys.version_info[0] >= 3 and sys.version_info[1] >=
9):
# library loading
# ----------------------------
+# Known per-backend runtime DSOs that, when present, are loaded with
+# RTLD_GLOBAL so their static initializers register the device backend.
+_BACKEND_RUNTIME_LIBS = ["cuda", "vulkan", "opencl", "metal", "rocm",
"hexagon", "extra"]
+
+
+def load_backend_libs(runtime_lib_path: str) -> None:
+ """Try to load each known backend runtime DSO; failures are silent."""
+ runtime_dir = Path(runtime_lib_path).resolve().parent
+ for backend in _BACKEND_RUNTIME_LIBS:
+ try:
+ load_lib_ctypes(
+ package="tvm",
+ target_name=f"tvm_runtime_{backend}",
+ mode="RTLD_GLOBAL",
+ extra_lib_paths=[runtime_dir],
+ )
+ except (OSError, FileNotFoundError, RuntimeError):
+ pass
+
+
# The TVM C++ side is split into two shared libraries:
#
# - ``libtvm_runtime`` — runtime-only sources. Loaded with ``RTLD_GLOBAL`` so
@@ -46,8 +67,8 @@ if not (sys.version_info[0] >= 3 and sys.version_info[1] >=
9):
# ``libtvm_runtime``. Loaded with ``RTLD_LOCAL`` so compiler internals
# don't leak into the global symbol namespace.
#
-# If the environment variable ``TVM_USE_RUNTIME_LIB`` is truthy, or the
-# compiler library is simply not present (runtime-only wheel), only the
+# If the environment variable ``TVM_USE_RUNTIME_LIB`` is set to ``"1"``, or
+# the compiler library is simply not present (runtime-only wheel), only the
# runtime is loaded and ``_LIB`` aliases ``_LIB_RUNTIME``.
_extra_lib_paths = libinfo.package_lib_paths()
_LIB_RUNTIME = load_lib_ctypes(
@@ -59,9 +80,9 @@ _LIB_RUNTIME = load_lib_ctypes(
# with RTLD_GLOBAL so their static initializers register device backends.
# Failures are swallowed silently — a missing driver just means that backend
# is unavailable, not an error.
-libinfo.load_backend_libs(_LIB_RUNTIME._name)
+load_backend_libs(_LIB_RUNTIME._name)
-_RUNTIME_ONLY = libinfo.use_runtime_lib()
+_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1"
if _RUNTIME_ONLY:
_LIB = _LIB_RUNTIME
else:
diff --git a/python/tvm/contrib/hexagon/build.py
b/python/tvm/contrib/hexagon/build.py
index 6320377798..30b9c48458 100644
--- a/python/tvm/contrib/hexagon/build.py
+++ b/python/tvm/contrib/hexagon/build.py
@@ -34,8 +34,6 @@ import subprocess
import sys
import tempfile
-from tvm_ffi import libinfo
-
from tvm.contrib.hexagon.hexagon_profiler import HexagonProfiler
from .session import Session
@@ -74,14 +72,9 @@ def _get_hexagon_rpc_lib_dir() -> pathlib.Path:
The path to the Hexagon API directory.
"""
global HEXAGON_RPC_LIB_DIR
+ HEXAGON_RPC_LIB_DIR = os.environ.get("HEXAGON_RPC_LIB_DIR")
if HEXAGON_RPC_LIB_DIR is None:
- for path in libinfo.find_lib_path():
- rpc_dir = os.path.join(os.path.dirname(path), "hexagon_api_output")
- if os.path.isdir(rpc_dir):
- HEXAGON_RPC_LIB_DIR = rpc_dir
- break
- else:
- raise RuntimeError("hexagon_api binaries not found, please define
HEXAGON_RPC_LIB_DIR")
+ raise RuntimeError("hexagon_api binaries not found, please define
HEXAGON_RPC_LIB_DIR")
return pathlib.Path(HEXAGON_RPC_LIB_DIR)
diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py
index 46cda68106..acd86c4445 100644
--- a/python/tvm/contrib/tvmjs.py
+++ b/python/tvm/contrib/tvmjs.py
@@ -37,9 +37,8 @@ except ImportError:
ml_dtypes = None
import tvm
-from tvm.libinfo import find_lib_path
from tvm.runtime import DataType
-from tvm.support.emcc import create_tvmjs_wasm
+from tvm.support.emcc import create_tvmjs_wasm, find_wasm_lib
def _convert_f32_to_bf16(value):
@@ -409,11 +408,11 @@ def export_runtime(runtime_dir):
+ " cd /path/to/tvm/web; make; npm run bundle"
)
- jsbundle = find_lib_path("tvmjs.bundle.js", optional=True)
+ jsbundle = find_wasm_lib("tvmjs.bundle.js", optional=True)
if not jsbundle:
raise RuntimeError("Cannot find tvmjs.bundle.js, " + web_hint)
- wasi = find_lib_path("tvmjs_runtime.wasi.js", optional=True)
+ wasi = find_wasm_lib("tvmjs_runtime.wasi.js", optional=True)
if not wasi:
raise RuntimeError("Cannot find tvmjs_runtime.wasi.js, " + web_hint)
diff --git a/python/tvm/libinfo.py b/python/tvm/libinfo.py
index d91db1f572..d56021df4e 100644
--- a/python/tvm/libinfo.py
+++ b/python/tvm/libinfo.py
@@ -14,24 +14,24 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Library information."""
+"""Library and include path information for TVM."""
from __future__ import annotations
import os
-import sys
from pathlib import Path
-from tvm_ffi.libinfo import load_lib_ctypes
+from tvm_ffi import libinfo as tvm_ffi_libinfo
-def use_runtime_lib() -> bool:
- """Whether ``TVM_USE_RUNTIME_LIB`` requests runtime-only mode.
+def _rel_top_directory() -> Path:
+ """Top directory in the installed (wheel) layout: ``python/tvm/``."""
+ return Path(__file__).parent
- Recognises ``1`` / ``true`` / ``yes`` (case-insensitive) as truthy.
- Anything else — including ``0`` and the unset case — is False.
- """
- return os.environ.get("TVM_USE_RUNTIME_LIB", "0").lower() in ("1", "true",
"yes")
+
+def _dev_top_directory() -> Path:
+ """Top directory in the source/dev layout: the worktree root."""
+ return _rel_top_directory() / ".." / ".."
def package_lib_paths() -> list[Path]:
@@ -44,283 +44,50 @@ def package_lib_paths() -> list[Path]:
pick the basenames they want (e.g. ``libtvm_runtime.so``) and the load
mode; this function only returns the search path.
"""
- pkg = Path(__file__).parent # python/tvm/
+ pkg = _rel_top_directory() # python/tvm/
paths: list[Path] = []
if os.environ.get("TVM_LIBRARY_PATH"):
paths.append(Path(os.environ["TVM_LIBRARY_PATH"]))
paths += [
pkg / "lib", # wheel layout
- pkg.parent.parent / "build" / "lib", # dev: <worktree>/build/lib
- pkg.parent.parent / "lib", # dev: <worktree>/lib
+ _dev_top_directory() / "build" / "lib", # dev: <worktree>/build/lib
+ _dev_top_directory() / "lib", # dev: <worktree>/lib
]
return paths
-_BACKEND_RUNTIME_LIBS = ["cuda", "vulkan", "opencl", "metal", "rocm",
"hexagon", "extra"]
-
-
-def load_backend_libs(runtime_lib_path: str) -> None:
- """Try to load each known backend runtime DSO; failures are silent."""
- runtime_dir = Path(runtime_lib_path).resolve().parent
- for backend in _BACKEND_RUNTIME_LIBS:
- try:
- load_lib_ctypes(
- package="tvm",
- target_name=f"tvm_runtime_{backend}",
- mode="RTLD_GLOBAL",
- extra_lib_paths=[runtime_dir],
- )
- except (OSError, FileNotFoundError, RuntimeError):
- pass
-
-
-def split_env_var(env_var, split):
- """Splits environment variable string.
-
- Parameters
- ----------
- env_var : str
- Name of environment variable.
-
- split : str
- String to split env_var on.
+def find_libtvm_runtime() -> str:
+ """Return the path to the ``libtvm_runtime`` shared library.
Returns
-------
- splits : list(string)
- If env_var exists, split env_var. Otherwise, empty list.
- """
- if os.environ.get(env_var, None):
- return [p.strip() for p in os.environ[env_var].split(split)]
- return []
-
-
-def get_dll_directories():
- """Get the possible dll directories"""
- # NB: This will either be the source directory (if TVM is run
- # inplace) or the install directory (if TVM is installed).
- # An installed TVM's curr_path will look something like:
- # $PREFIX/lib/python3.6/site-packages/tvm/_ffi
- ffi_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
- source_dir = os.path.join(ffi_dir, "..", "..")
- install_lib_dir = os.path.join(ffi_dir, "..", "..", "..")
-
- dll_path = []
-
- if os.environ.get("TVM_LIBRARY_PATH", None):
- dll_path.append(os.environ["TVM_LIBRARY_PATH"])
-
- if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
- dll_path.extend(split_env_var("LD_LIBRARY_PATH", ":"))
- dll_path.extend(split_env_var("PATH", ":"))
- elif sys.platform.startswith("darwin"):
- dll_path.extend(split_env_var("DYLD_LIBRARY_PATH", ":"))
- dll_path.extend(split_env_var("PATH", ":"))
- elif sys.platform.startswith("win32"):
- dll_path.extend(split_env_var("PATH", ";"))
-
- # Pip lib directory
- dll_path.append(ffi_dir)
- dll_path.append(os.path.join(ffi_dir, "lib"))
- # Default CMake build directory: shared libs are placed under build/lib/
- # to mirror the tvm-ffi layout (so wheel install + dev-mode dlopen find
- # them via the same `lib/` subdir).
- dll_path.append(os.path.join(source_dir, "build", "lib"))
- dll_path.append(os.path.join(source_dir, "build", "lib", "Release"))
- dll_path.append(os.path.join(source_dir, "build"))
- dll_path.append(os.path.join(source_dir, "build", "Release"))
- # Default make build directory
- dll_path.append(os.path.join(source_dir, "lib"))
-
- dll_path.append(install_lib_dir)
-
- # use extra TVM_HOME environment for finding libraries.
- if os.environ.get("TVM_HOME", None):
- tvm_source_home_dir = os.environ["TVM_HOME"]
- else:
- tvm_source_home_dir = source_dir
-
- if os.path.isdir(tvm_source_home_dir):
- dll_path.append(os.path.join(tvm_source_home_dir, "web", "dist",
"wasm"))
- dll_path.append(os.path.join(tvm_source_home_dir, "web", "dist"))
-
- dll_path = [os.path.realpath(x) for x in dll_path]
- return [x for x in dll_path if os.path.isdir(x)]
-
-
-def find_lib_path(name=None, search_path=None, optional=False):
- """Find dynamic library files.
-
- Parameters
- ----------
- name : list of str
- List of names to be found.
+ path : str
+ The resolved path to the TVM runtime shared library.
- Returns
- -------
- lib_path : list(string)
- List of all found path to the libraries
+ Raises
+ ------
+ RuntimeError
+ If the runtime library cannot be found.
"""
- use_runtime = use_runtime_lib()
- dll_path = get_dll_directories()
- # When the caller asks for a specific ``name`` we honour it directly
- # regardless of TVM_USE_RUNTIME_LIB; that env var is interpreted by
- # ``base.py::_load_lib`` to choose which name to ask for. This avoids
- # the runtime/compiler dual-list logic below from making `name` paths
- # unreachable when the user sets TVM_USE_RUNTIME_LIB.
- if name is not None:
- use_runtime = False
-
- if search_path is not None:
- if isinstance(search_path, list):
- dll_path = dll_path + search_path
- else:
- dll_path.append(search_path)
-
- if name is not None:
- if isinstance(name, list):
- lib_dll_path = []
- for n in name:
- lib_dll_path += [os.path.join(p, n) for p in dll_path]
- else:
- lib_dll_path = [os.path.join(p, name) for p in dll_path]
- runtime_dll_path = []
- ext_lib_dll_path = []
- else:
- if sys.platform.startswith("win32"):
- lib_dll_names = ["libtvm_compiler.dll", "tvm_compiler.dll"]
- runtime_dll_names = ["libtvm_runtime.dll", "tvm_runtime.dll"]
- ext_lib_dll_names = [
-
"3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.dll",
- "3rdparty/libflash_attn/src/libflash_attn.dll",
- ]
- elif sys.platform.startswith("darwin"):
- lib_dll_names = ["libtvm_compiler.dylib"]
- runtime_dll_names = ["libtvm_runtime.dylib"]
- ext_lib_dll_names = []
- else:
- lib_dll_names = ["libtvm_compiler.so"]
- runtime_dll_names = ["libtvm_runtime.so"]
- ext_lib_dll_names = [
-
"3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.so",
- "3rdparty/libflash_attn/src/libflash_attn.so",
- ]
-
- name = lib_dll_names + runtime_dll_names + ext_lib_dll_names
- lib_dll_path = [
- os.path.join(p, name)
- for name in lib_dll_names
- for p in dll_path
- if not p.endswith("python/tvm")
- ]
- runtime_dll_path = [
- os.path.join(p, name)
- for name in runtime_dll_names
- for p in dll_path
- if not p.endswith("python/tvm")
- ]
- ext_lib_dll_path = [os.path.join(p, name) for name in
ext_lib_dll_names for p in dll_path]
- if not use_runtime:
- # try to find lib_dll_path
- lib_found = [p for p in lib_dll_path if os.path.exists(p) and
os.path.isfile(p)]
- lib_found += [p for p in runtime_dll_path if os.path.exists(p) and
os.path.isfile(p)]
- lib_found += [p for p in ext_lib_dll_path if os.path.exists(p) and
os.path.isfile(p)]
- else:
- # try to find runtime_dll_path
- use_runtime = True
- lib_found = [p for p in runtime_dll_path if os.path.exists(p) and
os.path.isfile(p)]
-
- if not lib_found:
- if not optional:
- message = (
- f"Cannot find libraries: {name}\n"
- + "List of candidates:\n"
- + "\n".join(lib_dll_path + runtime_dll_path)
- )
- raise RuntimeError(message)
- return None
-
- if use_runtime:
- sys.stderr.write(f"Loading runtime library {lib_found[0]}... exec
only\n")
- sys.stderr.flush()
- return lib_found
-
-
-def find_include_path(name=None, search_path=None, optional=False):
- """Find header files for C compilation.
-
- Parameters
- ----------
- name : list of str
- List of directory names to be searched.
-
- Returns
- -------
- include_path : list(string)
- List of all found paths to header files.
- """
- if os.environ.get("TVM_SOURCE_DIR", None):
- source_dir = os.environ["TVM_SOURCE_DIR"]
- elif os.environ.get("TVM_HOME", None):
- source_dir = os.environ["TVM_HOME"]
- else:
- ffi_dir =
os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
- for source_dir in ["..", "../..", "../../.."]:
- source_dir = os.path.join(ffi_dir, source_dir)
- if os.path.isdir(os.path.join(source_dir, "include")):
- break
- else:
- raise AssertionError(f"Cannot find the source directory given
ffi_dir: {ffi_dir}")
- third_party_dir = os.path.join(source_dir, "3rdparty")
-
- header_path = []
-
- if os.environ.get("TVM_INCLUDE_PATH", None):
- header_path.append(os.environ["TVM_INCLUDE_PATH"])
-
- header_path.append(source_dir)
- header_path.append(third_party_dir)
-
- header_path = [os.path.abspath(x) for x in header_path]
- if search_path is not None:
- if isinstance(search_path, list):
- header_path = header_path + search_path
- else:
- header_path.append(search_path)
- if name is not None:
- if isinstance(name, list):
- tvm_include_path = []
- for n in name:
- tvm_include_path += [os.path.join(p, n) for p in header_path]
- else:
- tvm_include_path = [os.path.join(p, name) for p in header_path]
- dlpack_include_path = []
- else:
- tvm_include_path = [os.path.join(p, "include") for p in header_path]
- tvm_ffi_include_path = [
- os.path.join(p, "3rdparty", "tvm-ffi", "include") for p in
header_path
- ]
- dlpack_include_path = [
- os.path.join(p, "3rdparty", "tvm-ffi", "3rdparty", "dlpack",
"include")
- for p in header_path
- ]
-
- # try to find include path
- include_found = [p for p in tvm_include_path if os.path.exists(p) and
os.path.isdir(p)]
- include_found += [p for p in tvm_ffi_include_path if os.path.exists(p)
and os.path.isdir(p)]
- include_found += [p for p in dlpack_include_path if os.path.exists(p)
and os.path.isdir(p)]
-
- if not include_found:
- message = (
- "Cannot find the files.\n"
- + "List of candidates:\n"
- + str("\n".join(tvm_include_path + dlpack_include_path))
- )
- if not optional:
- raise RuntimeError(message)
- return None
-
- return include_found
+ candidate = tvm_ffi_libinfo._find_library_by_basename(
+ "tvm", "tvm_runtime", extra_lib_paths=package_lib_paths()
+ )
+ if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda
_: True):
+ return ret
+ raise RuntimeError("Cannot find libtvm_runtime")
+
+
+def find_include_path() -> str:
+ """Return the path to TVM's own ``include/`` directory."""
+ if ret := tvm_ffi_libinfo._resolve_and_validate(
+ paths=[
+ _rel_top_directory() / "include",
+ _dev_top_directory() / "include",
+ ],
+ cond=lambda p: (p / "tvm" / "runtime").is_dir(),
+ ):
+ return ret
+ raise RuntimeError("Cannot find TVM include path.")
# The version is written by setuptools_scm into _version.py at build time
diff --git a/python/tvm/rpc/minrpc.py b/python/tvm/rpc/minrpc.py
index f3819c2791..d46f2a2faf 100644
--- a/python/tvm/rpc/minrpc.py
+++ b/python/tvm/rpc/minrpc.py
@@ -48,7 +48,7 @@ def find_minrpc_server_libpath(server="posix_popen_server"):
return minrpc_dir, path
-def with_minrpc(compile_func, server="posix_popen_server",
runtime="libtvm_runtime"):
+def with_minrpc(compile_func, server="posix_popen_server"):
"""Attach the compiler function with minrpc related options.
Parameters
@@ -59,16 +59,13 @@ def with_minrpc(compile_func, server="posix_popen_server",
runtime="libtvm_runti
server : str
The server type.
- runtime : str
- The runtime library.
-
Returns
-------
fcompile : function
The return compilation.
"""
minrpc_dir, server_path = find_minrpc_server_libpath(server)
- runtime_path = libinfo.find_lib_path([runtime, runtime + ".so", runtime +
".dylib"])[0]
+ runtime_path = libinfo.find_libtvm_runtime()
tvm_ffi_path = tvm_ffi.libinfo.find_libtvm_ffi()
runtime_dir = os.path.abspath(os.path.dirname(runtime_path))
@@ -86,7 +83,12 @@ def with_minrpc(compile_func, server="posix_popen_server",
runtime="libtvm_runti
# inside libtvm_runtime; with the default ``--as-needed`` those
# registrations would never run.
options += ["-Wl,--no-as-needed"]
- options += ["-I" + path for path in libinfo.find_include_path()]
+ default_include_paths = [
+ libinfo.find_include_path(),
+ tvm_ffi.libinfo.find_include_path(),
+ tvm_ffi.libinfo.find_dlpack_include_path(),
+ ]
+ options += ["-I" + path for path in default_include_paths]
options += ["-I" + minrpc_dir]
fcompile = cc.cross_compiler(
compile_func, options=options, add_files=[server_path, runtime_path,
tvm_ffi_path]
diff --git a/python/tvm/rpc/server.py b/python/tvm/rpc/server.py
index 099cb8f1f1..4a28da6726 100644
--- a/python/tvm/rpc/server.py
+++ b/python/tvm/rpc/server.py
@@ -38,11 +38,11 @@ import struct
import sys
import threading
import time
+from pathlib import Path
import tvm_ffi
from tvm.base import py_str
-from tvm.libinfo import find_lib_path
from tvm.runtime.module import load_module as _load_module
from tvm.support import utils
from tvm.support.popen_pool import PopenWorker
@@ -115,7 +115,14 @@ def _server_env(load_library, work_path=None):
libs = []
load_library = load_library.split(":") if load_library else []
for file_name in load_library:
- file_name = find_lib_path(file_name)[0]
+ # Resolve the literal library file name against the current working
+ # directory, preserving the exact filename the caller requested.
+ resolved = tvm_ffi.libinfo._resolve_and_validate(
+ [Path.cwd() / file_name], cond=lambda p: p.is_file()
+ )
+ if resolved is None:
+ raise RuntimeError(f"Cannot find library {file_name}")
+ file_name = resolved
libs.append(ctypes.CDLL(file_name, ctypes.RTLD_GLOBAL))
logger.info("Load additional library %s", file_name)
temp.libs = libs
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 8ecda773b9..8fde8cc443 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -28,6 +28,7 @@ import numpy as np
from tvm_ffi import (
Module as _Module,
)
+from tvm_ffi import libinfo as tvm_ffi_libinfo
from tvm_ffi import (
load_module as _load_module,
)
@@ -38,8 +39,8 @@ from tvm_ffi import (
system_lib,
)
+import tvm.libinfo
from tvm.base import _RUNTIME_ONLY
-from tvm.libinfo import find_include_path
from . import _ffi_api
@@ -311,7 +312,12 @@ class Module(_Module):
if "options" in kwargs:
opts = kwargs["options"]
options = opts if isinstance(opts, list | tuple) else [opts]
- opts = options + ["-I" + path for path in find_include_path()]
+ default_include_paths = [
+ tvm.libinfo.find_include_path(),
+ tvm_ffi_libinfo.find_include_path(),
+ tvm_ffi_libinfo.find_dlpack_include_path(),
+ ]
+ opts = options + ["-I" + path for path in default_include_paths]
kwargs.update({"options": opts})
return fcompile(file_name, files, **kwargs)
diff --git a/python/tvm/support/emcc.py b/python/tvm/support/emcc.py
index c2f5e6dbbc..9bd6d24036 100644
--- a/python/tvm/support/emcc.py
+++ b/python/tvm/support/emcc.py
@@ -21,8 +21,46 @@ import os
import subprocess
from pathlib import Path
+from tvm import libinfo
from tvm.base import py_str
-from tvm.libinfo import find_lib_path
+
+
+def find_wasm_lib(name, optional=False):
+ """Find a wasm/web asset (e.g. a ``.bc`` or ``.js`` file) under
``web/dist``.
+
+ Specializes the asset search to the wasm build outputs that live under
+ ``web/dist/wasm`` and ``web/dist``, anchored on ``TVM_HOME`` (when set) or
+ the TVM source root otherwise.
+
+ Parameters
+ ----------
+ name : str
+ The asset file name to look for.
+
+ optional : bool
+ When True, return None instead of raising if the asset is not found.
+
+ Returns
+ -------
+ found : list(str) or None
+ List of matching paths, or None when ``optional`` and nothing is found.
+ """
+ # use extra TVM_HOME environment for finding libraries.
+ if os.environ.get("TVM_HOME", None):
+ tvm_source_home_dir = os.environ["TVM_HOME"]
+ else:
+ tvm_source_home_dir = os.fspath(libinfo._dev_top_directory())
+
+ dll_path = [
+ os.path.join(tvm_source_home_dir, "web", "dist", "wasm"),
+ os.path.join(tvm_source_home_dir, "web", "dist"),
+ ]
+ found = [os.path.join(p, name) for p in dll_path if
os.path.isfile(os.path.join(p, name))]
+ if not found:
+ if optional:
+ return None
+ raise RuntimeError(f"Cannot find {name}; searched {dll_path}")
+ return found
def create_tvmjs_wasm(output, objects, options=None, cc="emcc", libs=None):
@@ -72,10 +110,10 @@ def create_tvmjs_wasm(output, objects, options=None,
cc="emcc", libs=None):
all_libs = []
if not with_runtime:
- all_libs += [find_lib_path("wasm_runtime.bc")[0]]
+ all_libs += [find_wasm_lib("wasm_runtime.bc")[0]]
- all_libs += [find_lib_path("tvmjs_support.bc")[0]]
- all_libs += [find_lib_path("webgpu_runtime.bc")[0]]
+ all_libs += [find_wasm_lib("tvmjs_support.bc")[0]]
+ all_libs += [find_wasm_lib("webgpu_runtime.bc")[0]]
if libs:
if not isinstance(libs, list):
diff --git a/tests/python/relax/test_frontend_nn_extern_module.py
b/tests/python/relax/test_frontend_nn_extern_module.py
index a304766399..dba87c3fde 100644
--- a/tests/python/relax/test_frontend_nn_extern_module.py
+++ b/tests/python/relax/test_frontend_nn_extern_module.py
@@ -24,6 +24,7 @@ import numpy as np
import tvm_ffi
import tvm
+import tvm.libinfo
import tvm.testing
from tvm import relax
from tvm.relax.frontend import nn
@@ -122,12 +123,16 @@ def _check_ir_equality(mod):
def _compile_cc(src: Path, dst: Path):
# pylint: disable=import-outside-toplevel
from tvm.base import py_str
- from tvm.libinfo import find_include_path
# pylint: enable=import-outside-toplevel
cmd = ["g++", str(src)]
- for include_path in find_include_path():
+ default_include_paths = [
+ tvm.libinfo.find_include_path(),
+ tvm_ffi.libinfo.find_include_path(),
+ tvm_ffi.libinfo.find_dlpack_include_path(),
+ ]
+ for include_path in default_include_paths:
cmd += ["-I", include_path]
cmd += [
"-c",
