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 <[email protected]> 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()
