This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 80d645d82403170018c27e65c323a226433f8e3b Author: Tianqi Chen <[email protected]> AuthorDate: Fri Oct 16 16:25:56 2026 -0400 [TIRx] Simplify Python APIs and configuration discovery (#20766) Organize TIRx's Python APIs by moving IndexMap into its own module and LLVM helpers under `codegen.llvm`. Use the existing s_tir attribute checks directly, remove the empty dialect backend package, and place public operator declarations before private helpers. Centralize configuration discovery in `libinfo.py` and share it with package and native configuration commands. Replace legacy exception classes with builtin exceptions while retaining FFI error messages and subprocess cleanup. Remove the environment-probe test module and move wheel library validation into the runtime suite with updated packaging and CI selectors. --- tirx/CMakeLists.txt | 9 +- .../__init__.py => cmake/tvm-tirx-config.in} | 10 +- tirx/pyproject.toml | 2 +- tirx/python/tvm_tirx/__init__.py | 28 +- tirx/python/tvm_tirx/backend/metal/target_tags.py | 4 +- tirx/python/tvm_tirx/codegen/__init__.py | 204 +------------ .../tvm_tirx/codegen/{__init__.py => llvm.py} | 72 +---- tirx/python/tvm_tirx/config.py | 152 +--------- tirx/python/tvm_tirx/error.py | 126 -------- tirx/python/tvm_tirx/ir/op.py | 230 +++++++------- tirx/python/tvm_tirx/libinfo.py | 149 ++++++++- tirx/python/tvm_tirx/support/clang.py | 2 +- tirx/python/tvm_tirx/support/rocm.py | 2 +- tirx/python/tvm_tirx/testing/env.py | 8 +- tirx/python/tvm_tirx/tirx/__init__.py | 4 +- tirx/python/tvm_tirx/tirx/function.py | 297 +----------------- .../tvm_tirx/tirx/{function.py => index_map.py} | 169 +---------- tirx/python/tvm_tirx/tirx/op/__init__.py | 4 +- tirx/python/tvm_tirx/tirx/script/ir_builder/op.py | 2 +- .../tirx/script/ir_builder/parser_protocol.py | 2 +- .../python/backend/metal/codegen/test_basic.py | 12 +- .../python/backend/webgpu/codegen/test_basic.py | 20 +- tirx/tests/python/codegen/llvm/test_aarch64.py | 36 +-- tirx/tests/python/codegen/llvm/test_basic.py | 22 +- tirx/tests/python/codegen/llvm/test_riscv.py | 2 +- .../python/codegen/llvm/test_scalable_vector.py | 8 +- tirx/tests/python/codegen/llvm/test_x86.py | 2 +- .../test_validate_wheel_library.py} | 4 +- tirx/tests/python/script/test_basic_usage.py | 4 +- tirx/tests/python/support/test_config.py | 11 +- .../tests/python/target/test_llvm_features_info.py | 34 +-- tirx/tests/python/target/test_riscv_features.py | 6 +- tirx/tests/python/target/test_x86_features.py | 8 +- tirx/tests/python/testing/test_env.py | 332 --------------------- .../tirx/analysis/test_check_well_defined.py | 24 +- tirx/tests/python/tirx/base/test_imm_values.py | 4 +- .../python/tirx/base/test_scalable_datatype.py | 8 +- tirx/tests/python/tirx/codegen/test_cuda.py | 6 +- tirx/tests/python/tirx/codegen/test_hopper.py | 2 +- tirx/tests/python/tirx/codegen/test_ptx_addr.py | 4 +- tirx/tests/python/tirx/codegen/test_ptx_dialect.py | 2 +- tirx/tests/python/tirx/script/test_basic_usage.py | 8 +- tirx/tests/python/tirx/test_cuda_backend_config.py | 2 +- tirx/tests/python/tirx/tile/test_dispatcher.py | 2 +- .../tests/python/tirx/transform/test_lower_tirx.py | 2 +- .../python/tirx/transform/test_make_packed_api.py | 4 +- .../tirx/transform/test_split_host_device.py | 2 +- tirx/tests/python/tirx/transform/test_vectorize.py | 16 +- 48 files changed, 448 insertions(+), 1615 deletions(-) diff --git a/tirx/CMakeLists.txt b/tirx/CMakeLists.txt index 4001e955d9..d2581b2b5f 100644 --- a/tirx/CMakeLists.txt +++ b/tirx/CMakeLists.txt @@ -292,14 +292,19 @@ install(TARGETS tvm_tirx_header tvm_tirx_compiler tvm_runtime tvm_runtime_static # CMake package and dependency helpers install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm_tirx-config.cmake" "${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_tirx/config.py" + "${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_tirx/libinfo.py" DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx" COMPONENT tvm_tirx) install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/" DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx/utils" COMPONENT tvm_tirx FILES_MATCHING PATTERN "Find*.cmake") if(NOT TVM_TIRX_BUILD_PYTHON_MODULE) - install(PROGRAMS "${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_tirx/config.py" - DESTINATION "${CMAKE_INSTALL_BINDIR}" RENAME tvm-tirx-config COMPONENT tvm_tirx) + file(RELATIVE_PATH TVM_TIRX_CONFIG_PREFIX_RELATIVE + "${CMAKE_INSTALL_FULL_BINDIR}" "${CMAKE_INSTALL_PREFIX}") + configure_file("${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm-tirx-config.in" + "${CMAKE_CURRENT_BINARY_DIR}/tvm-tirx-config" @ONLY) + install(PROGRAMS "${CMAKE_CURRENT_BINARY_DIR}/tvm-tirx-config" + DESTINATION "${CMAKE_INSTALL_BINDIR}" COMPONENT tvm_tirx) endif() # Python package artifacts and licenses diff --git a/tirx/python/tvm_tirx/tirx/backend/__init__.py b/tirx/cmake/tvm-tirx-config.in similarity index 69% rename from tirx/python/tvm_tirx/tirx/backend/__init__.py rename to tirx/cmake/tvm-tirx-config.in index 6f5333629f..ebe607d13e 100644 --- a/tirx/python/tvm_tirx/tirx/backend/__init__.py +++ b/tirx/cmake/tvm-tirx-config.in @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 # 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 @@ -14,4 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""TIRx backend compatibility package.""" +"""Run the configuration helper installed with the native CMake package.""" + +import runpy +from pathlib import Path + +prefix = (Path(__file__).resolve().parent / "@TVM_TIRX_CONFIG_PREFIX_RELATIVE@").resolve() +config = runpy.run_path(str(prefix / "@CMAKE_INSTALL_LIBDIR@" / "cmake/tvm_tirx/config.py")) +config["__main__"]() diff --git a/tirx/pyproject.toml b/tirx/pyproject.toml index 4ef79dcb0c..333c1e8696 100644 --- a/tirx/pyproject.toml +++ b/tirx/pyproject.toml @@ -143,7 +143,7 @@ markers = [ skip = "*-win32 *-manylinux_i686 *-musllinux*" build-verbosity = 1 test-requires = ["pytest", "numpy"] -test-command = "python -I -c \"import importlib.util; assert importlib.util.find_spec('tvm') is None\" && python -I -m pytest -vvs {package}/tests/python/all-platform-minimal-test {package}/tests/python/runtime/test_tensor.py {package}/tests/python/ir/test_type.py {package}/tests/python/script/test_ir_builder.py {package}/tests/python/codegen/llvm/test_basic.py::test_llvm_vadd_pipeline {package}/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile" +test-command = "python -I -c \"import importlib.util; assert importlib.util.find_spec('tvm') is None\" && python -I -m pytest -vvs {package}/tests/python/runtime/test_validate_wheel_library.py {package}/tests/python/runtime/test_tensor.py {package}/tests/python/ir/test_type.py {package}/tests/python/script/test_ir_builder.py {package}/tests/python/codegen/llvm/test_basic.py::test_llvm_vadd_pipeline {package}/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile" [tool.cibuildwheel.linux] repair-wheel-command = "auditwheel repair --exclude libtvm_ffi.so --exclude libtvm_runtime_cuda.so --exclude 'libcuda.so.*' --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' --exclude 'libnvrtc-builtins.so.*' -w {dest_dir} {wheel}" diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index fcdeb1f0bd..8244906d9c 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -49,9 +49,6 @@ if TYPE_CHECKING or not _is_config_mode(): from .runtime._tensor import vpi, rocm, ext_dev, hexagon from .runtime import DataType, DataTypeCode - # tvm.error - from . import error - # tvm.ir from .ir import IRModule, transform from . import ir @@ -83,35 +80,12 @@ if TYPE_CHECKING or not _is_config_mode(): # Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup from .support import rocm as _rocm, nvcc as _nvcc - # 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") - - 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." - ) - - return in_pytest or tvm_backtrace - def tvm_wrap_excepthook(exception_hook): """Wrap given excepthook with TVM additional work.""" 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) + exception_hook(exctype, value, trbk) if hasattr(multiprocessing, "active_children"): # pylint: disable=not-callable diff --git a/tirx/python/tvm_tirx/backend/metal/target_tags.py b/tirx/python/tvm_tirx/backend/metal/target_tags.py index 30f406c963..b1073c0e1e 100644 --- a/tirx/python/tvm_tirx/backend/metal/target_tags.py +++ b/tirx/python/tvm_tirx/backend/metal/target_tags.py @@ -23,9 +23,9 @@ _METAL_HOST_TRIPLE = "arm64-apple-macos" def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu): try: - from tvm_tirx.codegen import llvm_is_valid_cpu + from tvm_tirx.codegen import llvm - if not llvm_is_valid_cpu(mcpu, _METAL_HOST_TRIPLE): + if not llvm.is_valid_cpu(mcpu, _METAL_HOST_TRIPLE): return except Exception: # pylint: disable=broad-except pass # LLVM not available; register unconditionally diff --git a/tirx/python/tvm_tirx/codegen/__init__.py b/tirx/python/tvm_tirx/codegen/__init__.py index afc5fc7c3f..5396028ac1 100644 --- a/tirx/python/tvm_tirx/codegen/__init__.py +++ b/tirx/python/tvm_tirx/codegen/__init__.py @@ -20,6 +20,7 @@ from tvm_ffi import Array, get_global_func from tvm_tirx.target import Target +from . import llvm from .executable import Executable, load_static_library @@ -65,206 +66,3 @@ def target_has_features(cpu_features, target=None): for feat in cpu_features: has_feats &= get_global_func("codegen.target_has_feature")(feat, target) return has_feats - - -def llvm_lookup_intrinsic_id(name): - """Lookup LLVM intrinsic id by name. - - Parameters - ---------- - name : str - The name of the intrinsic. - - Returns - ------- - intrin_id : int - The intrinsic id. - """ - return get_global_func("codegen.llvm_lookup_intrinsic_id")(name) - - -def llvm_get_intrinsic_name(intrin_id: int) -> str: - """Get the name of an intrinsic for a given id. - - Parameters - ---------- - intrin_id : int - The id of the intrinsic. - - Returns - ------- - name : str - The name of the intrinsic. - """ - return get_global_func("codegen.llvm_get_intrinsic_name")(intrin_id) - - -def llvm_get_system_x86_vendor(): - """Get system x86 vendor info. - - Parameters - ---------- - - Returns - ------- - vendor : str - The current system's cpu vendor. - """ - return get_global_func("codegen.llvm_get_system_x86_vendor")() - - -def llvm_get_system_triple(): - """Get system host triple. - - Parameters - ---------- - - Returns - ------- - triple : str - The current system's triple. - """ - return get_global_func("codegen.llvm_get_system_triple")() - - -def llvm_get_system_cpu(): - """Get system host cpu name. - - Parameters - ---------- - - Returns - ------- - cpu_name : str - The current system's cpu name. - """ - return get_global_func("codegen.llvm_get_system_cpu")() - - -def llvm_get_targets(): - """Get LLVM target list. - - Parameters - ---------- - - Returns - ------- - llvm_targets : list[str] - List of available LLVM targets. - """ - return get_global_func("codegen.llvm_get_targets")() - - -def llvm_get_cpu_archlist(target=None): - """Get CPU architectures for the target's `-mtriple`. - - Parameters - ---------- - target : Target - The TVM target. - - Returns - ------- - cpu_archlist : list[str] - List of available CPU architectures. - """ - assert isinstance(target, Target) or target is None - return get_global_func("codegen.llvm_get_cpu_archlist")(target) - - -def llvm_get_cpu_features(target=None): - """Get CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`. - - Parameters - ---------- - target : Target - The TVM target. - - Returns - ------- - cpu_features : list[str] - List of available CPU features. - """ - assert isinstance(target, Target) or target is None - feature_map = get_global_func("codegen.llvm_get_cpu_features")(target) - return set(feature_map.keys()) - - -def llvm_cpu_has_features(cpu_features, target=None): - """Check CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`. - - Parameters - ---------- - target : Target - The TVM target. - cpu_features : str or Array - CPU Feature(s) to check. - - Returns - ------- - has_features : bool - True if target CPU has the feature(s). - """ - assert isinstance(target, Target) or target is None - assert isinstance(cpu_features, Array | list | tuple | str) - has_feats = True - cpu_features = [cpu_features] if isinstance(cpu_features, str) else cpu_features - for feat in cpu_features: - has_feats &= get_global_func("codegen.llvm_cpu_has_feature")(feat, target) - return has_feats - - -def llvm_get_vector_width(target=None): - """Get vector width from LLVM target's `-mtriple` and `-mcpu` and considering `-mattr`. - - Parameters - ---------- - target : Target - The TVM target. - - Returns - ------- - vector_width : int - Vector with of target in number of bits, -1 on error. - """ - assert isinstance(target, Target) or target is None - return get_global_func("codegen.llvm_get_vector_width")(target) - - -def llvm_is_valid_cpu(cpu, triple): - """Check if a CPU name is valid for the given LLVM triple. - - Parameters - ---------- - cpu : str - The CPU name to check (e.g. "apple-m1"). - triple : str - The LLVM target triple (e.g. "arm64-apple-macos"). - - Returns - ------- - is_valid : bool - True if the CPU name is recognized by LLVM for the given triple. - """ - return get_global_func("codegen.llvm_is_valid_cpu")(cpu, triple) - - -def llvm_version_major(allow_none=False): - """Get the major LLVM version. - - Parameters - ---------- - allow_none : bool - Whether do we allow none. - - Returns - ------- - major : int - The major LLVM version. - """ - try: - return get_global_func("codegen.llvm_version_major")() - except (ValueError, AttributeError): - if allow_none: - return None - raise RuntimeError("LLVM version is not available, please check if you built TVM with LLVM") diff --git a/tirx/python/tvm_tirx/codegen/__init__.py b/tirx/python/tvm_tirx/codegen/llvm.py similarity index 75% copy from tirx/python/tvm_tirx/codegen/__init__.py copy to tirx/python/tvm_tirx/codegen/llvm.py index afc5fc7c3f..8734fea963 100644 --- a/tirx/python/tvm_tirx/codegen/__init__.py +++ b/tirx/python/tvm_tirx/codegen/llvm.py @@ -14,60 +14,14 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Code generation related functions.""" +"""LLVM code generation and target information.""" from tvm_ffi import Array, get_global_func from tvm_tirx.target import Target -from .executable import Executable, load_static_library - -def build_module(mod, target): - """Build IRModule into Module. - - Parameters - ---------- - mod : tvm_tirx.IRModule - The ir module. - - target : str - The target module type. - - Returns - ------- - module : runtime.Module - The corresponding module. - """ - target = Target(target) if isinstance(target, str) else target - return get_global_func("codegen.Build")(mod, target) - - -def target_has_features(cpu_features, target=None): - """Check CPU features for the target's `-mtriple` and `-mcpu` and `-mattr`. - - Parameters - ---------- - target : Target - The TVM target. - cpu_features : str or Array - CPU Feature(s) to check. - - Returns - ------- - has_features : bool - True if target has the feature(s). - """ - assert isinstance(target, Target) or target is None - assert isinstance(cpu_features, Array | list | tuple | str) - has_feats = True - cpu_features = [cpu_features] if isinstance(cpu_features, str) else cpu_features - for feat in cpu_features: - has_feats &= get_global_func("codegen.target_has_feature")(feat, target) - return has_feats - - -def llvm_lookup_intrinsic_id(name): +def lookup_intrinsic_id(name): """Lookup LLVM intrinsic id by name. Parameters @@ -83,7 +37,7 @@ def llvm_lookup_intrinsic_id(name): return get_global_func("codegen.llvm_lookup_intrinsic_id")(name) -def llvm_get_intrinsic_name(intrin_id: int) -> str: +def get_intrinsic_name(intrin_id: int) -> str: """Get the name of an intrinsic for a given id. Parameters @@ -99,7 +53,7 @@ def llvm_get_intrinsic_name(intrin_id: int) -> str: return get_global_func("codegen.llvm_get_intrinsic_name")(intrin_id) -def llvm_get_system_x86_vendor(): +def get_system_x86_vendor(): """Get system x86 vendor info. Parameters @@ -113,7 +67,7 @@ def llvm_get_system_x86_vendor(): return get_global_func("codegen.llvm_get_system_x86_vendor")() -def llvm_get_system_triple(): +def get_system_triple(): """Get system host triple. Parameters @@ -127,7 +81,7 @@ def llvm_get_system_triple(): return get_global_func("codegen.llvm_get_system_triple")() -def llvm_get_system_cpu(): +def get_system_cpu(): """Get system host cpu name. Parameters @@ -141,7 +95,7 @@ def llvm_get_system_cpu(): return get_global_func("codegen.llvm_get_system_cpu")() -def llvm_get_targets(): +def get_targets(): """Get LLVM target list. Parameters @@ -155,7 +109,7 @@ def llvm_get_targets(): return get_global_func("codegen.llvm_get_targets")() -def llvm_get_cpu_archlist(target=None): +def get_cpu_archlist(target=None): """Get CPU architectures for the target's `-mtriple`. Parameters @@ -172,7 +126,7 @@ def llvm_get_cpu_archlist(target=None): return get_global_func("codegen.llvm_get_cpu_archlist")(target) -def llvm_get_cpu_features(target=None): +def get_cpu_features(target=None): """Get CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`. Parameters @@ -190,7 +144,7 @@ def llvm_get_cpu_features(target=None): return set(feature_map.keys()) -def llvm_cpu_has_features(cpu_features, target=None): +def cpu_has_features(cpu_features, target=None): """Check CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`. Parameters @@ -214,7 +168,7 @@ def llvm_cpu_has_features(cpu_features, target=None): return has_feats -def llvm_get_vector_width(target=None): +def get_vector_width(target=None): """Get vector width from LLVM target's `-mtriple` and `-mcpu` and considering `-mattr`. Parameters @@ -231,7 +185,7 @@ def llvm_get_vector_width(target=None): return get_global_func("codegen.llvm_get_vector_width")(target) -def llvm_is_valid_cpu(cpu, triple): +def is_valid_cpu(cpu, triple): """Check if a CPU name is valid for the given LLVM triple. Parameters @@ -249,7 +203,7 @@ def llvm_is_valid_cpu(cpu, triple): return get_global_func("codegen.llvm_is_valid_cpu")(cpu, triple) -def llvm_version_major(allow_none=False): +def version_major(allow_none=False): """Get the major LLVM version. Parameters diff --git a/tirx/python/tvm_tirx/config.py b/tirx/python/tvm_tirx/config.py index 466a30fe14..da829bb6dc 100644 --- a/tirx/python/tvm_tirx/config.py +++ b/tirx/python/tvm_tirx/config.py @@ -22,88 +22,21 @@ there with an exact library directory, independently of Python package search. """ import argparse -import importlib.machinery import importlib.util import json -import os import shlex import sys from pathlib import Path -__version__ = "0.1.0a0" - -_LIBRARY_NAMES = ("tvm_tirx_compiler", "tvm_runtime") -_COMPONENTS = { - "compiler": "tvm_tirx_compiler", - "runtime": "tvm_runtime", - "cuda": "tvm_runtime_cuda", - "metal": "tvm_runtime_metal", - "opencl": "tvm_runtime_opencl", - "vulkan": "tvm_runtime_vulkan", - "rocm": "tvm_runtime_rocm", -} - - -def _find_prefix(start: Path) -> Path: - for candidate in (start, *start.parents): - if (candidate / "include/tvm/runtime/base.h").is_file(): - return candidate - raise RuntimeError(f"Cannot find TIRx public headers above {start}") - - -def _paths(prefix: Path | None, library_dir: Path | None) -> tuple[Path, Path]: - explicit_prefix = prefix is not None - if library_dir is not None: - library_dir = library_dir.resolve() - prefix = prefix.resolve() if prefix is not None else _find_prefix(library_dir) - else: - prefix = prefix.resolve() if prefix is not None else _find_prefix(Path(__file__).parent) - library_dir = prefix / "lib" - candidates = [prefix / "lib", prefix / "lib64", prefix / "build/lib"] - if prefix.name == "tirx": - candidates.append(prefix.parent / "build/lib") - if not explicit_prefix and os.environ.get("TVM_TIRX_LIBRARY_PATH"): - candidates.insert(0, Path(os.environ["TVM_TIRX_LIBRARY_PATH"])) - for candidate in candidates: - if all( - any( - (candidate / filename).is_file() - for filename in (f"lib{name}.so", f"lib{name}.dylib", f"{name}.dll") - ) - for name in ("tvm_runtime",) - ): - library_dir = candidate - break - if not (prefix / "include/tvm/runtime/base.h").is_file(): - raise RuntimeError(f"Cannot find TIRx public headers in {prefix / 'include'}") - return prefix, library_dir - - -def _libraries(library_dir: Path, platform: str, names=_LIBRARY_NAMES) -> dict[str, dict[str, str]]: - libraries = {} - windows = platform in ("Windows", "win32") - apple = platform in ("Darwin", "iOS", "darwin") - for name in names: - if windows: - filenames = (f"{name}.dll",) - elif apple: - filenames = (f"lib{name}.dylib", f"lib{name}.so") - else: - filenames = (f"lib{name}.so",) - location = next((library_dir / f for f in filenames if (library_dir / f).is_file()), None) - if location is None: - raise RuntimeError(f"Cannot find TIRx library {name} in {library_dir}") - implib = "" - if windows: - import_names = (f"{name}.lib", f"lib{name}.dll.a") - import_library = next( - (library_dir / f for f in import_names if (library_dir / f).is_file()), None - ) - if import_library is None: - raise RuntimeError(f"Missing import library for {location}") - implib = str(import_library) - libraries[name] = {"location": str(location), "implib": implib} - return libraries +if __package__: + from . import libinfo +else: + # The native CMake package installs the pure Python helper beside this file. + spec = importlib.util.spec_from_file_location( + "_tvm_tirx_libinfo", Path(__file__).with_name("libinfo.py") + ) + libinfo = importlib.util.module_from_spec(spec) + spec.loader.exec_module(libinfo) def configuration( @@ -112,62 +45,7 @@ def configuration( platform: str = sys.platform, ) -> dict: """Return the public include and core shared-library paths for one installation.""" - prefix, library_dir = _paths(prefix, library_dir) - return { - "version": __version__, - "include_dir": str(prefix / "include"), - "library_dir": str(library_dir), - "cmake_dir": str(library_dir / "cmake/tvm_tirx"), - "libraries": _libraries(library_dir, platform), - } - - -def _ffi_info(): - """Read the selected FFI package's pure Python helper without importing its runtime.""" - spec = importlib.machinery.PathFinder.find_spec("tvm_ffi") - if spec is None or spec.origin is None: - raise RuntimeError("Cannot find the matching apache-tvm-ffi Python package") - helper = Path(spec.origin).parent / "libinfo.py" - spec = importlib.util.spec_from_file_location("_tvm_ffi_libinfo", helper) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -def _available(library_dir, platform): - available = [] - for component, name in _COMPONENTS.items(): - try: - _libraries(library_dir, platform, (name,)) - except RuntimeError: - continue - available.append(component) - return available - - -def _link_inputs(library_dir, platform, components): - available = _available(library_dir, platform) - selected = list(dict.fromkeys(components or ["all"])) - if "all" in selected: - selected = [*available, *(name for name in selected if name != "all")] - for name in selected: - if name != "ffi" and name not in available: - raise RuntimeError( - f"Component '{name}' is unavailable; available: {' '.join(available)} ffi" - ) - # Dependents precede their dependencies even if the command line lists runtime first. - ordered = [name for name in selected if name not in ("runtime", "ffi")] - if any(name != "ffi" for name in selected): - ordered.append("runtime") - libraries = _libraries(library_dir, platform, [_COMPONENTS[name] for name in ordered]) - ffi = _ffi_info() - ffi_library = ( - ffi.find_windows_implib() if platform in ("Windows", "win32") else ffi.find_libtvm_ffi() - ) - return [ - *(Path(lib["implib"] or lib["location"]) for lib in libraries.values()), - Path(ffi_library), - ] + return libinfo.configuration(prefix, library_dir, platform) def __main__() -> None: @@ -216,10 +94,10 @@ def __main__() -> None: "complete static dependency closure (including TVM FFI and its dependencies)" ) if args.show_version: - print(__version__) + print(libinfo.__version__) if all(not getattr(args, "show_" + output) for output in outputs if output != "version"): return - prefix, library_dir = _paths(args.install_prefix, args.library_dir) + prefix, library_dir = libinfo._paths(args.install_prefix, args.library_dir) paths = { "prefix": prefix, "includedir": prefix / "include", @@ -232,11 +110,11 @@ def __main__() -> None: if getattr(args, "show_" + output): print(path) if args.show_components: - print(" ".join([*_available(library_dir, args.platform), "ffi"])) + print(" ".join([*libinfo._available(library_dir, args.platform), "ffi"])) if args.show_json: print(json.dumps(configuration(prefix, library_dir, args.platform))) if args.show_cflags or args.show_cxxflags: - ffi = _ffi_info() + ffi = libinfo._ffi_info() includes = dict.fromkeys( [str(prefix / "include"), ffi.find_include_path(), ffi.find_dlpack_include_path()] ) @@ -246,7 +124,7 @@ def __main__() -> None: if args.show_cxxflags: print(shlex.join([*flags, "-std=c++20"])) if any((args.show_ldflags, args.show_libs, args.show_libnames, args.show_libfiles)): - libraries = _link_inputs(library_dir, args.platform, args.components) + libraries = libinfo._link_inputs(library_dir, args.platform, args.components) apple = args.platform in ("Darwin", "iOS", "darwin") if args.show_ldflags: print( diff --git a/tirx/python/tvm_tirx/error.py b/tirx/python/tvm_tirx/error.py deleted file mode 100644 index e1a5fdd9b2..0000000000 --- a/tirx/python/tvm_tirx/error.py +++ /dev/null @@ -1,126 +0,0 @@ -# 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. -"""Structured error classes in TVM. - -Each error class takes an error message as its input. -See the example sections for suggested message conventions. -To make the code more readable, we recommended developers to -copy the examples and raise errors with the same message convention. - -.. note:: - - Please also refer to :ref:`error-handling-guide`. -""" - -from tvm_ffi import register_error - - -@register_error -class InternalError(RuntimeError): - """Internal error in the system. - - Examples - -------- - .. code :: c++ - - // Example code C++ - LOG(FATAL) << "InternalError: internal error detail."; - - .. code :: python - - # Example code in python - raise InternalError("internal error detail") - """ - - -@register_error -class RPCError(RuntimeError): - """Error thrown by the remote server handling the RPC call.""" - - -@register_error -class RPCSessionTimeoutError(RPCError, TimeoutError): - """Error thrown by the remote server when the RPC session has expired.""" - - -@register_error -class OpError(RuntimeError): - """Base class of all operator errors in frontends.""" - - -@register_error -class OpNotImplemented(OpError, NotImplementedError): - """Operator is not implemented. - - Example - ------- - .. code:: python - - raise OpNotImplemented( - "Operator {} is not supported in {} frontend".format( - missing_op, frontend_name)) - """ - - -@register_error -class OpAttributeRequired(OpError, AttributeError): - """Required attribute is not found. - - Example - ------- - .. code:: python - - raise OpAttributeRequired( - "Required attribute {} not found in operator {}".format( - attr_name, op_name)) - """ - - -@register_error -class OpAttributeInvalid(OpError, AttributeError): - """Attribute value is invalid when taking in a frontend operator. - - Example - ------- - .. code:: python - - raise OpAttributeInvalid( - "Value {} in attribute {} of operator {} is not valid".format( - value, attr_name, op_name)) - """ - - -@register_error -class OpAttributeUnImplemented(OpError, NotImplementedError): - """Attribute is not supported in a certain frontend. - - Example - ------- - .. code:: python - - raise OpAttributeUnImplemented( - "Attribute {} is not supported in operator {}".format( - attr_name, op_name)) - """ - - -@register_error -class DiagnosticError(RuntimeError): - """Error diagnostics were reported during the execution of a pass. - - See the configured diagnostic renderer for detailed error information. - """ diff --git a/tirx/python/tvm_tirx/ir/op.py b/tirx/python/tvm_tirx/ir/op.py index 7773e33129..e9c5ac10ca 100644 --- a/tirx/python/tvm_tirx/ir/op.py +++ b/tirx/python/tvm_tirx/ir/op.py @@ -30,121 +30,6 @@ from . import _ffi_api from .expr import Expr -def _make_op_api(op, module_name): - """Build a callable whose operands and result are governed by its Op.""" - from .attrs import make_node # pylint: disable=import-outside-toplevel - from .expr import Call # pylint: disable=import-outside-toplevel - - def call( - *args, - attrs=None, - ty_args=None, - loc: Location = UNKNOWN_LOC, - ty=None, - **kwargs, - ): - # Bind named operands using the registered signature, without inventing - # defaults or interpreting arbitrary Python wrapper signatures. - is_region = op.get_attr("FRegionGetBodyParams") is not None - body_params = kwargs.pop("body_params", None) if is_region else None - operands = list(args) - for index, info in enumerate(op.args_info): - if index < len(args): - if info.name in kwargs: - raise TypeError(f"{op.name}: multiple values for {info.name!r}") - elif info.name in kwargs: - operands.append(kwargs.pop(info.name)) - else: - raise TypeError(f"{op.name}: missing operand {info.name!r}") - if kwargs or (attrs is None and op.attrs_type_key): - if not op.attrs_type_key: - raise TypeError(f"{op.name}: unexpected keyword operands {tuple(kwargs)}") - if attrs is not None: - raise TypeError(f"{op.name}: cannot mix attrs with attribute keywords") - attrs = make_node(op.attrs_type_key, **kwargs) - if is_region: - from tvm_tirx.script.ir_builder.stmt import region - - if ( - ty is not None - or ty_args is not None - or not isinstance(loc, Location) - or not loc.same_as(UNKNOWN_LOC) - ): - raise TypeError( - f"{op.name}: region builders do not accept ty, ty_args or an explicit loc" - ) - return region(op, operands, body_params=body_params, attrs=attrs) - return Call(op, operands, attrs=attrs, ty_args=ty_args, loc=loc, ty=ty) - - call.__name__ = op.name.rsplit(".", 1)[-1] - call.__module__ = module_name - call.__doc__ = op.doc or f"Construct a call to {op.name}." - return call - - -def _init_op_api(namespace, target_module_name=None, *, wrapper=None): - """Initialize registered Op callables in an already loaded Python module. - - Like :func:`tvm_ffi.init_ffi_api`, the registry prefix comes first and the - target defaults to that module name. For example, a backend module uses - ``_init_op_api("tirx.cuda", __name__)``. Dotted suffixes require existing - module or SimpleNamespace containers; underscores are ordinary name parts. - - Generated functions accept registered positional/named operands plus - ``attrs``, ``ty_args``, ``loc`` and ``ty``. - Region operations construct a RegionFrame and accept ``body_params`` instead - of Call result/type arguments. - Omitting the result invokes an available Op inference hook; without one, - Call retains a missing type. - Explicit results and inference errors are preserved; Call.validate checks - the operator contract separately. Attribute - keywords construct the registered attrs schema, when one is declared. - - Existing Python callables retain ownership of their names; only missing - exposures are generated. Existing functions must accept printed calls or - have an appropriate exceptional printer hook. Non-callable collisions fail - before any functions are installed. An optional ``wrapper`` adapts newly - generated Call constructors at export time; existing exposures and region - builders are untouched. - Returns None. - """ - if not namespace or any( - not part.isidentifier() or keyword.iskeyword(part) for part in namespace.split(".") - ): - raise ValueError(f"Invalid Op namespace {namespace!r}") - target = sys.modules[target_module_name or namespace] - pending = [] - destinations = set() - for name in sorted(Op.list_op_names()): - if not name.startswith(namespace + "."): - continue - parts = name[len(namespace) + 1 :].split(".") - if any(not part.isidentifier() or keyword.iskeyword(part) for part in parts): - raise ValueError(f"Op {name!r} has no Python attribute spelling") - container = target - for part in parts[:-1]: - container = vars(container).get(part) - if not isinstance(container, ModuleType | SimpleNamespace): - raise ValueError(f"Op {name!r} requires an existing namespace at {part!r}") - destination = (id(container), parts[-1]) - if destination in destinations: - raise ValueError(f"Op {name!r} aliases another exposure destination") - destinations.add(destination) - op = Op.get(name) - if parts[-1] in vars(container): - current = vars(container)[parts[-1]] - if not callable(current): - raise ValueError(f"Op {name!r} conflicts with an existing Python attribute") - else: - call = _make_op_api(op, target.__name__) - if wrapper is not None and op.get_attr("FRegionGetBodyParams") is None: - call = wrapper(call) - pending.append((container, parts[-1], call)) - for container, name, call in pending: - setattr(container, name, call) - - @tvm_ffi.register_object("ir.Op") class Op(Expr): """Primitive operator in the IR.""" @@ -280,3 +165,118 @@ def register_op_attr(op_name, attr_key, value=None, override=False): return v return _register(value) if value is not None else _register + + +def _make_op_api(op, module_name): + """Build a callable whose operands and result are governed by its Op.""" + from .attrs import make_node # pylint: disable=import-outside-toplevel + from .expr import Call # pylint: disable=import-outside-toplevel + + def call( + *args, + attrs=None, + ty_args=None, + loc: Location = UNKNOWN_LOC, + ty=None, + **kwargs, + ): + # Bind named operands using the registered signature, without inventing + # defaults or interpreting arbitrary Python wrapper signatures. + is_region = op.get_attr("FRegionGetBodyParams") is not None + body_params = kwargs.pop("body_params", None) if is_region else None + operands = list(args) + for index, info in enumerate(op.args_info): + if index < len(args): + if info.name in kwargs: + raise TypeError(f"{op.name}: multiple values for {info.name!r}") + elif info.name in kwargs: + operands.append(kwargs.pop(info.name)) + else: + raise TypeError(f"{op.name}: missing operand {info.name!r}") + if kwargs or (attrs is None and op.attrs_type_key): + if not op.attrs_type_key: + raise TypeError(f"{op.name}: unexpected keyword operands {tuple(kwargs)}") + if attrs is not None: + raise TypeError(f"{op.name}: cannot mix attrs with attribute keywords") + attrs = make_node(op.attrs_type_key, **kwargs) + if is_region: + from tvm_tirx.script.ir_builder.stmt import region + + if ( + ty is not None + or ty_args is not None + or not isinstance(loc, Location) + or not loc.same_as(UNKNOWN_LOC) + ): + raise TypeError( + f"{op.name}: region builders do not accept ty, ty_args or an explicit loc" + ) + return region(op, operands, body_params=body_params, attrs=attrs) + return Call(op, operands, attrs=attrs, ty_args=ty_args, loc=loc, ty=ty) + + call.__name__ = op.name.rsplit(".", 1)[-1] + call.__module__ = module_name + call.__doc__ = op.doc or f"Construct a call to {op.name}." + return call + + +def _init_op_api(namespace, target_module_name=None, *, wrapper=None): + """Initialize registered Op callables in an already loaded Python module. + + Like :func:`tvm_ffi.init_ffi_api`, the registry prefix comes first and the + target defaults to that module name. For example, a backend module uses + ``_init_op_api("tirx.cuda", __name__)``. Dotted suffixes require existing + module or SimpleNamespace containers; underscores are ordinary name parts. + + Generated functions accept registered positional/named operands plus + ``attrs``, ``ty_args``, ``loc`` and ``ty``. + Region operations construct a RegionFrame and accept ``body_params`` instead + of Call result/type arguments. + Omitting the result invokes an available Op inference hook; without one, + Call retains a missing type. + Explicit results and inference errors are preserved; Call.validate checks + the operator contract separately. Attribute + keywords construct the registered attrs schema, when one is declared. + + Existing Python callables retain ownership of their names; only missing + exposures are generated. Existing functions must accept printed calls or + have an appropriate exceptional printer hook. Non-callable collisions fail + before any functions are installed. An optional ``wrapper`` adapts newly + generated Call constructors at export time; existing exposures and region + builders are untouched. + Returns None. + """ + if not namespace or any( + not part.isidentifier() or keyword.iskeyword(part) for part in namespace.split(".") + ): + raise ValueError(f"Invalid Op namespace {namespace!r}") + target = sys.modules[target_module_name or namespace] + pending = [] + destinations = set() + for name in sorted(Op.list_op_names()): + if not name.startswith(namespace + "."): + continue + parts = name[len(namespace) + 1 :].split(".") + if any(not part.isidentifier() or keyword.iskeyword(part) for part in parts): + raise ValueError(f"Op {name!r} has no Python attribute spelling") + container = target + for part in parts[:-1]: + container = vars(container).get(part) + if not isinstance(container, ModuleType | SimpleNamespace): + raise ValueError(f"Op {name!r} requires an existing namespace at {part!r}") + destination = (id(container), parts[-1]) + if destination in destinations: + raise ValueError(f"Op {name!r} aliases another exposure destination") + destinations.add(destination) + op = Op.get(name) + if parts[-1] in vars(container): + current = vars(container)[parts[-1]] + if not callable(current): + raise ValueError(f"Op {name!r} conflicts with an existing Python attribute") + else: + call = _make_op_api(op, target.__name__) + if wrapper is not None and op.get_attr("FRegionGetBodyParams") is None: + call = wrapper(call) + pending.append((container, parts[-1], call)) + for container, name, call in pending: + setattr(container, name, call) diff --git a/tirx/python/tvm_tirx/libinfo.py b/tirx/python/tvm_tirx/libinfo.py index 79345f706b..cbe13ae515 100644 --- a/tirx/python/tvm_tirx/libinfo.py +++ b/tirx/python/tvm_tirx/libinfo.py @@ -18,11 +18,12 @@ from __future__ import annotations +import importlib.machinery +import importlib.util 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_tirx/``.""" @@ -72,17 +73,17 @@ def find_libtvm_runtime() -> str: RuntimeError If the runtime library cannot be found. """ - candidate = tvm_ffi_libinfo._find_library_by_basename( + candidate = _ffi_info()._find_library_by_basename( "tvm_tirx", "tvm_runtime", extra_lib_paths=package_lib_paths() ) - if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda _: True): + if ret := _ffi_info()._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( + if ret := _ffi_info()._resolve_and_validate( paths=[ _rel_top_directory() / "include", _dev_top_directory() / "include", @@ -94,3 +95,141 @@ def find_include_path() -> str: __version__ = "0.1.0a0" + + +_LIBRARY_NAMES = ("tvm_tirx_compiler", "tvm_runtime") +_COMPONENTS = { + "compiler": "tvm_tirx_compiler", + "runtime": "tvm_runtime", + "cuda": "tvm_runtime_cuda", + "metal": "tvm_runtime_metal", + "opencl": "tvm_runtime_opencl", + "vulkan": "tvm_runtime_vulkan", + "rocm": "tvm_runtime_rocm", +} + + +def _find_prefix(start: Path) -> Path: + for candidate in (start, *start.parents): + if (candidate / "include/tvm/runtime/base.h").is_file(): + return candidate + raise RuntimeError(f"Cannot find TIRx public headers above {start}") + + +def _paths(prefix: Path | None, library_dir: Path | None) -> tuple[Path, Path]: + explicit_prefix = prefix is not None + if library_dir is not None: + library_dir = library_dir.resolve() + prefix = prefix.resolve() if prefix is not None else _find_prefix(library_dir) + else: + prefix = prefix.resolve() if prefix is not None else _find_prefix(Path(__file__).parent) + library_dir = prefix / "lib" + candidates = [prefix / "lib", prefix / "lib64", prefix / "build/lib"] + if prefix.name == "tirx": + candidates.append(prefix.parent / "build/lib") + if not explicit_prefix and os.environ.get("TVM_TIRX_LIBRARY_PATH"): + candidates.insert(0, Path(os.environ["TVM_TIRX_LIBRARY_PATH"])) + for candidate in candidates: + if all( + any( + (candidate / filename).is_file() + for filename in (f"lib{name}.so", f"lib{name}.dylib", f"{name}.dll") + ) + for name in ("tvm_runtime",) + ): + library_dir = candidate + break + if not (prefix / "include/tvm/runtime/base.h").is_file(): + raise RuntimeError(f"Cannot find TIRx public headers in {prefix / 'include'}") + return prefix, library_dir + + +def _libraries(library_dir: Path, platform: str, names=_LIBRARY_NAMES) -> dict[str, dict[str, str]]: + libraries = {} + windows = platform in ("Windows", "win32") + apple = platform in ("Darwin", "iOS", "darwin") + for name in names: + if windows: + filenames = (f"{name}.dll",) + elif apple: + filenames = (f"lib{name}.dylib", f"lib{name}.so") + else: + filenames = (f"lib{name}.so",) + location = next((library_dir / f for f in filenames if (library_dir / f).is_file()), None) + if location is None: + raise RuntimeError(f"Cannot find TIRx library {name} in {library_dir}") + implib = "" + if windows: + import_names = (f"{name}.lib", f"lib{name}.dll.a") + import_library = next( + (library_dir / f for f in import_names if (library_dir / f).is_file()), None + ) + if import_library is None: + raise RuntimeError(f"Missing import library for {location}") + implib = str(import_library) + libraries[name] = {"location": str(location), "implib": implib} + return libraries + + +def configuration( + prefix: Path | None = None, + library_dir: Path | None = None, + platform: str = sys.platform, +) -> dict: + """Return the public include and core shared-library paths for one installation.""" + prefix, library_dir = _paths(prefix, library_dir) + return { + "version": __version__, + "include_dir": str(prefix / "include"), + "library_dir": str(library_dir), + "cmake_dir": str(library_dir / "cmake/tvm_tirx"), + "libraries": _libraries(library_dir, platform), + } + + +def _ffi_info(): + """Read the selected FFI package's pure Python helper without importing its runtime.""" + spec = importlib.machinery.PathFinder.find_spec("tvm_ffi") + if spec is None or spec.origin is None: + raise RuntimeError("Cannot find the matching apache-tvm-ffi Python package") + helper = Path(spec.origin).parent / "libinfo.py" + spec = importlib.util.spec_from_file_location("_tvm_ffi_libinfo", helper) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _available(library_dir, platform): + available = [] + for component, name in _COMPONENTS.items(): + try: + _libraries(library_dir, platform, (name,)) + except RuntimeError: + continue + available.append(component) + return available + + +def _link_inputs(library_dir, platform, components): + available = _available(library_dir, platform) + selected = list(dict.fromkeys(components or ["all"])) + if "all" in selected: + selected = [*available, *(name for name in selected if name != "all")] + for name in selected: + if name != "ffi" and name not in available: + raise RuntimeError( + f"Component '{name}' is unavailable; available: {' '.join(available)} ffi" + ) + # Dependents precede their dependencies even if the command line lists runtime first. + ordered = [name for name in selected if name not in ("runtime", "ffi")] + if any(name != "ffi" for name in selected): + ordered.append("runtime") + libraries = _libraries(library_dir, platform, [_COMPONENTS[name] for name in ordered]) + ffi = _ffi_info() + ffi_library = ( + ffi.find_windows_implib() if platform in ("Windows", "win32") else ffi.find_libtvm_ffi() + ) + return [ + *(Path(lib["implib"] or lib["location"]) for lib in libraries.values()), + Path(ffi_library), + ] diff --git a/tirx/python/tvm_tirx/support/clang.py b/tirx/python/tvm_tirx/support/clang.py index c0763dfcd0..09120be9e5 100644 --- a/tirx/python/tvm_tirx/support/clang.py +++ b/tirx/python/tvm_tirx/support/clang.py @@ -44,7 +44,7 @@ def find_clang(required=True): matches the major llvm version that built with tvm """ cc_list = [] - major = tvm_tirx.codegen.llvm_version_major(allow_none=True) + major = tvm_tirx.codegen.llvm.version_major(allow_none=True) if major is not None: cc_list += [f"clang-{major}.0"] cc_list += [f"clang-{major}"] diff --git a/tirx/python/tvm_tirx/support/rocm.py b/tirx/python/tvm_tirx/support/rocm.py index 8da215a0fe..962b999787 100644 --- a/tirx/python/tvm_tirx/support/rocm.py +++ b/tirx/python/tvm_tirx/support/rocm.py @@ -49,7 +49,7 @@ def find_lld(required=True): matches the major llvm version that built with tvm """ lld_list = [] - major = tvm_tirx.codegen.llvm_version_major(allow_none=True) + major = tvm_tirx.codegen.llvm.version_major(allow_none=True) if major is not None: lld_list += [f"ld.lld-{major}.0"] lld_list += [f"ld.lld-{major}"] diff --git a/tirx/python/tvm_tirx/testing/env.py b/tirx/python/tvm_tirx/testing/env.py index 7e58d5c269..6ce5d3cd34 100644 --- a/tirx/python/tvm_tirx/testing/env.py +++ b/tirx/python/tvm_tirx/testing/env.py @@ -404,7 +404,7 @@ def has_nvcc_version(major: int, minor: int = 0, release: int = 0) -> bool: def _llvm_version_major() -> int: """Return the major LLVM version, or 0 if LLVM is unavailable.""" try: - return int(tvm_tirx.codegen.llvm_version_major()) + return int(tvm_tirx.codegen.llvm.version_major()) except Exception: # pylint: disable=broad-except return 0 @@ -451,7 +451,7 @@ def has_hexagon_toolchain() -> bool: try: return ( _runtime_enabled("llvm") - and "hexagon" in tvm_tirx.codegen.llvm_get_targets() + and "hexagon" in tvm_tirx.codegen.llvm.get_targets() and "HEXAGON_TOOLCHAIN" in os.environ ) except Exception: # pylint: disable=broad-except @@ -485,8 +485,8 @@ def _has_cpu_feature(features) -> bool: """True if the host CPU advertises the given LLVM target ``features``.""" try: codegen = tvm_tirx.codegen - cpu = codegen.llvm_get_system_cpu() - triple = codegen.llvm_get_system_triple() + cpu = codegen.llvm.get_system_cpu() + triple = codegen.llvm.get_system_triple() target = tvm_tirx.target.Target({"kind": "llvm", "mtriple": triple, "mcpu": cpu}) return bool(codegen.target_has_features(features, target)) except Exception: # pylint: disable=broad-except diff --git a/tirx/python/tvm_tirx/tirx/__init__.py b/tirx/python/tvm_tirx/tirx/__init__.py index c6653bfa1f..cbbde24567 100644 --- a/tirx/python/tvm_tirx/tirx/__init__.py +++ b/tirx/python/tvm_tirx/tirx/__init__.py @@ -44,7 +44,8 @@ from .expr import Select, TensorLoad from .expr import Let -from .function import Function, IndexMap, renew_def +from .function import Function, renew_def +from .index_map import IndexMap from .op import call_packed_lowered, call_cpacked_lowered, register_intrin_lowering from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern @@ -84,7 +85,6 @@ from .thread_scope import ThreadScope from .layout import TileLayout, Layout, ComposeLayout from . import transform from . import analysis -from . import backend # Compiler-only submodules. Skip under `TVM_USE_RUNTIME_LIB=1` since they # perform compiler-side FFI at module load (schema engine looks up diff --git a/tirx/python/tvm_tirx/tirx/function.py b/tirx/python/tvm_tirx/tirx/function.py index 2a7ce72203..daa9b20d30 100644 --- a/tirx/python/tvm_tirx/tirx/function.py +++ b/tirx/python/tvm_tirx/tirx/function.py @@ -17,20 +17,16 @@ # pylint: disable=unrecognized-inline-option """Function data types.""" -import collections -import inspect -from collections.abc import Callable, Mapping, Sequence -from typing import Optional +from collections.abc import Mapping, Sequence import tvm_ffi import tvm_tirx import tvm_tirx.runtime -from tvm_tirx.ir import BaseFunc, Range, Scriptable, SeqStmt, Stmt, Var +from tvm_tirx.ir import BaseFunc, Scriptable, SeqStmt, Stmt, Var from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.runtime import Object -from ..runtime._tensor import Tensor from . import _ffi_api from .expr import Expr @@ -88,11 +84,6 @@ class Function(BaseFunc, Scriptable): loc, ) # type: ignore - @property - def is_tirx(self): - """Whether this primitive function uses the TIRx dialect.""" - return not bool(self.attrs.get("s_tir", False)) - def with_body(self, new_body: Stmt | Sequence[Stmt] | None, loc: Location = UNKNOWN_LOC): """Create a new Function with the same set signatures but a new body. @@ -191,287 +182,3 @@ def renew_def(func: Function) -> Function: The function with fresh variable definitions. """ return _ffi_api.RenewDef(func) - - -@tvm_ffi.register_object("tirx.IndexMap") -class IndexMap(Object): - """A mapping from multi-dimensional indices to another set of multi-dimensional indices - - Parameters - ---------- - initial_indices : List[Var] - Variables representing the indices prior to remapping. - final_indices : List[Expr] - Expressions defining the indices after remapping. - inverse_index_map : Union[Callable, Optional[IndexMap]] - The optional pre-defined inverse index map. - When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. - Otherwise, the inverse index map will be computed on the fly. - It is the user's responsibility to ensure the correctness of the pre-defined inverse - index map. - """ - - initial_indices: list[Var] - final_indices: list[Expr] - - def __init__(self, initial_indices, final_indices, inverse_index_map): - if isinstance(inverse_index_map, Callable): - inverse_index_map = IndexMap.from_func(inverse_index_map) - self.__init_handle_by_constructor__( - _ffi_api.IndexMap, initial_indices, final_indices, inverse_index_map - ) - - @staticmethod - def from_func( - mapping_function: Callable, - ndim: int | None = None, - inverse_index_map: Callable | Optional["IndexMap"] = None, - *, - index_dtype: str = "int64", - ): - """Create an index map from a function - - Parameters - ---------- - mapping_function : Callable - - The function to map from source indices to target indices. - The function should accept `tirx.Var` parameters and return - a either a `tirx.Expr`, or a list of `tirx.Expr`. - Returning a `tirx.Expr` is equivalent to returning a - list of length 1 containing that `tirx.Expr`. - - ndim: Optional[int] - - The dimensionality of the tensor to which this - transformation should be applied. If mapping_function uses - variadic argument `*args`, `ndim` must be specified. If - mapping_function does not use variadic arguments, ndim is - optional. - - inverse_index_map : Union[Callable, Optional[IndexMap]] - The optional pre-defined inverse index map. - When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. - Otherwise, the inverse index map will be computed on the fly. - It is the user's responsibility to ensure the correctness of the pre-defined inverse - index map. - - index_dtype : str - The default index dtype to use for input iters in the mapping function. - - Returns - ------- - index_map: IndexMap - - Returns an IndexMap representing the `mapping_function`. - - """ - params = inspect.signature(mapping_function).parameters - - args = [] - var_arg_name = None - kwargs = collections.OrderedDict() - - for name, param in params.items(): - if param.kind in [ - inspect.Parameter.POSITIONAL_ONLY, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - ]: - args.append(tvm_tirx.tirx.Var(name, index_dtype)) - - elif param.kind == inspect.Parameter.VAR_POSITIONAL: - var_arg_name = name - - elif param.kind == inspect.Parameter.KEYWORD_ONLY: - kwargs[name] = tvm_tirx.tirx.Var(name, index_dtype) - - else: - raise ValueError("transform_layout mapping may not have *args") - - # Now that all the named arguments have been collected, - # everything that remains should go to the *args, if - # specified. - if var_arg_name is not None: - assert ndim is not None, "ndim must be specified when *args is used" - num_var_args = ndim - len(args) - len(kwargs) - for i in range(num_var_args): - args.append(tvm_tirx.tirx.Var(f"{var_arg_name}_{i}", index_dtype)) - - mapping = mapping_function(*args, **kwargs) - - initial_indices = args + list(kwargs.values()) - - final_indices = [] - - if tvm_tirx.ir.is_prim_expr(mapping): - is_iterable = False - else: - try: - iter(mapping) - is_iterable = True - except TypeError: - is_iterable = False - - if is_iterable: - for val in mapping: - if tvm_tirx.ir.is_prim_expr(val): - final_indices.append(val) - else: - raise TypeError( - "Expected mapping function to return tvm_tirx.ir.Expr " - "or a list of tvm_tirx.ir.Expr. " - f"Instead received {val} of type {type(val)}." - ) - else: - final_indices.append(mapping) - - return IndexMap(initial_indices, final_indices, inverse_index_map) - - def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool: - """Return if the index maps are equivalent. - - Parameters - ---------- - other_map: IndexMap - - The IndexMap to which the comparison should be made. - - analyzer : Optional[tvm_tirx.sym.Analyzer] - - The analyzer to use while comparing the mapped indices. When - provided, its accumulated bindings and constraints are reused so - that maps that are only equivalent under those bindings can be - proven equal. - - Returns - ------- - is_equivalent: bool - - True if the two mappings represent the same - transformation, otherwise False - """ - if len(self.initial_indices) != len(other_map.initial_indices): - return False - if len(self.final_indices) != len(other_map.final_indices): - return False - - if analyzer is None: - analyzer = tvm_tirx.sym.Analyzer() - - mapped_other_final_indices = other_map.map_indices(self.initial_indices, analyzer=analyzer) - for self_index, other_index in zip(self.final_indices, mapped_other_final_indices): - if not analyzer.can_prove_equal(self_index, other_index): - return False - - return True - - def map_indices(self, indices: list[Expr], analyzer=None) -> list[Expr]: - """Apply the index map to a set of indices - - Parameters - ---------- - indices : List[Expr] - The indices to be mapped - analyzer : Optional[tvm_tirx.sym.Analyzer] - The analyzer to use while simplifying mapped indices. - - Returns - ------- - result : List[Expr] - The mapped indices - """ - return _ffi_api.IndexMapMapIndices(self, indices, analyzer) - - def map_shape(self, shape: list[Expr], analyzer=None) -> list[Expr]: - """Apply the index map to a tensor shape - - Parameters - ---------- - shape : List[Expr] - The tensor shape to be mapped - analyzer : Optional[tvm_tirx.sym.Analyzer] - The analyzer to use while simplifying mapped shape expressions. - - Returns - ------- - result : List[Expr] - The mapped shape - """ - return _ffi_api.IndexMapMapShape(self, shape, analyzer) - - def map_tensor(self, arr_src: Tensor) -> Tensor: - """Apply thie index map to transform the layout of the input Tensor - - Parameters - ---------- - arr_src : runtime.Tensor - The Tensor to be transformed - - Returns - ------- - arr_dst : runtime.Tensor - The transformed Tensor - """ - return _ffi_api.IndexMapMapTensor(self, arr_src) - - def inverse(self, shape: list[Range | Expr], analyzer=None) -> "IndexMap": - """Return the inverse of the map - - Throws an error if the function is not bijective. - - Parameters - ---------- - shape: List[Union[Range,Expr]] - - The region over which the inverse should be determined. - Used for validating that the mapping is bijective over - this range. - analyzer : Optional[tvm_tirx.sym.Analyzer] - The analyzer to use while deriving and validating the inverse. - - Returns - ------- - inverse : IndexMap - - The inverse - """ - - shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape] - return _ffi_api.IndexMapInverse(self, shape, analyzer) - - def non_surjective_inverse( - self, shape: list[Range | Expr], analyzer=None - ) -> tuple["IndexMap", Expr]: - """Return the inverse of the map - - Can be applied to transformations that introduce padding. - - Parameters - ---------- - shape: List[Union[Range,Expr]] - - The region over which the inverse should be determined. - Used for determining the predicate. - analyzer : Optional[tvm_tirx.sym.Analyzer] - The analyzer to use while deriving the inverse and padding predicate. - - Returns - ------- - result : Tuple[IndexMap, Expr] - - The inverse, and a predicate for which the inverse maps to - a valid index in the input range. - - Examples - -------- - - .. code-block:: python - - index_map = IndexMap.from_func(lambda i: [i//4, i%4]) - inverse_map, predicate = index_map.non_surjective_inverse([14]) - assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k]) - print(predicate) # Prints "(axis0==3) && (axis2 >= 2)" - """ - - shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape] - return _ffi_api.IndexMapNonSurjectiveInverse(self, shape, analyzer) diff --git a/tirx/python/tvm_tirx/tirx/function.py b/tirx/python/tvm_tirx/tirx/index_map.py similarity index 69% copy from tirx/python/tvm_tirx/tirx/function.py copy to tirx/python/tvm_tirx/tirx/index_map.py index 2a7ce72203..0ddc175520 100644 --- a/tirx/python/tvm_tirx/tirx/function.py +++ b/tirx/python/tvm_tirx/tirx/index_map.py @@ -15,182 +15,21 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=unrecognized-inline-option -"""Function data types.""" +"""Mappings between multi-dimensional indices.""" import collections import inspect -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable from typing import Optional import tvm_ffi import tvm_tirx -import tvm_tirx.runtime -from tvm_tirx.ir import BaseFunc, Range, Scriptable, SeqStmt, Stmt, Var -from tvm_tirx.ir.location import UNKNOWN_LOC, Location +from tvm_tirx.ir import Expr, Range, Var from tvm_tirx.runtime import Object +from tvm_tirx.runtime._tensor import Tensor -from ..runtime._tensor import Tensor from . import _ffi_api -from .expr import Expr - - -@tvm_ffi.register_object("tirx.Function") -class Function(BaseFunc, Scriptable): - """A function declaration expression. - - Parameters - ---------- - params: List[Union[tvm_tirx.tirx.Var, tvm_tirx.tirx.Var]] - List of input parameters to the function. - - body: Stmt | Sequence[Stmt] | None - The body of the function, or None for a declaration. - - ret_type: tvm_tirx.ir.Type - The return type annotation of the function. - - attrs: Optional[tvm.Attrs] - Attributes of the function, can be None - - loc : Location, optional - The location of this itervar in the source code. - """ - - body: SeqStmt | None - - def __init__( - self, - params, - body: Stmt | Sequence[Stmt] | None, - ret_type=None, - attrs=None, - loc: Location = UNKNOWN_LOC, - ): - if ret_type is None: - ret_type = tvm_tirx.ir.Type.missing() - param_list = [] - for x in params: - x = tvm_tirx.runtime.convert(x) if not isinstance(x, Object) else x - if not isinstance(x, Var): - raise TypeError("params can only contain Var or Var") - param_list.append(x) - - if attrs is None: - attrs = tvm_tirx.ir.make_node("ir.DictAttrs") - - self.__init_handle_by_constructor__( - _ffi_api.Function, - param_list, - body, - ret_type, - attrs, - loc, - ) # type: ignore - - @property - def is_tirx(self): - """Whether this primitive function uses the TIRx dialect.""" - return not bool(self.attrs.get("s_tir", False)) - - def with_body(self, new_body: Stmt | Sequence[Stmt] | None, loc: Location = UNKNOWN_LOC): - """Create a new Function with the same set signatures but a new body. - - Parameters - ---------- - new_body : Stmt | Sequence[Stmt] | None - The new body, or None for a declaration. - - loc : Location, optional - The location of this itervar in the source code. - - Returns - ------- - new_func : Function - The created new function. - """ - return Function( - self.params, - new_body, - ret_type=self.ret_type, - attrs=self.attrs, - loc=loc, - ) - - def specialize(self, param_map: Mapping[Var, Expr | Var]): - """Specialize parameters of Function - - Parameters - ---------- - - param_map : Mapping[Var, Union[Expr, Var]] - The mapping from function params to the instance - - Examples - -------- - We can define a TIRX function with symbolic shape: - - .. code-block:: python - - from __future__ import annotations - - @T.function - def mem_copy( - A: T.Tensor((m, n), "float32"), - B: T.Tensor((m, n), "float32"), - m: T.int32, - n: T.int32, - ) -> None: - - for i, j in T.grid(m, n): - B[i, j] = A[i, j] - - Then we can make it specialized with given shapes or tensors. - - .. code-block:: python - - a, _, m, n = mem_copy.params - func = mem_copy.specialize({a: tirx.decl_tensor((16, 16))}) - # or - func = mem_copy.specialize({n: 16, m: 16}) - - The specialized function: - - .. code-block:: python - - @T.function - def mem_copy_16_16( - A: T.Tensor((16, 16), "float32"), B: T.Tensor((16, 16), "float32") - ) -> None: - - for i, j in T.grid(16, 16): - B[i, j] = A[i, j] - - Returns - ------- - func : Function - The new function with parameter specialized - """ - return _ffi_api.Specialize(self, param_map) # type: ignore - - -def renew_def(func: Function) -> Function: - """Renew variable definitions and remap their uses in a function. - - Definition regions follow the registered structural hooks. Free variables - and function attributes are preserved, and the function signature is rebuilt. - - Parameters - ---------- - func : Function - The function to renew. - - Returns - ------- - Function - The function with fresh variable definitions. - """ - return _ffi_api.RenewDef(func) @tvm_ffi.register_object("tirx.IndexMap") diff --git a/tirx/python/tvm_tirx/tirx/op/__init__.py b/tirx/python/tvm_tirx/tirx/op/__init__.py index a167f8cb62..9ba7b024e5 100644 --- a/tirx/python/tvm_tirx/tirx/op/__init__.py +++ b/tirx/python/tvm_tirx/tirx/op/__init__.py @@ -515,7 +515,7 @@ def call_llvm_intrin(name, *args, ty, loc: Location = UNKNOWN_LOC): from tvm_tirx import codegen if isinstance(name, str): - llvm_id = codegen.llvm_lookup_intrinsic_id(name) + llvm_id = codegen.llvm.lookup_intrinsic_id(name) elif isinstance(name, IntImm): llvm_id = name.value else: @@ -559,7 +559,7 @@ def call_llvm_pure_intrin(name, *args, ty, loc: Location = UNKNOWN_LOC): from tvm_tirx import codegen if isinstance(name, str): - llvm_id = codegen.llvm_lookup_intrinsic_id(name) + llvm_id = codegen.llvm.lookup_intrinsic_id(name) elif isinstance(name, IntImm): llvm_id = name.value else: diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/op.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/op.py index fd6887e564..e788d7d90c 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/op.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/op.py @@ -26,7 +26,7 @@ from tvm_tirx import tirx as _tir from tvm_tirx import tirx as tir # pylint: disable=unused-import -from tvm_tirx.codegen import llvm_lookup_intrinsic_id +from tvm_tirx.codegen.llvm import lookup_intrinsic_id as llvm_lookup_intrinsic_id from tvm_tirx.ir import Call, CallEffectKind, Var from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.prim import _ffi_api as _prim_ffi diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py index a0fd6bcdbd..8ef324f802 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py @@ -203,7 +203,7 @@ def check_well_formed_(function: _tir.Function) -> None: def _check_module_well_formed(module: _ir.IRModule) -> None: """Validate completed functions belonging to the TIRx dialect.""" for function in module.functions.values(): - if isinstance(function, _tir.Function) and function.is_tirx: + if isinstance(function, _tir.Function) and not function.attrs.get("s_tir", False): check_well_formed_(function) diff --git a/tirx/tests/python/backend/metal/codegen/test_basic.py b/tirx/tests/python/backend/metal/codegen/test_basic.py index d582ef1cc0..dc8934d83c 100644 --- a/tirx/tests/python/backend/metal/codegen/test_basic.py +++ b/tirx/tests/python/backend/metal/codegen/test_basic.py @@ -316,7 +316,7 @@ def test_bound_symbolic_stack_allocation(bounded): assert "thread float scratch[128]" in source else: with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation extent requires a finite compile-time upper bound", ): _build_metal(Module) @@ -363,7 +363,7 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded): assert "scratch[31] =" in source else: with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation extent requires a finite compile-time upper bound", ): _build_metal(Module) @@ -411,7 +411,7 @@ def test_unbounded_symbolic_stack_allocation_rejected(): T.evaluate(scratch[0]) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation extent requires a finite compile-time upper bound", ): _build_metal(Module) @@ -436,7 +436,7 @@ def test_unbounded_uint64_symbolic_stack_allocation_rejected(): T.evaluate(scratch[0]) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation extent requires a finite compile-time upper bound", ): _build_metal(Module) @@ -461,7 +461,7 @@ def test_nonpositive_stack_allocation_rejected(extent): T.evaluate(scratch.data) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation extent requires a positive compile-time upper bound", ): _build_metal(Module) @@ -493,7 +493,7 @@ def test_stack_allocation_element_count_overflow_rejected(): T.evaluate(scratch.data) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="Metal allocation element count is too large to represent", ): _build_metal(Module) diff --git a/tirx/tests/python/backend/webgpu/codegen/test_basic.py b/tirx/tests/python/backend/webgpu/codegen/test_basic.py index 6dc852436d..20bc5b7595 100644 --- a/tirx/tests/python/backend/webgpu/codegen/test_basic.py +++ b/tirx/tests/python/backend/webgpu/codegen/test_basic.py @@ -100,7 +100,7 @@ def test_bound_symbolic_allocation(scope, bounded): assert len(re.findall(declaration + r"\s+\w+\s*:\s*array<f32,\s*128>;", source)) == 2 else: with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="WebGPU allocation extent requires a finite compile-time upper bound", ): _build_webgpu(Module) @@ -146,7 +146,7 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded): assert "scratch[31" in source else: with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="WebGPU allocation extent requires a finite compile-time upper bound", ): _build_webgpu(Module) @@ -177,7 +177,7 @@ def test_bound_symbolic_workgroup_allocation_respects_target_limit(target_limit) assert re.search(r"var<workgroup>\s+\w+\s*:\s*array<f32,\s*128>;", source) else: with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match=r"WebGPU workgroup allocations use 512 bytes, .* supports only 496 bytes", ): _build_webgpu(Module, target) @@ -201,7 +201,7 @@ def test_unbounded_symbolic_stack_allocation_rejected(): T.evaluate(scratch[0]) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="WebGPU allocation extent requires a finite compile-time upper bound", ): _build_webgpu(Module) @@ -225,7 +225,7 @@ def test_nonpositive_stack_allocation_rejected(extent): T.evaluate(scratch.data) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="WebGPU allocation extent requires a positive compile-time upper bound", ): _build_webgpu(Module) @@ -256,7 +256,7 @@ def test_stack_allocation_element_count_overflow_rejected(): T.evaluate(scratch.data) with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match="WebGPU allocation element count is too large to represent", ): _build_webgpu(Module) @@ -282,9 +282,7 @@ def test_stack_allocation_byte_size_overflow_rejected(): ) T.evaluate(scratch.data) - with pytest.raises( - tvm_tirx.error.InternalError, match="WebGPU allocation byte size is too large to represent" - ): + with pytest.raises(RuntimeError, match="WebGPU allocation byte size is too large to represent"): _build_webgpu(Module) @@ -327,7 +325,7 @@ def test_total_workgroup_allocation_above_target_limit_rejected(): second[0] = 2.0 with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match=r"WebGPU workgroup allocations use 32784 bytes, .* supports only 32768 bytes", ): _build_webgpu(Module) @@ -352,7 +350,7 @@ def test_workgroup_allocation_accounts_for_declaration_alignment(): second[0] = 2.0 with pytest.raises( - tvm_tirx.error.InternalError, + RuntimeError, match=r"WebGPU workgroup allocations use 32 bytes, .* supports only 16 bytes", ): _build_webgpu(Module, {"kind": "webgpu", "max_shared_memory_per_block": 16}) diff --git a/tirx/tests/python/codegen/llvm/test_aarch64.py b/tirx/tests/python/codegen/llvm/test_aarch64.py index 0e98e3e6a7..1db9d6540e 100644 --- a/tirx/tests/python/codegen/llvm/test_aarch64.py +++ b/tirx/tests/python/codegen/llvm/test_aarch64.py @@ -24,13 +24,13 @@ import re import pytest import tvm_tirx -from tvm_tirx.codegen import llvm_version_major +from tvm_tirx.codegen import llvm from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -68,7 +68,7 @@ def test_mul(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -106,7 +106,7 @@ def test_add(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -144,7 +144,7 @@ def test_sub(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -181,7 +181,7 @@ def test_muladd(dtype): assembly, ) - if llvm_version_major() >= 18 and dtype in ("float", "float16"): + if llvm.version_major() >= 18 and dtype in ("float", "float16"): # Newer LLVM cost models (observed with LLVM 20) prefer a fixed-width # NEON main loop over a scalable SVE loop for floating-point fmuladd # on generic AArch64 targets, so also accept the NEON form. @@ -193,7 +193,7 @@ def test_muladd(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -235,7 +235,7 @@ def test_max(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -277,7 +277,7 @@ def test_min(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -315,7 +315,7 @@ def test_div(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"] @@ -352,7 +352,7 @@ def test_mod(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -393,7 +393,7 @@ def test_eq(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", @@ -434,7 +434,7 @@ def test_neq(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"] @@ -471,7 +471,7 @@ def test_or(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"] @@ -508,7 +508,7 @@ def test_and(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.parametrize( "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"] @@ -541,7 +541,7 @@ def test_not(dtype): @pytest.mark.skipif( - llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" + llvm.version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE" ) @pytest.mark.xfail( reason="Awaiting llvm support for gathered loads", @@ -578,7 +578,7 @@ def test_memcpy(dtype): @pytest.mark.skipif( - llvm_version_major() < 13, + llvm.version_major() < 13, reason="Function attribute vscale_range() is not supported in earlier versions of LLVM", ) @pytest.mark.parametrize( @@ -589,7 +589,7 @@ def test_memcpy(dtype): # Since LLVM 19, SVE/SVE2 are optional extensions of Armv9.0-A, so # "+v9a" no longer implies "+sve" and no vscale_range() is added: # https://releases.llvm.org/19.1.0/docs/ReleaseNotes.html#changes-to-the-aarch64-backend - ("+v9a", llvm_version_major() < 19), + ("+v9a", llvm.version_major() < 19), ("+sme", True), ], ) diff --git a/tirx/tests/python/codegen/llvm/test_basic.py b/tirx/tests/python/codegen/llvm/test_basic.py index e786094d86..3d41c705b1 100644 --- a/tirx/tests/python/codegen/llvm/test_basic.py +++ b/tirx/tests/python/codegen/llvm/test_basic.py @@ -25,7 +25,7 @@ import pytest import tvm_tirx import tvm_tirx.testing -from tvm_tirx.codegen import llvm_get_intrinsic_name, llvm_lookup_intrinsic_id +from tvm_tirx.codegen import llvm from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T from tvm_tirx.support import utils @@ -46,9 +46,7 @@ def test_duplicate_function_global_symbol_diagnostic(): T.func_attr({"global_symbol": "dup_symbol", "tirx.noalias": True}) A[0] = T.float32(2) - with pytest.raises( - tvm_tirx.error.InternalError, match="Duplicate Function global_symbol 'dup_symbol'" - ) as err: + with pytest.raises(RuntimeError, match="Duplicate Function global_symbol 'dup_symbol'") as err: tvm_tirx.tirx.compile(Module, target="llvm") assert "first_unique_key" in str(err.value) assert "second_unique_key" in str(err.value) @@ -98,8 +96,8 @@ def test_llvm_void_intrin(): @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") def test_llvm_intrinsic_id(): orig_name = "llvm.x86.sse2.pmadd.wd" - intrin_id = llvm_lookup_intrinsic_id(orig_name) - name = llvm_get_intrinsic_name(intrin_id) + intrin_id = llvm.lookup_intrinsic_id(orig_name) + name = llvm.get_intrinsic_name(intrin_id) assert orig_name == name @@ -107,7 +105,7 @@ def test_llvm_intrinsic_id(): def test_llvm_overloaded_intrin(): # Name lookup for overloaded intrinsics in LLVM 4- requires a name # that includes the overloaded types. - if tvm_tirx.codegen.llvm_version_major() < 5: + if tvm_tirx.codegen.llvm.version_major() < 5: return # int1 is the type for the is_zero_undef parameter @@ -448,7 +446,7 @@ def test_alignment(): return True return False - if tvm_tirx.codegen.llvm_version_major() >= 5: + if tvm_tirx.codegen.llvm.version_major() >= 5: assert has_param_alignment() # Check for assume intrinsics. This isn't 100% accurate, since it just @@ -681,9 +679,9 @@ def test_dwarf_debug_information(): C[i0_0 * 4 + i0_1] = A[i0_0 * 4 + i0_1] + B[i0_0 * 4 + i0_1] def check_llvm_object(): - if tvm_tirx.codegen.llvm_version_major() < 5: + if tvm_tirx.codegen.llvm.version_major() < 5: return - if tvm_tirx.codegen.llvm_version_major() > 6: + if tvm_tirx.codegen.llvm.version_major() > 6: return # build two functions mod = tvm_tirx.IRModule( @@ -719,9 +717,9 @@ def test_dwarf_debug_information(): assert re.search(r"""DW_AT_name.*fadd2""", str(output)) def check_llvm_ir(): - if tvm_tirx.codegen.llvm_version_major() < 5: + if tvm_tirx.codegen.llvm.version_major() < 5: return - if tvm_tirx.codegen.llvm_version_major() > 6: + if tvm_tirx.codegen.llvm.version_major() > 6: return # build two functions mod = tvm_tirx.IRModule( diff --git a/tirx/tests/python/codegen/llvm/test_riscv.py b/tirx/tests/python/codegen/llvm/test_riscv.py index 5ff26a28d0..4182e84925 100644 --- a/tirx/tests/python/codegen/llvm/test_riscv.py +++ b/tirx/tests/python/codegen/llvm/test_riscv.py @@ -118,7 +118,7 @@ def test_rvv_vscale_llvm_dbginfo(target): redsum = T.call_llvm_intrin( 'llvm.riscv.vfredusum', T.Broadcast(T.float32(0.0), T.vscale() * 2), prod, zero, T.uint64(7), T.uint64(8), ty='float32xvscalex2') # fmt: on - # tvm.error.InternalError: Can't fetch the lanes of a scalable vector at a compile time. + # tvm.RuntimeError: Can't fetch the lanes of a scalable vector at a compile time. with tvm_tirx.target.Target(target): f = tvm_tirx.tirx.compile(rvv_with_vscale, target).mod diff --git a/tirx/tests/python/codegen/llvm/test_scalable_vector.py b/tirx/tests/python/codegen/llvm/test_scalable_vector.py index 8bb365a728..0e703fbc4a 100644 --- a/tirx/tests/python/codegen/llvm/test_scalable_vector.py +++ b/tirx/tests/python/codegen/llvm/test_scalable_vector.py @@ -25,12 +25,12 @@ import pytest import tvm_tirx import tvm_tirx.testing -from tvm_tirx.codegen import llvm_version_major +from tvm_tirx.codegen import llvm from tvm_tirx.script import tirx as T @pytest.mark.skipif( - llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" + llvm.version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" ) @pytest.mark.parametrize( "target", @@ -63,7 +63,7 @@ def test_codegen_vscale(target): @pytest.mark.skipif( - llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" + llvm.version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" ) @pytest.mark.parametrize( "target", @@ -96,7 +96,7 @@ def test_scalable_buffer_load_store(target): @pytest.mark.skipif( - llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" + llvm.version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM" ) @pytest.mark.parametrize( "target", diff --git a/tirx/tests/python/codegen/llvm/test_x86.py b/tirx/tests/python/codegen/llvm/test_x86.py index 4c92a3722c..ea32df9b35 100644 --- a/tirx/tests/python/codegen/llvm/test_x86.py +++ b/tirx/tests/python/codegen/llvm/test_x86.py @@ -25,7 +25,7 @@ from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T from tvm_tirx.testing import env -llvm_version = tvm_tirx.codegen.llvm_version_major() +llvm_version = tvm_tirx.codegen.llvm.version_major() machine = platform.machine() if machine not in ["x86_64", "AMD64", "amd64"]: diff --git a/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py b/tirx/tests/python/runtime/test_validate_wheel_library.py similarity index 96% rename from tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py rename to tirx/tests/python/runtime/test_validate_wheel_library.py index ea92272ea6..eb0a5cc52d 100644 --- a/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py +++ b/tirx/tests/python/runtime/test_validate_wheel_library.py @@ -14,9 +14,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Post-install runtime and AOT checks. +"""Post-install wheel library and AOT checks. -These live in ``tests/python/all-platform-minimal-test`` so the standard suite and +These live in ``tests/python/runtime`` so the standard suite and the cibuildwheel ``test-command`` run a single pytest invocation. The assertions for wheel packaging check things the rest of the suite cannot check -- that LLVM is enabled (the other LLVM test merely *skips* when LLVM is absent) and that the CUDA diff --git a/tirx/tests/python/script/test_basic_usage.py b/tirx/tests/python/script/test_basic_usage.py index 458d6d2e2c..ecfdda3171 100644 --- a/tirx/tests/python/script/test_basic_usage.py +++ b/tirx/tests/python/script/test_basic_usage.py @@ -28,7 +28,7 @@ import pytest from minilang import Value import tvm_tirx -from tvm_tirx import DataType, error, ir +from tvm_tirx import DataType, ir from tvm_tirx.ir import prim from tvm_tirx.ir._overload_prim_expr import EqualOp from tvm_tirx.script import ir as I @@ -291,7 +291,7 @@ def test_loop_control_validation_preserves_valid_and_unchecked_ir(): invalid = tirx.Function(params=[], body=tvm_tirx.ir.Break()) ir.assert_structural_equal(invalid.body, tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Break()])) assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False) - with pytest.raises(error.InternalError, match="requires an enclosing loop"): + with pytest.raises(RuntimeError, match="requires an enclosing loop"): tirx.analysis.verify_well_formed(invalid) @T.function diff --git a/tirx/tests/python/support/test_config.py b/tirx/tests/python/support/test_config.py index 009847d985..f689b3bbaa 100644 --- a/tirx/tests/python/support/test_config.py +++ b/tirx/tests/python/support/test_config.py @@ -139,6 +139,7 @@ def test_cmake_targets_without_ffi(tmp_path): libdir = _installation(tmp_path / "install") cmake_dir = libdir / "cmake/tvm_tirx" shutil.copy2(_CONFIG_PATH, cmake_dir / "config.py") + shutil.copy2(_CONFIG_PATH.with_name("libinfo.py"), cmake_dir / "libinfo.py") shutil.copy2( _CONFIG_PATH.parents[2] / "cmake/tvm_tirx-config.cmake", cmake_dir / "tvm_tirx-config.cmake", @@ -186,13 +187,13 @@ def _check_component_closure_and_artifacts(tmp_path, monkeypatch): for filename in ("libtvm_runtime_cuda.so", "libtvm_runtime_metal.so", "libtvm_ffi.so"): (libdir / filename).touch() monkeypatch.setattr( - config, + config.libinfo, "_ffi_info", lambda: SimpleNamespace(find_libtvm_ffi=lambda: str(libdir / "libtvm_ffi.so")), ) def names(components): - return [path.name for path in config._link_inputs(libdir, "linux", components)] + return [path.name for path in config.libinfo._link_inputs(libdir, "linux", components)] assert names(["runtime", "cuda", "runtime", "ffi"]) == [ "libtvm_runtime_cuda.so", @@ -214,7 +215,7 @@ def _check_cli_link_flags_and_raw_inputs(tmp_path, monkeypatch, capsys, platform for name in ("tvm_runtime", "tvm_tirx_compiler", "tvm_ffi"): (libdir / f"lib{name}.{suffix}").touch() monkeypatch.setattr( - config, + config.libinfo, "_ffi_info", lambda: SimpleNamespace(find_libtvm_ffi=lambda: str(libdir / f"libtvm_ffi.{suffix}")), ) @@ -280,5 +281,5 @@ def _check_runtime_only_lib64_and_explicit_prefix_ignore_environment(tmp_path, m (libdir / "libtvm_tirx_compiler.so").unlink() unrelated = _installation(tmp_path / "unrelated") monkeypatch.setenv("TVM_TIRX_LIBRARY_PATH", str(unrelated)) - assert config._paths(prefix, None) == (prefix, libdir) - assert config._available(libdir, "linux") == ["runtime"] + assert config.libinfo._paths(prefix, None) == (prefix, libdir) + assert config.libinfo._available(libdir, "linux") == ["runtime"] diff --git a/tirx/tests/python/target/test_llvm_features_info.py b/tirx/tests/python/target/test_llvm_features_info.py index 4eb07fd99d..df2b9643d0 100644 --- a/tirx/tests/python/target/test_llvm_features_info.py +++ b/tirx/tests/python/target/test_llvm_features_info.py @@ -20,7 +20,7 @@ import tvm_tirx from tvm_tirx import codegen from tvm_tirx.target import Target -LLVM_VERSION = codegen.llvm_version_major() +LLVM_VERSION = codegen.llvm.version_major() def test_llvm_targets(capfd): @@ -29,27 +29,27 @@ def test_llvm_targets(capfd): ## # check blank results - assert len(codegen.llvm_get_targets()) - assert len(codegen.llvm_get_system_cpu()) - assert len(codegen.llvm_get_system_triple()) - assert len(codegen.llvm_get_system_x86_vendor()) + assert len(codegen.llvm.get_targets()) + assert len(codegen.llvm.get_system_cpu()) + assert len(codegen.llvm.get_system_triple()) + assert len(codegen.llvm.get_system_x86_vendor()) # check ffi vs python assert ( - codegen.llvm_get_system_cpu() == tvm_tirx.get_global_func("codegen.llvm_get_system_cpu")() + codegen.llvm.get_system_cpu() == tvm_tirx.get_global_func("codegen.llvm_get_system_cpu")() ) assert ( - codegen.llvm_get_system_triple() + codegen.llvm.get_system_triple() == tvm_tirx.get_global_func("codegen.llvm_get_system_triple")() ) assert ( - codegen.llvm_get_system_x86_vendor() + codegen.llvm.get_system_x86_vendor() == tvm_tirx.get_global_func("codegen.llvm_get_system_x86_vendor")() ) - assert str(codegen.llvm_get_targets()) == str( + assert str(codegen.llvm.get_targets()) == str( tvm_tirx.get_global_func("codegen.llvm_get_targets")() ) - tvm_tirx.codegen.llvm_get_cpu_features( + tvm_tirx.codegen.llvm.get_cpu_features( tvm_tirx.target.Target({"kind": "llvm", "mtriple": "x86_64-linux-gnu", "mcpu": "dummy"}) ) expected_str = ( @@ -86,7 +86,7 @@ def test_target_features(min_llvm_version, llvm_target, cpu_arch, cpu_features, ## legalize llvm_target ## - assert llvm_target in codegen.llvm_get_targets() + assert llvm_target in codegen.llvm.get_targets() ## ## legalize cpu_arch @@ -94,11 +94,11 @@ def test_target_features(min_llvm_version, llvm_target, cpu_arch, cpu_features, ### with context with target: - assert cpu_arch in codegen.llvm_get_cpu_archlist() + assert cpu_arch in codegen.llvm.get_cpu_archlist() ### no context but with expicit target - assert cpu_arch in codegen.llvm_get_cpu_archlist(target) + assert cpu_arch in codegen.llvm.get_cpu_archlist(target) # check ffi vs python - assert str(codegen.llvm_get_cpu_archlist(target)) == str( + assert str(codegen.llvm.get_cpu_archlist(target)) == str( tvm_tirx.get_global_func("codegen.llvm_get_cpu_archlist")(target) ) @@ -108,11 +108,11 @@ def test_target_features(min_llvm_version, llvm_target, cpu_arch, cpu_features, ### with context with target: - assert codegen.llvm_cpu_has_features(cpu_features) == is_supported + assert codegen.llvm.cpu_has_features(cpu_features) == is_supported ### no context but with expicit target - assert codegen.llvm_cpu_has_features(cpu_features, target) == is_supported + assert codegen.llvm.cpu_has_features(cpu_features, target) == is_supported # check ffi vs python for feat in cpu_features: - assert str(codegen.llvm_cpu_has_features(feat, target)) == str( + assert str(codegen.llvm.cpu_has_features(feat, target)) == str( tvm_tirx.get_global_func("codegen.llvm_cpu_has_feature")(feat, target) ) diff --git a/tirx/tests/python/target/test_riscv_features.py b/tirx/tests/python/target/test_riscv_features.py index 3a805c4270..c81ca2e4c4 100644 --- a/tirx/tests/python/target/test_riscv_features.py +++ b/tirx/tests/python/target/test_riscv_features.py @@ -19,10 +19,10 @@ import pytest import tvm_tirx from tvm_tirx import codegen -from tvm_tirx.codegen import llvm_get_vector_width, target_has_features +from tvm_tirx.codegen import llvm, target_has_features from tvm_tirx.target import Target -LLVM_VERSION = codegen.llvm_version_major() +LLVM_VERSION = codegen.llvm.version_major() # fmt: off @@ -63,4 +63,4 @@ def test_riscv_rvv_features(min_llvm_version, tvm_target, vec_width): return with Target(tvm_target): - assert llvm_get_vector_width() == vec_width + assert llvm.get_vector_width() == vec_width diff --git a/tirx/tests/python/target/test_x86_features.py b/tirx/tests/python/target/test_x86_features.py index d638d202c6..59da3b2015 100644 --- a/tirx/tests/python/target/test_x86_features.py +++ b/tirx/tests/python/target/test_x86_features.py @@ -22,7 +22,7 @@ from tvm_tirx import codegen from tvm_tirx.codegen import target_has_features from tvm_tirx.target import Target -LLVM_VERSION = codegen.llvm_version_major() +LLVM_VERSION = codegen.llvm.version_major() # Some x86 features have been removed from upstream LLVM. Tests for these # features only meaningfully run on LLVM versions that still recognise them. @@ -197,7 +197,7 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo try: assert target_has_features(x86_feature) == is_supported assert False - except tvm_tirx.error.InternalError as e: + except RuntimeError as e: msg = str(e) assert msg.find("Check failed: (allow_not_defined) is false: Target context required") != -1 @@ -209,7 +209,7 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo == is_supported ) assert False - except tvm_tirx.error.InternalError as e: + except RuntimeError as e: msg = str(e) assert ( msg.find("Check failed: (allow_not_defined) is false: Target context required") @@ -253,4 +253,4 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo ) == is_supported # check for feature in target's llvm full x86 CPU feature list if (not list(Target(tvm_target).attrs.get("mattr", []))) and isinstance(x86_feature, str): - assert (x86_feature in codegen.llvm_get_cpu_features()) == is_supported + assert (x86_feature in codegen.llvm.get_cpu_features()) == is_supported diff --git a/tirx/tests/python/testing/test_env.py b/tirx/tests/python/testing/test_env.py deleted file mode 100644 index 2f7105b6be..0000000000 --- a/tirx/tests/python/testing/test_env.py +++ /dev/null @@ -1,332 +0,0 @@ -# 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. -"""Tests for the thin ``tvm_tirx.testing.env`` capability probes.""" - -import subprocess -import sys -import textwrap -from types import SimpleNamespace - -import pytest - -import tvm_tirx -import tvm_tirx.testing -from tvm_tirx.testing import env - -# Probes that take no arguments and must return a plain bool without raising. -_BOOL_PROBES = [ - # runtime device - env.has_cuda, - env.has_rocm, - env.has_vulkan, - env.has_metal, - env.has_opencl, - env.has_nvptx, - env.has_llvm, - env.has_gpu, - # build support - env.has_cudnn, - env.has_cublas, - env.has_nccl, - env.has_hipblas, - env.has_nvshmem, - # version / capability - env.has_matrixcore, - env.has_cudagraph, - # toolchain / environment - env.has_hexagon, - env.has_hexagon_toolchain, - env.has_adreno_opencl, - # cpu features - env.has_x86_vnni, - env.has_x86_avx512, - # host architecture -] - - [email protected]("probe", _BOOL_PROBES, ids=lambda p: p.__name__) -def test_probe_returns_bool(probe): - """Every probe returns a real bool and never raises during collection/run.""" - assert isinstance(probe(), bool) - - -def test_has_cuda_implies_device(): - """has_cuda() requires a device (it also requires the kind to be enabled).""" - if env.has_cuda(): - assert tvm_tirx.cuda().exist - - -def test_has_gpu_is_raw_any_device(): - """has_gpu() is the disjunction of the raw device checks (no target gating).""" - any_device = ( - env._device_exists("cuda") # pylint: disable=protected-access - or env._device_exists("rocm") # pylint: disable=protected-access - or env._device_exists("opencl") # pylint: disable=protected-access - or env._device_exists("metal") # pylint: disable=protected-access - or env._device_exists("vulkan") # pylint: disable=protected-access - ) - assert env.has_gpu() == any_device - - -def test_target_enabled_respects_tvm_test_targets(monkeypatch): - """A device kind excluded from TVM_TEST_TARGETS is reported as not enabled.""" - env._target_enabled.cache_clear() # pylint: disable=protected-access - monkeypatch.setenv("TVM_TEST_TARGETS", "cuda;llvm") - try: - assert env._target_enabled("cuda") # pylint: disable=protected-access - assert env._target_enabled("llvm") # pylint: disable=protected-access - assert not env._target_enabled("opencl") # pylint: disable=protected-access - assert not env._target_enabled("metal") # pylint: disable=protected-access - finally: - env._target_enabled.cache_clear() # pylint: disable=protected-access - - -def test_cuda_compute_is_monotonic(): - """has_cuda_compute is monotone in the requested version.""" - if not env.has_cuda(): - # Without a CUDA device every query is False, including the (0, 0) floor. - assert not env.has_cuda_compute(1, 0) - assert not env.has_cuda_compute(0, 0) - return - # A device that satisfies (major, minor) also satisfies anything lower. - assert env.has_cuda_compute(1, 0) - assert env.has_cuda_compute(0, 0) - - [email protected]( - ("compute_version", "expected"), - [ - ("8.0", "sm_80"), - ("9.0", "sm_90a"), - ("10.0", "sm_100a"), - ("10.3", "sm_103a"), - ("invalid", None), - ], -) -def test_cuda_arch_from_compute_version(compute_version, expected): - assert env._cuda_arch_from_compute_version(compute_version) == expected - - -def test_has_cuda_arch_matches_detected_device(): - actual = env.cuda_arch() - if actual is None: - assert not env.has_cuda_arch("sm_100a") - else: - assert env.has_cuda_arch(actual) - - -def test_cuda_arch_uses_requested_device(monkeypatch): - class Device: - exist = True - compute_version = "10.3" - - requested = [] - - def cuda(device_id): - requested.append(device_id) - return Device() - - env.cuda_arch.cache_clear() - monkeypatch.setattr(tvm_tirx, "cuda", cuda) - try: - assert env.cuda_arch(7) == "sm_103a" - assert requested == [7] - finally: - env.cuda_arch.cache_clear() - - -def test_has_multi_gpu_is_bool(): - assert isinstance(env.has_multi_gpu(), bool) - assert isinstance(env.has_multi_gpu(1), bool) - # Requiring a single device is at least as permissive as requiring two. - assert env.has_multi_gpu(1) or not env.has_multi_gpu(2) - - -def test_llvm_min_version_is_monotone(): - if not env.has_llvm(): - assert not env.has_llvm_min_version(1) - return - # An LLVM that satisfies a higher floor also satisfies a lower one. - assert env.has_llvm_min_version(1) - - -def test_probes_are_memoized(): - """Probes are cached so the driver/subprocess is hit once per process.""" - env.has_cuda() - info = env._device_exists.cache_info() # pylint: disable=protected-access - assert info.hits + info.misses >= 1 - - [email protected] -def clear_env_cache(): - probes = ( - env._target_enabled, - env._runtime_enabled, - env._device_exists, - env.build_flag_enabled, - env._llvm_version_major, - env.has_hexagon_toolchain, - env.has_hexagon, - ) - for probe in probes: - probe.cache_clear() - yield - for probe in probes: - probe.cache_clear() - - -def test_unavailable_selected_target_does_not_enable_llvm(monkeypatch, clear_env_cache): - monkeypatch.setenv("TVM_TEST_TARGETS", "cuda") - monkeypatch.setattr(tvm_tirx.runtime, "enabled", lambda kind: kind == "llvm") - - def unexpected_device_lookup(target): - pytest.fail(f"Excluded or unavailable target was probed: {target}") - - monkeypatch.setattr(tvm_tirx, "device_from_target", unexpected_device_lookup) - assert not tvm_tirx.testing.device_enabled("cuda") - assert not tvm_tirx.testing.device_enabled("llvm") - - [email protected]("target_form", ["string", "dict", "target"]) -def test_device_enabled_accepts_target_forms(monkeypatch, clear_env_cache, target_form): - monkeypatch.setenv("TVM_TEST_TARGETS", '{"kind": "llvm"};cuda') - monkeypatch.setattr(tvm_tirx.runtime, "enabled", lambda kind: kind == "llvm") - monkeypatch.setattr(tvm_tirx, "device_from_target", lambda target: SimpleNamespace(exist=True)) - target = { - "string": "llvm -num-cores=2", - "dict": {"kind": "llvm", "num-cores": 2}, - "target": tvm_tirx.target.Target({"kind": "llvm", "num-cores": 2}), - }[target_form] - assert tvm_tirx.testing.device_enabled(target) - - [email protected]("kind", ["llvm", "cuda"]) [email protected]( - ("runtime_enabled", "device_exists", "expected"), - [(False, True, False), (True, False, False), (True, True, True)], -) -def test_device_enabled_requires_runtime_and_device( - monkeypatch, clear_env_cache, kind, runtime_enabled, device_exists, expected -): - # In particular, the CPU can exist even when LLVM was not built. - monkeypatch.setenv("TVM_TEST_TARGETS", kind) - monkeypatch.setattr(tvm_tirx.runtime, "enabled", lambda target: runtime_enabled) - monkeypatch.setattr( - tvm_tirx, "device_from_target", lambda target: SimpleNamespace(exist=device_exists) - ) - assert tvm_tirx.testing.device_enabled(kind) is expected - - [email protected]("cudnn_available", [None, False, True]) -def test_device_enabled_cudnn_runtime_probe(monkeypatch, clear_env_cache, cudnn_available): - monkeypatch.setenv("TVM_TEST_TARGETS", '{"kind": "cuda", "libs": ["cudnn"]}') - monkeypatch.setattr(tvm_tirx.runtime, "enabled", lambda kind: True) - monkeypatch.setattr(tvm_tirx, "device", lambda *args: SimpleNamespace(exist=True)) - monkeypatch.setattr(tvm_tirx.support, "libinfo", lambda: {"USE_CUDNN": "ON"}) - exists = None if cudnn_available is None else lambda: cudnn_available - monkeypatch.setattr(tvm_tirx, "get_global_func", lambda name, allow_missing: exists) - assert tvm_tirx.testing.device_enabled("cuda") is (cudnn_available is True) - - [email protected]( - ("runtime", "toolchain", "llvm_target", "serial", "can_compile", "can_run"), - [ - (False, True, True, True, True, False), - (True, False, True, True, False, False), - (True, True, False, True, False, False), - (True, True, True, False, True, False), - (True, True, True, True, True, True), - (True, True, None, True, False, False), - ], -) -def test_hexagon_environment_requirements( - monkeypatch, clear_env_cache, runtime, toolchain, llvm_target, serial, can_compile, can_run -): - monkeypatch.setenv("TVM_TEST_TARGETS", "hexagon") - for variable, present in ( - ("HEXAGON_TOOLCHAIN", toolchain), - ("ANDROID_SERIAL_NUMBER", serial), - ): - if present: - monkeypatch.setenv(variable, "test-value") - else: - monkeypatch.delenv(variable, raising=False) - monkeypatch.setattr( - tvm_tirx.runtime, - "enabled", - lambda kind: runtime if kind == "hexagon" else llvm_target is not None, - ) - monkeypatch.setattr( - tvm_tirx.codegen, "llvm_get_targets", lambda: ["hexagon"] if llvm_target else [] - ) - assert env.has_hexagon_toolchain() is can_compile - assert env.has_hexagon() is can_run - - -def test_core_testing_imports_without_domain_modules(): - # Let TVM finish its own startup, then make the domain modules unavailable to testing. - code = textwrap.dedent( - """ - import os - import sys - import tvm_tirx - - for module in list(sys.modules): - if module == "tvm_tirx.testing" or module.startswith("tvm_tirx.testing."): - del sys.modules[module] - if hasattr(tvm_tirx, "testing"): - del tvm_tirx.testing - - for name in ("te", "contrib", "s_tir"): - prefix = "tvm_tirx." + name - for module in list(sys.modules): - if module == prefix or module.startswith(prefix + "."): - sys.modules[module] = None - sys.modules[prefix] = None - if hasattr(tvm_tirx, name): - delattr(tvm_tirx, name) - - import tvm_tirx.testing - from tvm_tirx.testing import env - - tvm_tirx.testing.assert_allclose([1], [1]) - tvm_tirx.runtime.enabled = lambda kind: True - tvm_tirx.codegen.llvm_get_targets = lambda: ["hexagon"] - os.environ["HEXAGON_TOOLCHAIN"] = "test-toolchain" - os.environ["ANDROID_SERIAL_NUMBER"] = "test-device" - assert env.has_hexagon_toolchain() - assert env.has_hexagon() - """ - ) - subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True) - - -# --- demonstration of the target idiom ------------------------------------- -# -# The standard gating idiom: a plain registered pytest marker (for ``-m`` -# selection) plus a skipif backed by a thin env probe (for runtime gating). - - [email protected] [email protected](not env.has_cuda(), reason="need cuda") -def test_thin_cuda_idiom(): - dev = tvm_tirx.cuda() - assert dev.exist - - -if __name__ == "__main__": - tvm_tirx.testing.main() diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py b/tirx/tests/python/tirx/analysis/test_check_well_defined.py index 6545216b5d..5e35f9ff2e 100644 --- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py +++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py @@ -66,7 +66,7 @@ def test_error_for_out_of_scope_usage(): func = tvm_tirx.tirx.Function([], body) with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), + (ValueError, RuntimeError), match="Invalid use of undefined variable i at .* no longer in-scope.", ): tvm_tirx.tirx.analysis.verify_well_formed(func) @@ -84,7 +84,7 @@ def test_error_for_nested_rebind_usage(): T.evaluate(i) with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), + (ValueError, RuntimeError), match="ill-formed, due to multiple nested definitions of variable i", ): tvm_tirx.tirx.analysis.verify_well_formed(func) @@ -108,7 +108,7 @@ def test_error_for_repeated_binding(): T.evaluate(i) with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), + (ValueError, RuntimeError), match="multiple nested definitions of variable i", ): tvm_tirx.tirx.analysis.verify_well_formed(func) @@ -131,9 +131,7 @@ def test_error_for_cross_function_reuse(): T.bind(42, var=i) T.evaluate(i) - with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), match="multiple definitions of variable i" - ): + with pytest.raises((ValueError, RuntimeError), match="multiple definitions of variable i"): tvm_tirx.tirx.analysis.verify_well_formed(mod) @@ -177,7 +175,7 @@ def test_reuse_of_region_parameter_across_functions_is_ill_formed(): A[threadIdx_x] = A[threadIdx_x] + T.float32(1) with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), + (ValueError, RuntimeError), match="multiple definitions of variable threadIdx_x", ): tvm_tirx.tirx.analysis.verify_well_formed(mod) @@ -222,7 +220,7 @@ def test_error_message_without_previous_definition_location(): T.bind(99, var=x) # This should trigger the error T.evaluate(x) - with pytest.raises((ValueError, tvm_tirx.error.InternalError)) as exc_info: + with pytest.raises((ValueError, RuntimeError)) as exc_info: tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -247,7 +245,7 @@ def test_error_message_with_previous_definition_location(): T.bind(99, var=x) # This should trigger the error T.evaluate(x) - with pytest.raises((ValueError, tvm_tirx.error.InternalError)) as exc_info: + with pytest.raises((ValueError, RuntimeError)) as exc_info: tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -278,7 +276,7 @@ def test_sequential_redefinition_with_location(): T.bind(2, var=x) # This should trigger the error T.evaluate(x) - with pytest.raises((ValueError, tvm_tirx.error.InternalError)) as exc_info: + with pytest.raises((ValueError, RuntimeError)) as exc_info: tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -349,7 +347,7 @@ def test_tensor_load_asserted_type_matches_source_and_indices(): malformed = tvm_tirx.ir.load_json(json.dumps(graph)) assert not tvm_tirx.tirx.analysis.verify_well_formed(malformed, assert_mode=False) - with pytest.raises(tvm_tirx.error.InternalError, match="asserts result type"): + with pytest.raises(RuntimeError, match="asserts result type"): tvm_tirx.tirx.analysis.verify_well_formed(malformed) @@ -368,13 +366,13 @@ def test_tensor_load_malformed_indices_return_false_without_asserting(): rank_mismatch_indices.pop() rank_mismatch = tvm_tirx.ir.load_json(json.dumps(rank_mismatch_graph)) assert not tvm_tirx.tirx.analysis.verify_well_formed(rank_mismatch, assert_mode=False) - with pytest.raises(tvm_tirx.error.InternalError, match="indexes 2-dimensional tensor"): + with pytest.raises(RuntimeError, match="indexes 2-dimensional tensor"): tvm_tirx.tirx.analysis.verify_well_formed(rank_mismatch) indices[0], indices[1] = indices[1], indices[0] non_final_vector = tvm_tirx.ir.load_json(json.dumps(graph)) assert not tvm_tirx.tirx.analysis.verify_well_formed(non_final_vector, assert_mode=False) - with pytest.raises(tvm_tirx.error.InternalError, match="only the final index"): + with pytest.raises(RuntimeError, match="only the final index"): tvm_tirx.tirx.analysis.verify_well_formed(non_final_vector) diff --git a/tirx/tests/python/tirx/base/test_imm_values.py b/tirx/tests/python/tirx/base/test_imm_values.py index f267b63bce..1eb03c33cd 100644 --- a/tirx/tests/python/tirx/base/test_imm_values.py +++ b/tirx/tests/python/tirx/base/test_imm_values.py @@ -59,8 +59,8 @@ def test_tir_invalid_intimm(dtype, literals): for l in literals: # Out-of-range positive literals raise a builtin ValueError from # the IntImm range check; negative-into-unsigned raises an - # InternalError ("cannot make uint from negative value") which is a - # RuntimeError subclass. Accept either. + # InternalError ("cannot make uint from negative value") mapped to + # RuntimeError. Accept either. with pytest.raises((RuntimeError, ValueError)): tirx.const(l, dtype) diff --git a/tirx/tests/python/tirx/base/test_scalable_datatype.py b/tirx/tests/python/tirx/base/test_scalable_datatype.py index 45bb8f4553..a5a11a9759 100644 --- a/tirx/tests/python/tirx/base/test_scalable_datatype.py +++ b/tirx/tests/python/tirx/base/test_scalable_datatype.py @@ -19,7 +19,7 @@ import pytest import tvm_tirx from tvm_tirx import tirx -from tvm_tirx.codegen import llvm_version_major +from tvm_tirx.codegen import llvm from tvm_tirx.script import tirx as T """ @@ -36,11 +36,11 @@ def test_create_scalable_data_type_python_api(): # the old name from the intrinsic table: # https://releases.llvm.org/20.1.0/docs/ReleaseNotes.html _STEPVECTOR_NAME = ( - "llvm.stepvector" if llvm_version_major() >= 20 else "llvm.experimental.stepvector" + "llvm.stepvector" if llvm.version_major() >= 20 else "llvm.experimental.stepvector" ) [email protected](llvm_version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.") [email protected](llvm.version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.") def test_create_scalable_tir_intrin(): intrin = tirx.call_llvm_intrin(_STEPVECTOR_NAME, ty="int32xvscalex4") assert intrin.ty.dtype == "int32xvscalex4" @@ -49,7 +49,7 @@ def test_create_scalable_tir_intrin(): ) [email protected](llvm_version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.") [email protected](llvm.version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.") def test_tvm_script_create_scalable_tir_intrin(): @T.function def my_func(): diff --git a/tirx/tests/python/tirx/codegen/test_cuda.py b/tirx/tests/python/tirx/codegen/test_cuda.py index 3918130d40..1bb95ed34e 100644 --- a/tirx/tests/python/tirx/codegen/test_cuda.py +++ b/tirx/tests/python/tirx/codegen/test_cuda.py @@ -187,9 +187,7 @@ def test_legacy_tensorcore_ops_require_downstream_codegen(op_name): "global_symbol", "main" ) target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}) - with pytest.raises( - tvm_tirx.error.InternalError, match="Unresolved call.*" + re.escape(call.op.name) - ): + with pytest.raises(RuntimeError, match="Unresolved call.*" + re.escape(call.op.name)): tvm_tirx.get_global_func("target.build.cuda")(tvm_tirx.IRModule.from_expr(func), target) @@ -204,7 +202,7 @@ def test_warp_storage_requires_explicit_local_layout(target_kind): output[0] = fragment[0] target = tvm_tirx.target.Target(target_kind) - with pytest.raises(tvm_tirx.error.InternalError, match="warp storage is unsupported"): + with pytest.raises(RuntimeError, match="warp storage is unsupported"): tvm_tirx.get_global_func("target.build." + target_kind)( tvm_tirx.IRModule.from_expr(main), target ) diff --git a/tirx/tests/python/tirx/codegen/test_hopper.py b/tirx/tests/python/tirx/codegen/test_hopper.py index 23cbff0d9c..fad00405a7 100644 --- a/tirx/tests/python/tirx/codegen/test_hopper.py +++ b/tirx/tests/python/tirx/codegen/test_hopper.py @@ -608,7 +608,7 @@ def test_cp_async_bulk_tensor_global_to_shared_unicast(dtype, inputs): ], ) def test_tensormap_encode_tiled_runtime_validation(shape, dtype, encode_args, error_msg): - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): _run_tensormap_encode(shape, dtype, encode_args) diff --git a/tirx/tests/python/tirx/codegen/test_ptx_addr.py b/tirx/tests/python/tirx/codegen/test_ptx_addr.py index 5d9cf124ee..e4bf18de4b 100644 --- a/tirx/tests/python/tirx/codegen/test_ptx_addr.py +++ b/tirx/tests/python/tirx/codegen/test_ptx_addr.py @@ -225,9 +225,7 @@ def test_ptx_addr_unrolled_expression_and_dynamic_rejection(): T.ptx.ld.global_.b32(value, T.ptx.addr(src.data, i * 4)) for func in (thread_dynamic, loop_dynamic): - with pytest.raises( - (ValueError, tvm_tirx.error.InternalError), match="must become a compile-time" - ): + with pytest.raises((ValueError, RuntimeError), match="must become a compile-time"): _cuda_source(func) diff --git a/tirx/tests/python/tirx/codegen/test_ptx_dialect.py b/tirx/tests/python/tirx/codegen/test_ptx_dialect.py index df7c3438ba..a30a0e516a 100644 --- a/tirx/tests/python/tirx/codegen/test_ptx_dialect.py +++ b/tirx/tests/python/tirx/codegen/test_ptx_dialect.py @@ -1842,7 +1842,7 @@ def test_ptx_logic_shift_dispatch(): d = T.local_scalar("uint32") T.ptx.lop3.b32(d, A[0], A[0], A[0], A[0]) - with pytest.raises((ValueError, tvm_tirx.error.InternalError), match="compile-time constants"): + with pytest.raises((ValueError, RuntimeError), match="compile-time constants"): _cuda_source(lut_runtime) @T.function diff --git a/tirx/tests/python/tirx/script/test_basic_usage.py b/tirx/tests/python/tirx/script/test_basic_usage.py index 7d8dce025f..b2c048f047 100644 --- a/tirx/tests/python/tirx/script/test_basic_usage.py +++ b/tirx/tests/python/tirx/script/test_basic_usage.py @@ -229,7 +229,7 @@ from tvm_tirx.script import tirx as T def func(): x: T.handle = T.int64(0) """ - with pytest.raises(tvm_tirx.error.InternalError): + with pytest.raises(RuntimeError): from_source(src_handle) # 3. Banned: non-PrimType annotation without T.let @@ -240,7 +240,7 @@ from tvm_tirx.ir import PtrType, PrimType def func(): x: T.Var(name="x", ty=PtrType(PrimType("float16"))) = T.int64(0) """ - with pytest.raises(tvm_tirx.error.InternalError): + with pytest.raises(RuntimeError): from_source(src_ptr) # 4. An explicit mutable scalar declaration retains updates — round-trip @@ -571,14 +571,14 @@ def test_thread_return_is_distinct_from_function_return(): def test_loop_control_validation_preserves_valid_and_unchecked_ir(): # Invalid loop placement must be rejected, while disabled checks preserve the original IR. - from tvm_tirx import error, ir, tirx + from tvm_tirx import ir, tirx invalid = tirx.Function(params=[], body=tvm_tirx.ir.Break()) # Direct construction retains the native statement for an explicit verifier pass. ir.assert_structural_equal(invalid.body, tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Break()])) assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False) - with pytest.raises(error.InternalError, match="requires an enclosing loop"): + with pytest.raises(RuntimeError, match="requires an enclosing loop"): tirx.analysis.verify_well_formed(invalid) @T.function diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py index 1eaba7425a..b8b044eee9 100644 --- a/tirx/tests/python/tirx/test_cuda_backend_config.py +++ b/tirx/tests/python/tirx/test_cuda_backend_config.py @@ -599,7 +599,7 @@ def test_legacy_payloads(): current = old_binary + string(snapshot.encode()) assert load(current).inspect_source("backend_config") == snapshot for malformed in (old_binary + b"x", current[:-1], current + b"x"): - with pytest.raises(tvm_tirx.error.InternalError): + with pytest.raises(RuntimeError): load(malformed) diff --git a/tirx/tests/python/tirx/tile/test_dispatcher.py b/tirx/tests/python/tirx/tile/test_dispatcher.py index c4c0ce11d6..eeb9000aec 100644 --- a/tirx/tests/python/tirx/tile/test_dispatcher.py +++ b/tirx/tests/python/tirx/tile/test_dispatcher.py @@ -158,7 +158,7 @@ def test_residual_tile_call_is_rejected(): call.op.name, "FTileOpDispatch", lambda call, context: Function([], Evaluate(0)) ) # Dispatch expands statement-boundary Calls; other positions must be rejected. - with pytest.raises(tvm_tirx.error.InternalError, match="Failed to lower the TIRx program"): + with pytest.raises(RuntimeError, match="Failed to lower the TIRx program"): _lower(Return(call)) diff --git a/tirx/tests/python/tirx/transform/test_lower_tirx.py b/tirx/tests/python/tirx/transform/test_lower_tirx.py index 51335b0bab..d27649f222 100644 --- a/tirx/tests/python/tirx/transform/test_lower_tirx.py +++ b/tirx/tests/python/tirx/transform/test_lower_tirx.py @@ -897,7 +897,7 @@ def test_lower_layout(): @pytest.mark.parametrize( "target,error,message", [ - (None, tvm_tirx.error.InternalError, "Target context required"), + (None, RuntimeError, "Target context required"), ("cuda", RuntimeError, "An explicit layout is required for shared tensors"), ], ids=["missing-target", "missing-shared-layout"], diff --git a/tirx/tests/python/tirx/transform/test_make_packed_api.py b/tirx/tests/python/tirx/transform/test_make_packed_api.py index 19367b4014..b7d807a8f7 100644 --- a/tirx/tests/python/tirx/transform/test_make_packed_api.py +++ b/tirx/tests/python/tirx/transform/test_make_packed_api.py @@ -692,9 +692,7 @@ def test_return_from_parallel_scope_is_rejected(): for i in T.parallel(2): return i - with pytest.raises( - tvm_tirx.error.InternalError, match="Return cannot be used in parallel scope" - ): + with pytest.raises(RuntimeError, match="Return cannot be used in parallel scope"): tvm_tirx.tirx.transform.MakePackedAPI()(Before) diff --git a/tirx/tests/python/tirx/transform/test_split_host_device.py b/tirx/tests/python/tirx/transform/test_split_host_device.py index c08edbc037..0fb2059989 100644 --- a/tirx/tests/python/tirx/transform/test_split_host_device.py +++ b/tirx/tests/python/tirx/transform/test_split_host_device.py @@ -169,7 +169,7 @@ def test_device_kernel_nonzero_return_is_rejected(): ) return 1 - with pytest.raises(tvm_tirx.error.InternalError, match="successful return"): + with pytest.raises(RuntimeError, match="successful return"): tvm_tirx.tirx.transform.SplitHostDevice()(Before) diff --git a/tirx/tests/python/tirx/transform/test_vectorize.py b/tirx/tests/python/tirx/transform/test_vectorize.py index d1f42b8b77..7093fbe203 100644 --- a/tirx/tests/python/tirx/transform/test_vectorize.py +++ b/tirx/tests/python/tirx/transform/test_vectorize.py @@ -18,7 +18,7 @@ import pytest import tvm_tirx import tvm_tirx.testing -from tvm_tirx.codegen import llvm_version_major +from tvm_tirx.codegen import llvm from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T @@ -84,7 +84,7 @@ def test_vectorize_vector_scalable_error(): A[T.ramp(j * 4, 1, 4)] = T.Broadcast(T.float32(1), 4) error_msg = "Creating scalable vectors from existing vectors is not supported." - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): tvm_tirx.tirx.transform.VectorizeLoop()(Module) @@ -97,7 +97,7 @@ def test_vectorize_vector_scalable_error2(): A[j] = T.Broadcast(T.float32(1), T.vscale() * 4) error_msg = "Vectorizing over scalable tensor elements is not supported in vectorizer." - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): tvm_tirx.tirx.transform.VectorizeLoop()(Module) @@ -113,7 +113,7 @@ def test_vectorize_vector_scalable_error3(): ) error_msg = "Vectorizing over existing scalable vectors is not supported." - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): tvm_tirx.tirx.transform.VectorizeLoop()(Module) @@ -129,7 +129,7 @@ def test_vectorize_vector_scalable_error4(): ) error_msg = "Creating scalable vectors from existing vectors is not supported." - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): tvm_tirx.tirx.transform.VectorizeLoop()(Module) @@ -509,7 +509,7 @@ def test_illegal_extent(): A[j] = 3 error_msg = r"Failed to vectorize loop with extent n for target None" - with pytest.raises(tvm_tirx.error.InternalError, match=error_msg): + with pytest.raises(RuntimeError, match=error_msg): tvm_tirx.tirx.transform.VectorizeLoop()(Mod) @@ -523,7 +523,7 @@ def test_illegal_vscale_in_non_sve_compilation(): A[j] = 13 msg = "Failed to vectorize loop with extent T.vscale\\(\\) \\* 4 for target" - with pytest.raises(tvm_tirx.error.InternalError, match=msg): + with pytest.raises(RuntimeError, match=msg): tvm_tirx.tirx.transform.VectorizeLoop()(Mod) @@ -877,7 +877,7 @@ def test_vectorize_llvm_pure_intrin_fail(extent, vec_str, target): # LLVM 20 added vector support for llvm.lround/llvm.llround. The IR Verifier's # "Intrinsic does not support vectors" check was removed in release/20.x, so # compilation only fails on LLVM <= 19. - if llvm_version_major() >= 20: + if llvm.version_major() >= 20: tvm_tirx.tirx.compile(mod, target=target) else: with pytest.raises(Exception, match="Intrinsic does not support vectors"):
