You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/15 12:35:14 UTC
(tvm) 02/03: [Python] Add lightweight runtime build configuration queries
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch jenkins-wheel-test-artifacts
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit d795529e717d2cc92ec2163648ce4ebbede92d1f
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 12:34:58 2026 +0000
[Python] Add lightweight runtime build configuration queries
---
pyproject.toml | 3 +
python/tvm/__init__.py | 182 +++++++++++++++++++-----------------
python/tvm/config.py | 63 +++++++++++++
python/tvm/libinfo.py | 39 ++++----
tests/python/runtime/test_config.py | 68 ++++++++++++++
5 files changed, 255 insertions(+), 100 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 5c7512032f..3a19c57d50 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,6 +57,9 @@ dependencies = [
"typing_extensions",
]
+[project.scripts]
+tvm-config = "tvm.config:__main__"
+
[project.optional-dependencies]
torch = ["torch"]
cuda = ["cuda-bindings"]
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index e6381cbe3b..2c35d829f2 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -18,120 +18,134 @@
# pylint: disable=redefined-builtin, wildcard-import
"""TVM: Open Deep Learning Compiler Stack."""
-import multiprocessing
import os
import sys
-# ffi module must load first
-from tvm_ffi import register_object, register_global_func, get_global_func
+from typing import TYPE_CHECKING
-# top-level alias
-from .libinfo import __version__
-from .base import _RUNTIME_ONLY
-# tvm.runtime
-from .runtime import Object
-from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
-from .runtime._tensor import vpi, rocm, ext_dev, hexagon
-from .runtime import DataType, DataTypeCode
+def _is_config_mode() -> bool:
+ """Recognize the config CLI before loading Python or native TVM components."""
+ if sys.argv[0].endswith("tvm-config"):
+ return True
+ argv = getattr(sys, "orig_argv", [])
+ return any(arg == "-m" and argv[i + 1 : i + 2] == ["tvm.config"] for i, arg in enumerate(argv))
-# tvm.error
-from . import error
-# tvm.ir
-from .ir import IRModule, transform
-from . import ir
+if TYPE_CHECKING or not _is_config_mode():
+ import multiprocessing
-# tvm.script — must be imported before any dialect package so that
-# tvm.script.register_dialect is reachable when dialect __init__.py files run.
-from . import script
+ # ffi module must load first
+ from tvm_ffi import register_object, register_global_func, get_global_func
-# tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
-from . import tirx
-from .ir import prim
+ # top-level alias
+ from .libinfo import __version__
+ from .base import _RUNTIME_ONLY
-# tvm.backend — owns backend Python load hooks
-from . import backend
+ # tvm.runtime
+ from .runtime import Object
+ from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
+ from .runtime._tensor import vpi, rocm, ext_dev, hexagon
+ from .runtime import DataType, DataTypeCode
-# tvm.target
-from . import target
-from .target import device_from_target
+ # tvm.error
+ from . import error
-# Shared code generation helpers
-from . import codegen
+ # tvm.ir
+ from .ir import IRModule, transform
+ from . import ir
-# tvm.te
-from . import te
+ # tvm.script — must be imported before any dialect package so that
+ # tvm.script.register_dialect is reachable when dialect __init__.py files run.
+ from . import script
-# tvm.driver
-from .driver import compile
+ # tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
+ from . import tirx
+ from .ir import prim
-# others
-from . import sym
+ # tvm.backend — owns backend Python load hooks
+ from . import backend
-# support infra
-from . import support
+ # tvm.target
+ from . import target
+ from .target import device_from_target
-# Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup
-from .support import rocm as _rocm, nvcc as _nvcc
+ # Shared code generation helpers
+ from . import codegen
-# Relax contain modules that are only available in compiler package
-# Do not import them if TVM is built with runtime only
-if not _RUNTIME_ONLY:
- # tvm.relax — registers itself via tvm.script.register_dialect in its __init__
- from . import relax
+ # tvm.te
+ from . import te
-# NOTE: This file should be python2 compatible so we can
-# raise proper error message when user run the package using
-# an older version of the python
+ # tvm.driver
+ from .driver import compile
+ # others
+ from . import sym
-def _should_print_backtrace():
- in_pytest = "PYTEST_CURRENT_TEST" in os.environ
- tvm_backtrace = os.environ.get("TVM_BACKTRACE", "0")
+ # support infra
+ from . import support
- try:
- tvm_backtrace = bool(int(tvm_backtrace))
- except ValueError:
- raise ValueError(f"invalid value for TVM_BACKTRACE {tvm_backtrace}, please set to 0 or 1.")
+ # Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup
+ from .support import rocm as _rocm, nvcc as _nvcc
+
+ # Relax contain modules that are only available in compiler package
+ # Do not import them if TVM is built with runtime only
+ if not _RUNTIME_ONLY:
+ # tvm.relax — registers itself via tvm.script.register_dialect in its __init__
+ from . import relax
- return in_pytest or tvm_backtrace
+ # NOTE: This file should be python2 compatible so we can
+ # raise proper error message when user run the package using
+ # an older version of the python
+ def _should_print_backtrace():
+ in_pytest = "PYTEST_CURRENT_TEST" in os.environ
+ tvm_backtrace = os.environ.get("TVM_BACKTRACE", "0")
-def tvm_wrap_excepthook(exception_hook):
- """Wrap given excepthook with TVM additional work."""
+ try:
+ tvm_backtrace = bool(int(tvm_backtrace))
+ except ValueError:
+ raise ValueError(
+ f"invalid value for TVM_BACKTRACE {tvm_backtrace}, please set to 0 or 1."
+ )
- def wrapper(exctype, value, trbk):
- """Clean subprocesses when TVM is interrupted."""
- if exctype is error.DiagnosticError and not _should_print_backtrace():
- # TODO(@jroesch): consider moving to C++?
- print("note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.")
- else:
- exception_hook(exctype, value, trbk)
+ return in_pytest or tvm_backtrace
- if hasattr(multiprocessing, "active_children"):
- # pylint: disable=not-callable
- for p in multiprocessing.active_children():
- p.terminate()
+ def tvm_wrap_excepthook(exception_hook):
+ """Wrap given excepthook with TVM additional work."""
- return wrapper
+ def wrapper(exctype, value, trbk):
+ """Clean subprocesses when TVM is interrupted."""
+ if exctype is error.DiagnosticError and not _should_print_backtrace():
+ # TODO(@jroesch): consider moving to C++?
+ print(
+ "note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace."
+ )
+ else:
+ exception_hook(exctype, value, trbk)
+ if hasattr(multiprocessing, "active_children"):
+ # pylint: disable=not-callable
+ for p in multiprocessing.active_children():
+ p.terminate()
-sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
+ return wrapper
-# Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt
-# into being loaded automatically at ``import tvm`` time by declaring an entry
-# point in the ``tvm.backends`` group:
-# [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
-# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
-from .backend._autoload_backends import _autoload_backends
+ sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
-_autoload_backends()
+ # Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt
+ # into being loaded automatically at ``import tvm`` time by declaring an entry
+ # point in the ``tvm.backends`` group:
+ # [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
+ # Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
+ from .backend._autoload_backends import _autoload_backends
-# Load the optional ORC JIT extension for its global FFI registrations. LLVM
-# function lookup reports a targeted installation error if this import fails;
-# importing TVM itself remains silent for AOT-only workflows.
-try:
- import tvm_ffi_orcjit as _tvm_ffi_orcjit
-except Exception: # pylint: disable=broad-exception-caught
- pass
+ _autoload_backends()
+
+ # Load the optional ORC JIT extension for its global FFI registrations. LLVM
+ # function lookup reports a targeted installation error if this import fails;
+ # importing TVM itself remains silent for AOT-only workflows.
+ try:
+ import tvm_ffi_orcjit as _tvm_ffi_orcjit
+ except Exception: # pylint: disable=broad-exception-caught
+ pass
diff --git a/python/tvm/config.py b/python/tvm/config.py
new file mode 100644
index 0000000000..8a75f4469b
--- /dev/null
+++ b/python/tvm/config.py
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Lightweight compiler and linker configuration for the TVM runtime."""
+
+import argparse
+import sys
+from pathlib import Path
+
+from . import libinfo
+
+
+def __main__() -> None:
+ """Print paths and flags using the same conventions as tvm_ffi.config."""
+ parser = argparse.ArgumentParser(description="Get TVM runtime build configuration")
+ for name, help_text in (
+ ("includedir", "Print include directory"),
+ ("libdir", "Print library directory"),
+ ("libfiles", "Fully qualified library filenames"),
+ ("libs", "Libraries to be linked"),
+ ("cflags", "Print C flags"),
+ ("cxxflags", "Print C++ flags"),
+ ("ldflags", "Print linker flags"),
+ ):
+ parser.add_argument(f"--{name}", action="store_true", help=help_text)
+ args = parser.parse_args()
+ if len(sys.argv) == 1:
+ parser.print_help()
+ return
+ if args.includedir:
+ print(libinfo.find_include_path())
+ if args.libdir or args.libfiles or args.libs or args.ldflags:
+ library = Path(libinfo.find_libtvm_runtime())
+ link_file = library.with_suffix(".lib") if sys.platform == "win32" else library
+ if args.libdir:
+ print(library.parent)
+ if args.libfiles:
+ print(link_file)
+ if args.cxxflags:
+ print(f"-I{libinfo.find_include_path()} -std=c++17")
+ if args.cflags:
+ print(f"-I{libinfo.find_include_path()}")
+ if args.libs:
+ print(link_file if sys.platform == "win32" else "-ltvm_runtime")
+ if args.ldflags and sys.platform != "win32":
+ print(f"-L{library.parent}")
+
+
+if __name__ == "__main__":
+ __main__()
diff --git a/python/tvm/libinfo.py b/python/tvm/libinfo.py
index 312d769f7e..e651277aec 100644
--- a/python/tvm/libinfo.py
+++ b/python/tvm/libinfo.py
@@ -19,10 +19,9 @@
from __future__ import annotations
import os
+import sys
from pathlib import Path
-from tvm_ffi import libinfo as tvm_ffi_libinfo
-
def _rel_top_directory() -> Path:
"""Top directory in the installed (wheel) layout: ``python/tvm/``."""
@@ -69,24 +68,32 @@ def find_libtvm_runtime() -> str:
RuntimeError
If the runtime library cannot be 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")
+ if sys.platform.startswith("win32"):
+ names = ["tvm_runtime.dll"]
+ search_var = "PATH"
+ elif sys.platform == "darwin":
+ names = ["libtvm_runtime.dylib", "libtvm_runtime.so"]
+ search_var = "DYLD_LIBRARY_PATH"
+ else:
+ names = ["libtvm_runtime.so"]
+ search_var = "LD_LIBRARY_PATH"
+ paths = package_lib_paths()
+ paths.extend(Path(p) for p in os.environ.get(search_var, "").split(os.pathsep) if p)
+ if search_var != "PATH":
+ paths.extend(Path(p) for p in os.environ.get("PATH", "").split(os.pathsep) if p)
+ for directory in paths:
+ for name in names:
+ candidate = directory / name
+ if candidate.is_file():
+ return str(candidate.resolve())
+ raise RuntimeError(f"Cannot find libtvm_runtime in {paths}")
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
+ for path in (_rel_top_directory() / "include", _dev_top_directory() / "include"):
+ if (path / "tvm" / "runtime").is_dir():
+ return str(path.resolve())
raise RuntimeError("Cannot find TVM include path.")
diff --git a/tests/python/runtime/test_config.py b/tests/python/runtime/test_config.py
new file mode 100644
index 0000000000..238c083e7d
--- /dev/null
+++ b/tests/python/runtime/test_config.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The config command must work without importing native runtime dependencies."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import tvm
+
+
+def test_config_without_native_imports(tmp_path):
+ # A startup hook rejects native loads and dependencies even when they are installed.
+ (tmp_path / "sitecustomize.py").write_text(
+ "import sys\n"
+ "class BlockNative:\n"
+ " def find_spec(self, fullname, path=None, target=None):\n"
+ " if fullname.split('.')[0] in ('tvm_ffi', 'torch', 'numpy') or "
+ "fullname in ('tvm.base', 'tvm.runtime', 'tvm.backend', 'tvm.driver'):\n"
+ " raise AssertionError(fullname)\n"
+ "sys.meta_path.insert(0, BlockNative())\n"
+ "def audit(event, args):\n"
+ " if event == 'ctypes.dlopen': raise AssertionError(args)\n"
+ "sys.addaudithook(audit)\n"
+ )
+ env = dict(os.environ)
+ env["PYTHONPATH"] = os.pathsep.join((str(tmp_path), str(Path(tvm.__file__).parent.parent)))
+ command = [sys.executable, "-m", "tvm.config"]
+ include = subprocess.check_output([*command, "--includedir"], env=env, text=True).strip()
+ assert (Path(include) / "tvm" / "runtime").is_dir()
+ library = subprocess.check_output([*command, "--libfiles"], env=env, text=True).strip()
+ assert Path(library).is_file()
+ flags = subprocess.check_output(
+ [*command, "--cflags", "--cxxflags", "--libs", "--ldflags"], env=env, text=True
+ )
+ assert f"-I{include}" in flags
+ assert "-std=c++17" in flags
+
+
+def test_runtime_library_path_fallback(tmp_path, monkeypatch):
+ library = tmp_path / Path(tvm.libinfo.find_libtvm_runtime()).name
+ library.touch()
+ monkeypatch.setattr(tvm.libinfo, "package_lib_paths", lambda: [])
+ monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
+ monkeypatch.delenv("DYLD_LIBRARY_PATH", raising=False)
+ monkeypatch.setenv("PATH", str(tmp_path))
+ assert tvm.libinfo.find_libtvm_runtime() == str(library)
+
+
+if __name__ == "__main__":
+ import tvm.testing
+
+ tvm.testing.main()