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 4e909f70810ece7949121628e777309281fff7ef Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 12:44:21 2026 -0400 [REFACTOR] Separate TIRX test helpers from parent integrations (#20792) Separate parent TVM integration and build-reporting helpers from TIRX runtime/device checks. Keep target-selection policy in parent TVM, move benchmark helpers into the TIRX testing package, replace parameter wrappers with native pytest fixtures, remove slow-test helpers, and give TIRX-owned environment controls explicit names while preserving shared runtime settings. Keep Vulkan/OpenCL compilation in the separate GPU-other build and exclude those backends from GPU test wheels. Preserve native Vulkan bool storage and align existing backend assertions with the current code generation and device-attribute APIs. --- tirx/python/tvm_tirx/__init__.py | 10 +- tirx/python/tvm_tirx/backend/_autoload_backends.py | 4 +- tirx/python/tvm_tirx/backend/cuda/iket.py | 2 +- tirx/python/tvm_tirx/backend/hexagon/linker.py | 15 ++ tirx/python/tvm_tirx/base.py | 4 +- tirx/python/tvm_tirx/script/printer/highlight.py | 4 +- tirx/python/tvm_tirx/support/__init__.py | 39 +-- tirx/python/tvm_tirx/support/cc.py | 4 +- tirx/python/tvm_tirx/support/libinfo.py | 44 ---- tirx/python/tvm_tirx/support/ndk.py | 22 +- tirx/python/tvm_tirx/testing/__init__.py | 2 +- tirx/python/tvm_tirx/testing/_pytest.py | 2 +- tirx/python/tvm_tirx/{tirx => testing}/bench.py | 0 tirx/python/tvm_tirx/testing/env.py | 244 ++----------------- tirx/python/tvm_tirx/testing/utils.py | 197 +-------------- tirx/python/tvm_tirx/tirx/__init__.py | 2 +- tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 2 +- tirx/src/backend/cuda/codegen/codegen_cuda.cc | 2 +- .../backend/cuda/codegen/cuda_fallback_module.h | 4 +- tirx/src/backend/cuda/transform/lower_iket.cc | 2 +- .../hexagon/codegen/hexagon_fallback_module.h | 4 +- .../backend/metal/codegen/metal_fallback_module.h | 4 +- .../opencl/codegen/opencl_fallback_module.h | 4 +- .../backend/rocm/codegen/rocm_fallback_module.h | 4 +- tirx/src/backend/vulkan/codegen/codegen_spirv.cc | 14 +- tirx/src/backend/vulkan/codegen/spirv_utils.cc | 4 +- .../vulkan/codegen/vulkan_fallback_module.h | 4 +- .../webgpu/codegen/webgpu_fallback_module.h | 4 +- .../python/backend/cuda/codegen/test_basic.py | 2 +- tirx/tests/python/backend/cuda/codegen/test_fp8.py | 13 +- .../python/backend/cuda/test_module_export.py | 4 +- .../python/backend/hexagon/codegen/test_basic.py | 8 +- .../python/backend/metal/test_module_export.py | 4 +- .../python/backend/opencl/codegen/test_basic.py | 4 +- .../python/backend/opencl/test_module_export.py | 4 +- .../python/backend/rocm/test_module_export.py | 4 +- .../python/backend/vulkan/codegen/test_basic.py | 21 +- .../python/backend/vulkan/test_module_export.py | 4 +- .../python/codegen/gpu_common/test_allreduce.py | 8 +- .../python/codegen/test_argument_validation.py | 6 +- tirx/tests/python/codegen/test_assert.py | 5 +- tirx/tests/python/codegen/test_nvshmem.py | 2 +- tirx/tests/python/runtime/test_memory_scope.py | 6 +- tirx/tests/python/runtime/test_module.py | 2 +- tirx/tests/python/sym/test_const_int_bound.py | 271 ++++++++++++++------- tirx/tests/python/sym/test_iter_affine_map.py | 19 +- tirx/tests/python/sym/test_rewrite_simplify.py | 237 ++++++++++++++---- tirx/tests/python/sym/test_simplify.py | 7 +- tirx/tests/python/target/test_target.py | 11 +- tirx/tests/python/tirx/base/test_index_map.py | 201 +++++++-------- tirx/tests/python/tirx/base/test_intrin.py | 4 +- tirx/tests/python/tirx/codegen/test_cuda.py | 4 +- tirx/tests/python/tirx/iket/test_orchestration.py | 6 +- tirx/tests/python/tirx/iket/test_profiler.py | 10 +- tirx/tests/python/tirx/test_axis_enum.py | 2 +- tirx/tests/python/tirx/test_bench_utils.py | 8 +- .../tests/python/tirx/test_compilation_pipeline.py | 4 +- tirx/tests/python/tirx/test_cuda_backend_config.py | 14 +- tirx/tests/python/tirx/test_thread_scope.py | 2 +- tirx/tests/python/tirx/test_tile_namespace.py | 2 +- .../python/tirx/transform/test_fp8_legalize.py | 12 +- tirx/tests/python/tirx/utils.py | 16 -- 62 files changed, 702 insertions(+), 872 deletions(-) diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index 9fcdf53125..d881a33586 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -76,7 +76,7 @@ if TYPE_CHECKING or not _is_config_mode(): # support infra from . import support - # Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup + # Side-effect imports: register CUDA/ROCm FFI callbacks at TIRX startup from .support import rocm as _rocm, nvcc as _nvcc def tvm_wrap_excepthook(exception_hook): @@ -97,16 +97,16 @@ if TYPE_CHECKING or not _is_config_mode(): # Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt # into being loaded automatically at ``import tvm_tirx`` time by declaring an entry - # point in the ``tvm.backends`` group: - # [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload". - # Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``. + # point in the ``tvm_tirx.backends`` group: + # [project.entry-points."tvm_tirx.backends"] tvm_foo = "tvm_foo:_autoload". + # Autoload can be disabled via ``TVM_TIRX_DEVICE_BACKEND_AUTOLOAD=0``. from .backend._autoload_backends import _autoload_backends _autoload_backends() # Load the optional ORC JIT extension for its global FFI registrations. LLVM # function lookup reports a targeted installation error if this import fails; - # importing TVM itself remains silent for AOT-only workflows. + # importing TIRX itself remains silent for AOT-only workflows. try: import tvm_ffi_orcjit as _tvm_ffi_orcjit except Exception: # pylint: disable=broad-exception-caught diff --git a/tirx/python/tvm_tirx/backend/_autoload_backends.py b/tirx/python/tvm_tirx/backend/_autoload_backends.py index 4e225f09a0..ad333064e9 100644 --- a/tirx/python/tvm_tirx/backend/_autoload_backends.py +++ b/tirx/python/tvm_tirx/backend/_autoload_backends.py @@ -54,7 +54,7 @@ def _autoload_backends() -> None: return _AUTO_LOAD_DONE = True - if os.environ.get("TVM_DEVICE_BACKEND_AUTOLOAD", "1") == "0": + if os.environ.get("TVM_TIRX_DEVICE_BACKEND_AUTOLOAD", "1") == "0": return from tvm_tirx import _RUNTIME_ONLY # pylint: disable=import-outside-toplevel @@ -65,7 +65,7 @@ def _autoload_backends() -> None: # Out-of-tree extensions opt into being loaded automatically at ``import tvm_tirx`` time # by declaring an entry point in the ``tvm_tirx.backends`` group: # [project.entry-points."tvm_tirx.backends"] tvm_foo = "tvm_foo:_autoload". - # Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``. + # Autoload can be disabled via ``TVM_TIRX_DEVICE_BACKEND_AUTOLOAD=0``. for entry_pt in entry_points(group="tvm_tirx.backends"): try: entry_pt.load()() diff --git a/tirx/python/tvm_tirx/backend/cuda/iket.py b/tirx/python/tvm_tirx/backend/cuda/iket.py index d3e520f6b0..008bddaf2a 100644 --- a/tirx/python/tvm_tirx/backend/cuda/iket.py +++ b/tirx/python/tvm_tirx/backend/cuda/iket.py @@ -49,7 +49,7 @@ from typing import Any import tvm_tirx from tvm_tirx.script import tirx as T -_INJECTED_CHILD_ENABLE_ENV = "TVM_IKET_INJECTED_CHILD_ENABLE" +_INJECTED_CHILD_ENABLE_ENV = "TVM_TIRX_IKET_INJECTED_CHILD_ENABLE" _POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all")) _INJECTION_CONFIG_ENV_VARS = ("IKET_INJECTION_CONFIG", "SMODEL_INJECTION_CONFIG") _INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", *_INJECTION_CONFIG_ENV_VARS) diff --git a/tirx/python/tvm_tirx/backend/hexagon/linker.py b/tirx/python/tvm_tirx/backend/hexagon/linker.py index 1903c17b74..892197d751 100644 --- a/tirx/python/tvm_tirx/backend/hexagon/linker.py +++ b/tirx/python/tvm_tirx/backend/hexagon/linker.py @@ -16,11 +16,13 @@ # under the License. """Hexagon shared-library linking and custom linker registration.""" +import functools import os from pathlib import Path import tvm_ffi +import tvm_tirx from tvm_tirx.ir import StringImm from tvm_tirx.support import cc @@ -110,3 +112,16 @@ def link_shared(so_name, objs, extra_args=None): cc=linker, ) return 0 + + [email protected] +def has_hexagon_toolchain() -> bool: + """True if the Hexagon toolchain is available for compilation.""" + try: + return ( + tvm_tirx.runtime.enabled("llvm") + and "hexagon" in tvm_tirx.codegen.llvm.get_targets() + and "HEXAGON_TOOLCHAIN" in os.environ + ) + except Exception: # pylint: disable=broad-except + return False diff --git a/tirx/python/tvm_tirx/base.py b/tirx/python/tvm_tirx/base.py index 17c4a05b7b..af8665125d 100644 --- a/tirx/python/tvm_tirx/base.py +++ b/tirx/python/tvm_tirx/base.py @@ -31,8 +31,8 @@ from . import libinfo # ---------------------------- # Whether only the runtime library is loaded (runtime-only wheel, or -# ``TVM_USE_RUNTIME_LIB=1``). Set during library loading below. -_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1" +# ``TVM_TIRX_USE_RUNTIME_LIB=1``). Set during library loading below. +_RUNTIME_ONLY = os.environ.get("TVM_TIRX_USE_RUNTIME_LIB") == "1" # Handles of the core libraries actually loaded, keyed by basename # (e.g. ``{"tvm_runtime": <CDLL>, "tvm_tirx_compiler": <CDLL>}``). Downstream / diff --git a/tirx/python/tvm_tirx/script/printer/highlight.py b/tirx/python/tvm_tirx/script/printer/highlight.py index 603b949793..bf1dd5d74f 100644 --- a/tirx/python/tvm_tirx/script/printer/highlight.py +++ b/tirx/python/tvm_tirx/script/printer/highlight.py @@ -51,7 +51,7 @@ def cprint( https://pygments.org/styles/ The default pygmentize style can also be set with the environment - variable "TVM_PYGMENTIZE_STYLE". + variable "TVM_TIRX_PYGMENTIZE_STYLE". """ if hasattr(printable, "script") and callable(getattr(printable, "script")): printable = printable.script() @@ -200,7 +200,7 @@ def _get_pygments_style( if style is not None: return style - style_from_environment = os.environ.get("TVM_PYGMENTIZE_STYLE", "").strip() + style_from_environment = os.environ.get("TVM_TIRX_PYGMENTIZE_STYLE", "").strip() if style_from_environment: return style_from_environment diff --git a/tirx/python/tvm_tirx/support/__init__.py b/tirx/python/tvm_tirx/support/__init__.py index abc46b92f2..6a6fc61664 100644 --- a/tirx/python/tvm_tirx/support/__init__.py +++ b/tirx/python/tvm_tirx/support/__init__.py @@ -14,41 +14,4 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""tvm_tirx.support — Python helpers that integrate TVM with external CLIs -and host-side tools (compilers, archivers, subprocess pools, build-info -queries). Distinct from `tvm.contrib`, which is reserved for optional -vendor SDK integrations and experimental features.""" - -import json -import os -import platform -import sys -import textwrap - -import tvm_ffi - -import tvm_tirx - -tvm_ffi.init_ffi_api("support", __name__) - -from .libinfo import libinfo - - -def describe(): - """ - Print out information about TVM and the current Python environment - """ - info = list((k, v) for k, v in libinfo().items()) - info = dict(sorted(info, key=lambda x: x[0])) - print("Python Environment") - sys_version = sys.version.replace("\n", " ") - uname = platform.uname() - uname = f"{uname.system} {uname.release} {uname.version} {uname.machine}" - lines = [ - f"TVM version = {tvm_tirx.__version__}", - f"Python version = {sys_version} ({sys.maxsize.bit_length() + 1} bit)", - f"uname = {uname}", - ] - print(textwrap.indent("\n".join(lines), prefix=" ")) - print("CMake Options:") - print(textwrap.indent(json.dumps(info, indent=2), prefix=" ")) +"""Host-side compilers, archivers, and subprocess helpers for TIRX.""" diff --git a/tirx/python/tvm_tirx/support/cc.py b/tirx/python/tvm_tirx/support/cc.py index 5080bc308c..c3c54fe3d0 100644 --- a/tirx/python/tvm_tirx/support/cc.py +++ b/tirx/python/tvm_tirx/support/cc.py @@ -374,8 +374,8 @@ def _linux_compile( def _windows_compile(output, objects, options, cwd=None, ccache_env=None): - compiler = os.getenv("TVM_WIN_CC", default="clang") - win_target = os.getenv("TVM_WIN_TARGET", default="x86_64") + compiler = os.getenv("TVM_TIRX_WIN_CC", default="clang") + win_target = os.getenv("TVM_TIRX_WIN_TARGET", default="x86_64") cmd = [compiler] cmd += ["-O2"] cmd += ["--target=" + win_target] diff --git a/tirx/python/tvm_tirx/support/libinfo.py b/tirx/python/tvm_tirx/support/libinfo.py deleted file mode 100644 index d191ef62e5..0000000000 --- a/tirx/python/tvm_tirx/support/libinfo.py +++ /dev/null @@ -1,44 +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. -"""Build-info query helpers for tvm_tirx.support.""" - -import os - - -def libinfo(): - """Returns a dictionary of compile-time info — minimal Python fallback. - - The native ``support.GetLibInfo`` global function is no longer registered - after the upstream sync, so we synthesize the values from build-time hints - instead. - """ - return { - "USE_CUDA": os.environ.get("TVM_USE_CUDA", "ON"), - "USE_LLVM": os.environ.get("TVM_USE_LLVM", "ON"), - "USE_NCCL": os.environ.get("TVM_USE_NCCL", "ON"), - "USE_NVTX": os.environ.get("TVM_USE_NVTX", "ON"), - "USE_NVSHMEM": os.environ.get("TVM_USE_NVSHMEM", "OFF"), - "USE_CUDNN": "OFF", - "USE_CUTLASS": "OFF", - "USE_VULKAN": "OFF", - "USE_OPENCL": "OFF", - "USE_METAL": "OFF", - "USE_ROCM": "OFF", - "USE_CLML": "OFF", - "USE_NNAPI_RUNTIME": "OFF", - "USE_NNAPI_CODEGEN": "OFF", - } diff --git a/tirx/python/tvm_tirx/support/ndk.py b/tirx/python/tvm_tirx/support/ndk.py index b525abef2a..188401fa67 100644 --- a/tirx/python/tvm_tirx/support/ndk.py +++ b/tirx/python/tvm_tirx/support/ndk.py @@ -45,11 +45,11 @@ def create_shared(output, objects, options=None): options : list of str, optional The additional options. """ - if "TVM_NDK_CC" not in os.environ: + if "TVM_TIRX_NDK_CC" not in os.environ: raise RuntimeError( - "Require environment variable TVM_NDK_CC to be the NDK standalone compiler" + "Require environment variable TVM_TIRX_NDK_CC to be the NDK standalone compiler" ) - compiler = os.environ["TVM_NDK_CC"] + compiler = os.environ["TVM_TIRX_NDK_CC"] cmd = [compiler] cmd += ["-o", output] @@ -73,7 +73,9 @@ def create_shared(output, objects, options=None): # assign output format create_shared.output_format = "so" create_shared.get_target_triple = ( - get_target_by_dump_machine(os.environ["TVM_NDK_CC"]) if "TVM_NDK_CC" in os.environ else None + get_target_by_dump_machine(os.environ["TVM_TIRX_NDK_CC"]) + if "TVM_TIRX_NDK_CC" in os.environ + else None ) @@ -88,9 +90,9 @@ def create_staticlib(output, inputs): inputs : list List of object files or tar files """ - if "TVM_NDK_CC" not in os.environ: + if "TVM_TIRX_NDK_CC" not in os.environ: raise RuntimeError( - "Require environment variable TVM_NDK_CC to be the NDK standalone compiler" + "Require environment variable TVM_TIRX_NDK_CC to be the NDK standalone compiler" ) output_name = os.path.basename(output) @@ -98,7 +100,7 @@ def create_staticlib(output, inputs): tmp_output = temp.relpath("lib" + output_name) objects = _tar.normalize_file_list_by_unpacking_tars(temp, inputs) - compiler = os.environ["TVM_NDK_CC"] + compiler = os.environ["TVM_TIRX_NDK_CC"] base_path = os.path.dirname(compiler) ar_path = os.path.join(base_path, "llvm-ar") cmd = [ar_path] @@ -146,12 +148,12 @@ def get_global_symbol_section_map(path, *, nm=None) -> dict[str, str]: symbol_section_map: Dict[str, str] A map from defined global symbol to their sections """ - if "TVM_NDK_CC" not in os.environ: + if "TVM_TIRX_NDK_CC" not in os.environ: raise RuntimeError( - "Require environment variable TVM_NDK_CC to be the NDK standalone compiler" + "Require environment variable TVM_TIRX_NDK_CC to be the NDK standalone compiler" ) if nm is None: - compiler = os.environ["TVM_NDK_CC"] + compiler = os.environ["TVM_TIRX_NDK_CC"] base_path = os.path.dirname(compiler) nm = os.path.join(base_path, "llvm-nm") return _cc.get_global_symbol_section_map(path, nm=nm) diff --git a/tirx/python/tvm_tirx/testing/__init__.py b/tirx/python/tvm_tirx/testing/__init__.py index 11936b2e6d..448d5a2a0e 100644 --- a/tirx/python/tvm_tirx/testing/__init__.py +++ b/tirx/python/tvm_tirx/testing/__init__.py @@ -17,7 +17,7 @@ # under the License. # pylint: disable=redefined-builtin, wildcard-import -"""Utility Python functions for TVM testing""" +"""TIRX test utilities and local capability probes.""" from ._ffi_api import ( echo, diff --git a/tirx/python/tvm_tirx/testing/_pytest.py b/tirx/python/tvm_tirx/testing/_pytest.py index 89bd9497ff..5e805a2fdc 100644 --- a/tirx/python/tvm_tirx/testing/_pytest.py +++ b/tirx/python/tvm_tirx/testing/_pytest.py @@ -118,7 +118,7 @@ def pytest_runtest_setup(item): def _source_codegen_without_cuda(monkeypatch): # CPU-only source checks must not depend on a CUDA compiler being installed. if not env.has_cuda(): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") @pytest.fixture(autouse=True) diff --git a/tirx/python/tvm_tirx/tirx/bench.py b/tirx/python/tvm_tirx/testing/bench.py similarity index 100% rename from tirx/python/tvm_tirx/tirx/bench.py rename to tirx/python/tvm_tirx/testing/bench.py diff --git a/tirx/python/tvm_tirx/testing/env.py b/tirx/python/tvm_tirx/testing/env.py index 2200d0817f..fbba72148e 100644 --- a/tirx/python/tvm_tirx/testing/env.py +++ b/tirx/python/tvm_tirx/testing/env.py @@ -14,74 +14,32 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Thin capability probes for test gating. - -This module exposes small ``has_*`` predicates that report whether the -current environment can run a given feature. They are meant to be used -with plain pytest markers and ``skipif``:: - - import pytest - import tvm_tirx.testing - - @pytest.mark.gpu - @pytest.mark.skipif(not tvm_tirx.testing.env.has_cuda(), reason="need cuda") - def test_my_cuda_kernel(): - ... - -Each probe's expensive query (device lookup, ``nvcc`` subprocess, libinfo) -is memoized with :func:`functools.cache`, so it runs at most once per -process even though ``skipif`` evaluates the predicate at import time for -every decorated test. Probes never raise: when support is absent they -return ``False`` (or a zero version tuple) rather than propagating an -error out of collection. - -Three kinds of probe live here: - -* **runtime device** probes (``has_cuda``, ``has_gpu`` …) ask whether a - usable device of a given kind is present; -* **build-support** probes (``has_cudnn`` …, ``build_flag_enabled`` …) ask whether - an optional library was compiled into the runtime; -* **version / capability** probes (``has_cuda_compute``, - ``has_nvcc_version`` …) ask about a finer capability of a present device - or toolchain. -""" +"""Local runtime, device, and compiler capability probes.""" import functools -import os import tvm_ffi import tvm_tirx __all__ = [ - "build_flag_enabled", "cuda_arch", - "has_adreno_opencl", # cpu features "has_cpu_feature", - "has_cublas", # runtime device "has_cuda", "has_cuda_arch", # version / capability "has_cuda_compute", "has_cudagraph", - # build support - "has_cudnn", "has_gpu", - # toolchain / environment - "has_hexagon", - "has_hexagon_toolchain", - "has_hipblas", "has_llvm", "has_llvm_min_version", "has_matrixcore", "has_metal", "has_multi_gpu", - "has_nccl", "has_nvcc_version", "has_nvptx", - "has_nvshmem", "has_opencl", "has_rocm", "has_vulkan", @@ -100,89 +58,6 @@ def _device_exists(kind: str, index: int = 0) -> bool: return False [email protected] -def build_flag_enabled(flag: str) -> bool: - """Return whether an optional build flag (e.g. ``USE_CUTLASS``) is on. - - A flag counts as enabled unless it is explicitly disabled, so library - flags carrying a path (rather than a boolean) still register as present. - Callers gate via ``@pytest.mark.skipif(not env.build_flag_enabled("USE_X"), ...)``. - """ - try: - value = tvm_tirx.support.libinfo().get(flag, "OFF") - return str(value).lower() not in ("off", "false", "0") - except Exception: # pylint: disable=broad-except - return False - - -DEFAULT_TEST_TARGETS = [ - "llvm", - "cuda", - "nvptx", - {"kind": "vulkan", "from_device": 0}, - "opencl", - {"kind": "opencl", "device": "mali"}, - {"kind": "opencl", "device": "intel_graphics"}, - "metal", - "rocm", - "hexagon", -] - - -def _parse_target_entry(entry): - """Parse a target entry from TVM_TEST_TARGETS env var. - - Entries can be plain kind names (e.g. "llvm") or JSON dicts - (e.g. '{"kind": "opencl", "device": "mali"}'). - """ - entry = entry.strip() - if entry.startswith("{"): - import json # pylint: disable=import-outside-toplevel - - return json.loads(entry) - return entry - - -def _tvm_test_targets(): - target_str = os.environ.get("TVM_TEST_TARGETS", "").strip() - if target_str: - # De-duplicate while preserving order. dict items can't be hashed - # directly, so use their str() form as the dedup key. - targets = [] - seen = set() - for t in target_str.split(";"): - t = t.strip() - if not t: - continue - parsed = _parse_target_entry(t) - key = str(parsed) - if key in seen: - continue - seen.add(key) - targets.append(parsed) - return targets - - return DEFAULT_TEST_TARGETS - - [email protected] -def _target_enabled(kind: str) -> bool: - """True if ``kind`` is selected by ``TVM_TEST_TARGETS`` (or the default set). - - Honors the ``TVM_TEST_TARGETS`` opt-out, so CI can exclude a flaky - backend (e.g. opencl) via ``TVM_TEST_TARGETS`` and have its tests skip - even when a device is physically present. - """ - try: - for target in _tvm_test_targets(): - k = target["kind"] if isinstance(target, dict) else str(target).split()[0] - if k == kind: - return True - return False - except Exception: # pylint: disable=broad-except - return True # fail open: the device check still gates - - @functools.cache def _runtime_enabled(kind: str) -> bool: """True if the runtime was built with support for target ``kind``. @@ -197,55 +72,46 @@ def _runtime_enabled(kind: str) -> bool: return False -def _device_usable(kind: str) -> bool: - """True if ``kind`` is enabled for this run and a ``kind`` device exists. - - The TVM_TEST_TARGETS opt-out is checked first so that an excluded backend - never probes a (possibly crashy) device. - """ - return _target_enabled(kind) and _device_exists(kind) - - # --- runtime device probes ------------------------------------------------- def has_cuda() -> bool: - """True if a CUDA device is present and enabled in TVM_TEST_TARGETS.""" - return _device_usable("cuda") + """True if a CUDA device is present.""" + return _device_exists("cuda") def has_rocm() -> bool: - """True if a ROCm device is present and enabled in TVM_TEST_TARGETS.""" - return _device_usable("rocm") + """True if a ROCm device is present.""" + return _device_exists("rocm") def has_vulkan() -> bool: - """True if a Vulkan device is present and enabled in TVM_TEST_TARGETS.""" - return _device_usable("vulkan") + """True if a Vulkan device is present.""" + return _device_exists("vulkan") def has_metal() -> bool: - """True if a Metal device is present and enabled in TVM_TEST_TARGETS.""" - return _device_usable("metal") + """True if a Metal device is present.""" + return _device_exists("metal") def has_opencl() -> bool: - """True if an OpenCL device is present and enabled in TVM_TEST_TARGETS.""" - return _device_usable("opencl") + """True if an OpenCL device is present.""" + return _device_exists("opencl") def has_nvptx() -> bool: """True if NVPTX is usable: a (CUDA) device, plus the LLVM backend it needs.""" - return _device_usable("nvptx") and has_llvm() + return _device_exists("nvptx") and has_llvm() def has_llvm() -> bool: - """True if the LLVM backend was built in and enabled in TVM_TEST_TARGETS. + """True if the LLVM backend was built in. Uses ``tvm_tirx.runtime.enabled`` rather than device existence: ``llvm`` maps to the CPU device, which exists even on a ``USE_LLVM=OFF`` build. """ - return _target_enabled("llvm") and _runtime_enabled("llvm") + return _runtime_enabled("llvm") def has_gpu() -> bool: @@ -268,53 +134,6 @@ def has_multi_gpu(count: int = 2) -> bool: return False -# --- build-support probes -------------------------------------------------- -# -# These wrap the optional-library build flags. Features that extend CUDA / -# ROCm additionally require the parent device to be present. - - -def has_cudnn() -> bool: - """True if cuDNN was built in and a CUDA device is present.""" - return has_cuda() and build_flag_enabled("USE_CUDNN") - - -def has_cublas() -> bool: - """True if cuBLAS was built in and a CUDA device is present.""" - return has_cuda() and build_flag_enabled("USE_CUBLAS") - - -def has_nccl() -> bool: - """True if NCCL was built in and a CUDA device is present.""" - try: - return has_cuda() and ( - tvm_ffi.get_global_func("runtime.disco.nccl.init_ccl", allow_missing=True) is not None - ) - except Exception: # pylint: disable=broad-except - return False - - -def has_hipblas() -> bool: - """True if hipBLAS was built in and a ROCm device is present.""" - return has_rocm() and build_flag_enabled("USE_HIPBLAS") - - [email protected] -def has_nvshmem() -> bool: - """True if the disco NVSHMEM runtime is available (requires CUDA). - - Probes the runtime global function rather than the ``USE_NVSHMEM`` build - flag, since the flag can be set in builds that do not ship the runtime. - """ - try: - return has_cuda() and ( - tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid", allow_missing=True) - is not None - ) - except Exception: # pylint: disable=broad-except - return False - - # --- version / capability probes ------------------------------------------- @@ -444,41 +263,6 @@ def has_cudagraph() -> bool: return False -# --- toolchain / environment probes ---------------------------------------- - - [email protected] -def has_hexagon_toolchain() -> bool: - """True if the Hexagon toolchain is available for compilation.""" - try: - return ( - _runtime_enabled("llvm") - and "hexagon" in tvm_tirx.codegen.llvm.get_targets() - and "HEXAGON_TOOLCHAIN" in os.environ - ) - except Exception: # pylint: disable=broad-except - return False - - [email protected] -def has_hexagon() -> bool: - """True if Hexagon can both compile and run (toolchain + attached device).""" - try: - return ( - _runtime_enabled("hexagon") - and has_hexagon_toolchain() - and "ANDROID_SERIAL_NUMBER" in os.environ - ) - except Exception: # pylint: disable=broad-except - return False - - [email protected] -def has_adreno_opencl() -> bool: - """True if remote Adreno OpenCL testing is configured (RPC_TARGET set).""" - return build_flag_enabled("USE_OPENCL") and os.environ.get("RPC_TARGET") is not None - - # --- cpu feature probes ---------------------------------------------------- diff --git a/tirx/python/tvm_tirx/testing/utils.py b/tirx/python/tvm_tirx/testing/utils.py index 1ddc894ce3..0a79ef9692 100644 --- a/tirx/python/tvm_tirx/testing/utils.py +++ b/tirx/python/tvm_tirx/testing/utils.py @@ -16,69 +16,20 @@ # under the License. # pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name -"""TVM testing utilities - -Organization -************ - -This file contains functions expected to be called directly by a user -while writing unit tests. Integrations with the pytest framework -for TVM's own test suite are in ``tests/python/conftest.py``. - -Testing Markers -*************** - -We use pytest markers to specify the requirements of test functions. -Currently there is a single distinction that matters for our testing -environment: does the test require a gpu. Tests that require a gpu are -tagged with the ``gpu`` pytest marker -- the only registered marker (see -the ``markers`` entry in ``pyproject.toml``). This lets us select the -gpu subset of tests with ``pytest -m gpu`` (and exclude them on cpu-only -nodes with ``pytest -m "not gpu"``). - -The ``gpu`` marker only controls which testing node a test runs on; it -does not check whether the required hardware or libraries are actually -present. To gate a test on a specific capability, combine the marker -with a ``skipif`` that consults the memoized environment probes in -:py:mod:`tvm_tirx.testing.env`: - -.. code-block:: python - - @pytest.mark.gpu - @pytest.mark.skipif(not tvm_tirx.testing.env.has_cuda(), reason="need cuda") - def test_cuda_vectorize_add(): - ... - -There is one ``has_*`` (or ``is_*``) probe per capability -- for example -:py:func:`tvm_tirx.testing.env.has_gpu`, :py:func:`tvm_tirx.testing.env.has_cuda`, -and :py:func:`tvm_tirx.testing.env.has_vulkan`. For optional Python packages, -prefer ``pytest.importorskip("pkg_name")`` instead of a ``skipif``. - -To run a test against a variety of targets, parametrize over ``target`` with -``@pytest.mark.parametrize("target", [...])`` -- tag GPU targets with -``pytest.mark.gpu`` so the CI routes them to GPU nodes, and skip an unavailable -target with ``pytest.mark.skipif(not tvm_tirx.testing.device_enabled(target))``. The -set of enabled targets is controlled by the ``TVM_TEST_TARGETS`` environment -variable, so the CI can run different targets on different testing nodes. - -""" +"""Numerical assertions, device availability, and TIRX test utilities.""" import inspect -import os import sys import time import ml_dtypes import numpy as np import pytest -import tvm_ffi import tvm_tirx from . import env -SKIP_SLOW_TESTS = os.getenv("SKIP_SLOW_TESTS", "").lower() in {"true", "1", "yes"} - def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True): """Version of np.testing.assert_allclose with `atol` and `rtol` fields set @@ -95,146 +46,20 @@ def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True): def device_enabled(target): - """Check if a target should be used when testing. - - Gate a device-specific test on this with - ``@pytest.mark.skipif(not tvm_tirx.testing.device_enabled(target))``. - - This allows the user to control which devices they are testing against. In - tests, this should be used to check if a device should be used when said - device is an optional part of the test. - - Parameters - ---------- - target : str or Dict[str, Any] or tvm_tirx.target.Target - Target string to check against - - Returns - ------- - bool - Whether or not the device associated with this target is enabled. - - Example - ------- - >>> @pytest.mark.gpu - >>> def test_mytest(): - >>> for target in ["cuda", "llvm"]: - >>> if device_enabled(target): - >>> test_body... - - Here, `test_body` will only be reached by with `target="cuda"` on gpu test - nodes and `target="llvm"` on cpu test nodes. - """ + """Return whether the target's runtime and local device are available.""" if isinstance(target, dict): - target_kind = target["kind"] + kind = target["kind"] elif hasattr(target, "kind"): - target_kind = target.kind.name + kind = target.kind.name else: assert isinstance(target, str), "device_enabled requires a target as a string" - # Target strings may include extra flags; only compare the kind. - target_kind = target.split(" ")[0] - # Check only explicitly selected candidates of the requested kind. In - # particular, an unavailable selection must not silently enable LLVM. - for candidate in env._tvm_test_targets(): - kind = candidate["kind"] if isinstance(candidate, dict) else candidate.split()[0] - if kind != target_kind: - continue - if kind == "hexagon": - return env._runtime_enabled("hexagon") and "ANDROID_SERIAL_NUMBER" in os.environ - if not env._runtime_enabled(kind): - continue - if kind == "cuda" and "cudnn" in tvm_tirx.target.Target(candidate).attrs.get("libs", []): - exists = tvm_ffi.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) - if env.has_cudnn() and exists is not None and exists(): - return True - elif tvm_tirx.device_from_target(candidate).exist: - return True - return False - - -slow = pytest.mark.skipif( - SKIP_SLOW_TESTS, - reason="Skipping slow test since the SKIP_SLOW_TESTS environment variable is 'true'", -) - - -def parameter(*values, ids=None, by_dict=None): - """Convenience function to define pytest parametrized fixtures. - - Declaring a variable using ``tvm_tirx.testing.parameter`` will define a - parametrized pytest fixture that can be used by test - functions. This is intended for cases that have no setup cost, - such as strings, integers, tuples, etc. For cases that have a - significant setup cost, use an ordinary pytest fixture - instead. - - If a test function accepts multiple parameters defined using - ``tvm_tirx.testing.parameter``, then the test will be run using every - combination of those parameters. - - The parameter definition applies to all tests in a module. If a - specific test should have different values for the parameter, that - test should be marked with ``@pytest.mark.parametrize``. - - Parameters - ---------- - values : Any - - A list of parameter values. A unit test that accepts this - parameter as an argument will be run once for each parameter - given. - - ids : List[str], optional - - A list of names for the parameters. If None, pytest will - generate a name from the value. These generated names may not - be readable/useful for composite types such as tuples. - - by_dict : Dict[str, Any] - - A mapping from parameter name to parameter value, to set both the - values and ids. - - Returns - ------- - function - A function output from pytest.fixture. - - Example - ------- - >>> size = tvm_tirx.testing.parameter(1, 10, 100) - >>> def test_using_size(size): - >>> ... # Test code here - - Or - - >>> shape = tvm_tirx.testing.parameter((5,10), (512,1024), ids=['small','large']) - >>> def test_using_size(shape): - >>> ... # Test code here - - Or - - >>> shape = tvm_tirx.testing.parameter(by_dict={'small': (5,10), 'large': (512,1024)}) - >>> def test_using_size(shape): - >>> ... # Test code here - - """ - - if by_dict is not None: - if values or ids: - raise RuntimeError( - "Use of the by_dict parameter cannot be used alongside positional arguments" - ) - - ids, values = zip(*by_dict.items()) - - # Optional cls parameter in case a parameter is defined inside a - # class scope. - @pytest.fixture(params=values, ids=ids, scope="session") - def as_fixture(*_cls, request): - return request.param - - return as_fixture + kind = target.split()[0] + if not env._runtime_enabled(kind): + return False + try: + return bool(tvm_tirx.device_from_target(target).exist) + except Exception: # pylint: disable=broad-except + return False def identity_after(x, sleep): diff --git a/tirx/python/tvm_tirx/tirx/__init__.py b/tirx/python/tvm_tirx/tirx/__init__.py index 79c23e4e88..b61906f327 100644 --- a/tirx/python/tvm_tirx/tirx/__init__.py +++ b/tirx/python/tvm_tirx/tirx/__init__.py @@ -85,7 +85,7 @@ from .layout import TileLayout, Layout, ComposeLayout from . import transform from . import analysis -# Compiler-only submodules. Skip under `TVM_USE_RUNTIME_LIB=1` since they +# Compiler-only submodules. Skip under `TVM_TIRX_USE_RUNTIME_LIB=1` since they # perform compiler-side FFI at module load (schema engine looks up # `ir.RegisterOp`; codegen registry hooks the build pipeline). from tvm_tirx.base import _RUNTIME_ONLY as _RUNTIME_ONLY_TIRX # pylint: disable=wrong-import-position diff --git a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py index bc50dcc2c6..0308c9bb0d 100644 --- a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py +++ b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py @@ -43,7 +43,7 @@ def tirx_pipeline(*, prepare_only=False, target=None): pass_ctx = tvm_tirx.ir.transform.PassContext.current() config = pass_ctx.config passes = [tirx.transform.TileDispatch()] - if "TVM_PRINT_AFTER_TIRX_DISPATCH_OPS" in os.environ: + if "TVM_TIRX_PRINT_AFTER_DISPATCH_OPS" in os.environ: passes.append(tvm_tirx.ir.transform.PrintIR()) passes += [ tirx.transform.LowerTensorLayouts(), diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda.cc b/tirx/src/backend/cuda/codegen/codegen_cuda.cc index 22cf57dfff..9308200f7f 100644 --- a/tirx/src/backend/cuda/codegen/codegen_cuda.cc +++ b/tirx/src/backend/cuda/codegen/codegen_cuda.cc @@ -1646,7 +1646,7 @@ void CodeGenCUDA::PrintVecElemLoadExpr(const PrimType& t, int i, const std::stri // Python postproc hook rewrite it, and hands the source bytes off to the // fallback-aware module factory. The factory may JIT to PTX/cubin via // `tvm_callback_cuda_compile` when -// USE_CUDA=ON; on USE_CUDA=OFF builds (or when TVM_COMPILE_FORCE_FALLBACK is +// USE_CUDA=ON; on USE_CUDA=OFF builds (or when TVM_TIRX_COMPILE_FORCE_FALLBACK is // set), it returns a `CUDAFallbackModuleObj` carrying the raw source for // later cross-compile. ffi::Module BuildCUDA(IRModule mod, Target target) { diff --git a/tirx/src/backend/cuda/codegen/cuda_fallback_module.h b/tirx/src/backend/cuda/codegen/cuda_fallback_module.h index c8dbb0c3c7..89ea39964f 100644 --- a/tirx/src/backend/cuda/codegen/cuda_fallback_module.h +++ b/tirx/src/backend/cuda/codegen/cuda_fallback_module.h @@ -60,14 +60,14 @@ ffi::Module CUDAFallbackModuleCreate(ffi::Bytes code, ffi::String fmt, * * - Registry hit (USE_CUDA=ON build) → real `CUDAModuleObj` returned. * - Registry miss (USE_CUDA=OFF build) → `CUDAFallbackModuleObj` returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). */ inline ffi::Module CUDAModuleCreateWithFallback(ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return CUDAFallbackModuleCreate(std::move(code), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/cuda/transform/lower_iket.cc b/tirx/src/backend/cuda/transform/lower_iket.cc index 08a62bcee7..70e395b47f 100644 --- a/tirx/src/backend/cuda/transform/lower_iket.cc +++ b/tirx/src/backend/cuda/transform/lower_iket.cc @@ -1183,7 +1183,7 @@ class InstrumentOfficialKernel : public StmtExprMutator { bool IketEnabled(const IRModule& module) { if (module->HasNonzeroAttr(tvm::backend::cuda::attr::kIKetEnabled.name)) return true; - const char* child_enable = std::getenv("TVM_IKET_INJECTED_CHILD_ENABLE"); + const char* child_enable = std::getenv("TVM_TIRX_IKET_INJECTED_CHILD_ENABLE"); const char* injection = std::getenv("CUDA_INJECTION64_PATH"); const char* injection_config = std::getenv("IKET_INJECTION_CONFIG"); if (!injection_config || injection_config[0] == '\0') { diff --git a/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h b/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h index 2e5ba80ead..463a58a132 100644 --- a/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h +++ b/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h @@ -66,14 +66,14 @@ ffi::Module HexagonFallbackModuleCreate(ffi::Bytes code, ffi::String fmt, * returned. * - Registry miss → `HexagonFallbackModuleObj` * returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). */ inline ffi::Module HexagonModuleCreateWithFallback( ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, HexagonSourceMap source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return HexagonFallbackModuleCreate(std::move(code), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/metal/codegen/metal_fallback_module.h b/tirx/src/backend/metal/codegen/metal_fallback_module.h index 4d8ed0eae6..f80fd690a1 100644 --- a/tirx/src/backend/metal/codegen/metal_fallback_module.h +++ b/tirx/src/backend/metal/codegen/metal_fallback_module.h @@ -62,7 +62,7 @@ ffi::Module MetalFallbackModuleCreate(ffi::Map<ffi::String, ffi::Bytes> smap, ff * returned. * - Registry miss (USE_METAL=OFF build) → `MetalFallbackModuleObj` * returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). */ @@ -70,7 +70,7 @@ inline ffi::Module MetalModuleCreateWithFallback(ffi::Map<ffi::String, ffi::Byte ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return MetalFallbackModuleCreate(std::move(smap), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/opencl/codegen/opencl_fallback_module.h b/tirx/src/backend/opencl/codegen/opencl_fallback_module.h index 6f1287abc0..9833ede6eb 100644 --- a/tirx/src/backend/opencl/codegen/opencl_fallback_module.h +++ b/tirx/src/backend/opencl/codegen/opencl_fallback_module.h @@ -61,14 +61,14 @@ ffi::Module OpenCLFallbackModuleCreate(ffi::Bytes code, ffi::String fmt, * - Registry hit (USE_OPENCL=ON build) → real OpenCL module returned. * - Registry miss (USE_OPENCL=OFF build) → `OpenCLFallbackModuleObj` * returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless * of registry state (used by per-backend fallback tests on a * USE_X=ON CI box). */ inline ffi::Module OpenCLModuleCreateWithFallback(ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return OpenCLFallbackModuleCreate(std::move(code), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/rocm/codegen/rocm_fallback_module.h b/tirx/src/backend/rocm/codegen/rocm_fallback_module.h index 1e88d506f0..763e49c446 100644 --- a/tirx/src/backend/rocm/codegen/rocm_fallback_module.h +++ b/tirx/src/backend/rocm/codegen/rocm_fallback_module.h @@ -60,14 +60,14 @@ ffi::Module ROCmFallbackModuleCreate(ffi::Bytes code, ffi::String fmt, * * - Registry hit (USE_ROCM=ON build) → real `ROCMModuleObj` returned. * - Registry miss (USE_ROCM=OFF build) → `ROCmFallbackModuleObj` returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). */ inline ffi::Module ROCmModuleCreateWithFallback(ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return ROCmFallbackModuleCreate(std::move(code), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc index f99a4d3f4a..96666e6fdc 100644 --- a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc +++ b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc @@ -670,14 +670,15 @@ spirv::Value CodeGenSPIRV::Dispatch_(const TensorLoadNode* op) { Var buffer_var = op->source.as_or_throw<tvm::tirx::TensorVar>(); PrimExpr prim_index = op->indices[0]; - PrimType desired_read_type = op->ty.as_or_throw<PrimType>(); - if (desired_read_type == PrimType::Bool()) { - desired_read_type = boolean_storage_type_.WithLanes(desired_read_type.lanes()); - } - auto it = storage_info_.find(buffer_var.get()); TVM_FFI_ICHECK(it != storage_info_.end()); StorageInfo& info = it->second; + + PrimType desired_read_type = op->ty.as_or_throw<PrimType>(); + // Local allocations retain OpTypeBool; external boolean buffers use byte storage. + if (desired_read_type == PrimType::Bool() && info.element_type != PrimType::Bool()) { + desired_read_type = boolean_storage_type_.WithLanes(desired_read_type.lanes()); + } info.CheckContentType(desired_read_type, prim_index.ty().lanes()); spirv::SType content_type = builder_->GetSType(info.element_type); @@ -696,8 +697,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const TensorLoadNode* op) { spirv::Value index = MakeValue(vec_index); spirv::Value ptr = builder_->StructArrayAccess(ptr_type, buffer, index); spirv::Value loaded = builder_->MakeValue(spv::OpLoad, content_type, ptr, mask); - // OpTypeBool have no physical address/storage. Here, cast from - // the storage type to an OpTypeBool. + // Convert externally stored bytes back to bool. Native bool loads are unchanged. if (op->ty.as_or_throw<PrimType>() == PrimType::Bool()) { auto spirv_bool = builder_->GetSType(PrimType::Bool()); loaded = builder_->Cast(spirv_bool, loaded); diff --git a/tirx/src/backend/vulkan/codegen/spirv_utils.cc b/tirx/src/backend/vulkan/codegen/spirv_utils.cc index d8239a4770..32b96f60df 100644 --- a/tirx/src/backend/vulkan/codegen/spirv_utils.cc +++ b/tirx/src/backend/vulkan/codegen/spirv_utils.cc @@ -143,7 +143,7 @@ std::pair<std::unordered_map<std::string, runtime::SPIRVShader>, std::string> Lo SPIRVShader shader = cg.BuildFunction(f, entry); - if (auto path = std::getenv("TVM_VULKAN_DEBUG_SHADER_SAVEPATH")) { + if (auto path = std::getenv("TVM_TIRX_VULKAN_DEBUG_SHADER_SAVEPATH")) { if (*path) { std::stringstream ss; ss << path << "/" << f_name << "_"; @@ -157,7 +157,7 @@ std::pair<std::unordered_map<std::string, runtime::SPIRVShader>, std::string> Lo } } - if (!support::GetEnv<bool>("TVM_VULKAN_DISABLE_SHADER_VALIDATION", false)) { + if (!support::GetEnv<bool>("TVM_TIRX_VULKAN_DISABLE_SHADER_VALIDATION", false)) { spirv_tools.ValidateShader(shader.data); } diff --git a/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h b/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h index 026ec69d56..f8c5417e7e 100644 --- a/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h +++ b/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h @@ -61,7 +61,7 @@ ffi::Module VulkanFallbackModuleCreate(ffi::Map<ffi::String, ffi::Bytes> smap, f * - Registry hit (USE_VULKAN=ON build) → real `VulkanModuleObj` returned. * - Registry miss (USE_VULKAN=OFF build) → `VulkanFallbackModuleObj` * returned. - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). */ @@ -69,7 +69,7 @@ inline ffi::Module VulkanModuleCreateWithFallback(ffi::Map<ffi::String, ffi::Byt ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return VulkanFallbackModuleCreate(std::move(smap), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h b/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h index da12b39ea6..f2722f40f6 100644 --- a/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h +++ b/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h @@ -63,7 +63,7 @@ ffi::Module WebGPUFallbackModuleCreate(ffi::Map<ffi::String, ffi::Bytes> smap, f * side; wasm runtime registers only the load-from-bytes loader). * - Registry miss → `WebGPUFallbackModuleObj` returned (the canonical * C++ side for WebGPU). - * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless + * - `TVM_TIRX_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless * (test hook, no behavioral change here since fallback is always * selected anyway). */ @@ -71,7 +71,7 @@ inline ffi::Module WebGPUModuleCreateWithFallback(ffi::Map<ffi::String, ffi::Byt ffi::String fmt, ffi::Map<ffi::String, runtime::FunctionInfo> fmap, ffi::Map<ffi::String, ffi::String> source) { - if (tvm::support::GetEnv<bool>("TVM_COMPILE_FORCE_FALLBACK", false)) { + if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) { return WebGPUFallbackModuleCreate(std::move(smap), std::move(fmt), std::move(fmap), std::move(source)); } diff --git a/tirx/tests/python/backend/cuda/codegen/test_basic.py b/tirx/tests/python/backend/cuda/codegen/test_basic.py index 38942610d8..0b2ee6d007 100644 --- a/tirx/tests/python/backend/cuda/codegen/test_basic.py +++ b/tirx/tests/python/backend/cuda/codegen/test_basic.py @@ -1142,7 +1142,7 @@ def test_device_host_call_same_func(*, backend_config): def test_thread_return(monkeypatch, *, backend_config): if not env.has_cuda(): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") @I.ir_module class Module: diff --git a/tirx/tests/python/backend/cuda/codegen/test_fp8.py b/tirx/tests/python/backend/cuda/codegen/test_fp8.py index 73536ecdc6..f94abb2f13 100644 --- a/tirx/tests/python/backend/cuda/codegen/test_fp8.py +++ b/tirx/tests/python/backend/cuda/codegen/test_fp8.py @@ -203,7 +203,9 @@ def test_fp8_vector_conversions(native_dtype, promoted_dtype, numpytype): ) -bcast_length = tvm_tirx.testing.parameter(2, 4, 6, 8) [email protected](params=[2, 4, 6, 8], scope="session") +def bcast_length(request): + return request.param @pytest.mark.gpu @@ -238,7 +240,9 @@ def test_half_broadcast(bcast_length): tvm_tirx.testing.assert_allclose(b.numpy(), b_np) -vector_length = tvm_tirx.testing.parameter(2, 4) [email protected](params=[2, 4], scope="session") +def vector_length(request): + return request.param @pytest.mark.gpu @@ -396,8 +400,9 @@ class BaseFP8E4M3QuantScaleOnly: class TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly): - # weight_shape = tvm.testing.parameter((32000, 4096), (4096, 14336)) - weight_shape = tvm_tirx.testing.parameter((128, 256), (128, 64)) + @pytest.fixture(params=[(128, 256), (128, 64)], scope="session") + def weight_shape(self, request): + return request.param @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9") diff --git a/tirx/tests/python/backend/cuda/test_module_export.py b/tirx/tests/python/backend/cuda/test_module_export.py index 9f4b75db9c..529707981c 100644 --- a/tirx/tests/python/backend/cuda/test_module_export.py +++ b/tirx/tests/python/backend/cuda/test_module_export.py @@ -127,11 +127,11 @@ def test_export_load_with_fallback(monkeypatch, tmp_path, *, backend_config): with T.launch_thread("threadIdx.x", 32) as i_1: B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0 - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") host_lib = tvm_tirx.tirx.compile( Module, target={"kind": "cuda", "backend_config": backend_config} ) - monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK") + monkeypatch.delenv("TVM_TIRX_COMPILE_FORCE_FALLBACK") lib_path = str(tmp_path / "lib.so") host_lib.export_library(lib_path) diff --git a/tirx/tests/python/backend/hexagon/codegen/test_basic.py b/tirx/tests/python/backend/hexagon/codegen/test_basic.py index b46b83427b..eda7c3e25f 100644 --- a/tirx/tests/python/backend/hexagon/codegen/test_basic.py +++ b/tirx/tests/python/backend/hexagon/codegen/test_basic.py @@ -24,9 +24,9 @@ from tvm_ffi import register_global_func import tvm_tirx import tvm_tirx.testing +from tvm_tirx.backend.hexagon.linker import has_hexagon_toolchain from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T -from tvm_tirx.testing import env @pytest.fixture(autouse=True) @@ -47,7 +47,7 @@ def register_linker(): register_global_func(name, original_linker, override=True) [email protected](not env.has_hexagon(), reason="need hexagon") [email protected](not has_hexagon_toolchain(), reason="need hexagon toolchain") def test_basic(): target = tvm_tirx.target.Target("qcom/hexagon-v66") @@ -69,7 +69,7 @@ def test_basic(): assert vadds, asm [email protected](not env.has_hexagon(), reason="need hexagon") [email protected](not has_hexagon_toolchain(), reason="need hexagon toolchain") def test_llvm_target_features(): target = tvm_tirx.target.Target("qcom/hexagon-v66") @@ -88,7 +88,7 @@ def test_llvm_target_features(): assert fs # Check that it's non-empty [email protected](not env.has_hexagon(), reason="need hexagon") [email protected](not has_hexagon_toolchain(), reason="need hexagon toolchain") def test_llvm_options(): target = tvm_tirx.target.Target( { diff --git a/tirx/tests/python/backend/metal/test_module_export.py b/tirx/tests/python/backend/metal/test_module_export.py index ed01f86302..2661d7e18a 100644 --- a/tirx/tests/python/backend/metal/test_module_export.py +++ b/tirx/tests/python/backend/metal/test_module_export.py @@ -38,9 +38,9 @@ def test_export_load_with_fallback(monkeypatch, tmp_path): with T.launch_thread("threadIdx.x", 32) as i_1: B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0 - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") host_lib = tvm_tirx.tirx.compile(Module, target="metal") - monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK") + monkeypatch.delenv("TVM_TIRX_COMPILE_FORCE_FALLBACK") lib_path = str(tmp_path / "lib.so") host_lib.export_library(lib_path) diff --git a/tirx/tests/python/backend/opencl/codegen/test_basic.py b/tirx/tests/python/backend/opencl/codegen/test_basic.py index 2e32e17696..a21f63647a 100644 --- a/tirx/tests/python/backend/opencl/codegen/test_basic.py +++ b/tirx/tests/python/backend/opencl/codegen/test_basic.py @@ -188,7 +188,7 @@ def test_opencl_type_casting(): fun = tvm_tirx.tirx.compile(Module, target=target).mod assembly = fun.imports[0].inspect_source() lcond = "convert_int4(((convert_uint4(((uint4)(((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3)))))" - rcond = "(convert_uint4(((((int4)(((convert_int(get_local_id(0))))+(1*0), ((convert_int(get_local_id(0))))+(1*1), ((convert_int(get_local_id(0))))+(1*2), ((convert_int(get_local_id(0))))+(1*3))) % ((int4)(3, 3, 3, 3))) == ((int4)(1, 1, 1, 1))))))))" + rcond = "(convert_uint4(((((int4)((((convert_int(get_local_id(0))) * 4))+(1*0), (((convert_int(get_local_id(0))) * 4))+(1*1), (((convert_int(get_local_id(0))) * 4))+(1*2), (((convert_int(get_local_id(0))) * 4))+(1*3))) % ((int4)(3, 3, 3, 3))) == ((int4)(1, 1, 1, 1))))))))" pattern_cond = f"({lcond} && {rcond})" assert assembly.count(pattern_cond) != 0 @@ -196,6 +196,8 @@ def test_opencl_type_casting(): dev = tvm_tirx.opencl(0) c = tvm_tirx.runtime.empty((n,), dtype, dev) fun(c) + expected = [float(i // 4 == 3 and i % 3 == 1) for i in range(n)] + tvm_tirx.testing.assert_allclose(c.numpy(), expected) tvm_tirx.testing.run_with_gpu_lock(run_and_check) diff --git a/tirx/tests/python/backend/opencl/test_module_export.py b/tirx/tests/python/backend/opencl/test_module_export.py index e115a99669..45f5f92f49 100644 --- a/tirx/tests/python/backend/opencl/test_module_export.py +++ b/tirx/tests/python/backend/opencl/test_module_export.py @@ -42,9 +42,9 @@ def test_export_load_with_fallback(monkeypatch, tmp_path): with T.launch_thread("threadIdx.x", 32) as i_1: B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0 - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") host_lib = tvm_tirx.tirx.compile(Module, target=target) - monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK") + monkeypatch.delenv("TVM_TIRX_COMPILE_FORCE_FALLBACK") lib_path = str(tmp_path / "lib.so") host_lib.export_library(lib_path) diff --git a/tirx/tests/python/backend/rocm/test_module_export.py b/tirx/tests/python/backend/rocm/test_module_export.py index e51fa88bb4..a37f3de895 100644 --- a/tirx/tests/python/backend/rocm/test_module_export.py +++ b/tirx/tests/python/backend/rocm/test_module_export.py @@ -39,9 +39,9 @@ def test_export_load_with_fallback(monkeypatch, tmp_path): with T.launch_thread("threadIdx.x", 32) as i_1: B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0 - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") host_lib = tvm_tirx.tirx.compile(Module, target="rocm") - monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK") + monkeypatch.delenv("TVM_TIRX_COMPILE_FORCE_FALLBACK") lib_path = str(tmp_path / "lib.so") host_lib.export_library(lib_path) diff --git a/tirx/tests/python/backend/vulkan/codegen/test_basic.py b/tirx/tests/python/backend/vulkan/codegen/test_basic.py index a8f01b124e..3ad300e14b 100644 --- a/tirx/tests/python/backend/vulkan/codegen/test_basic.py +++ b/tirx/tests/python/backend/vulkan/codegen/test_basic.py @@ -31,8 +31,15 @@ from tvm_tirx.script.ir_builder import IRBuilder from tvm_tirx.testing import env from tvm_tirx.tirx.script import ir_builder as T_builder -dtype = tvm_tirx.testing.parameter("float32", "int32", "float16", "int8") -fuzz_seed = tvm_tirx.testing.parameter(range(25)) + [email protected](params=["float32", "int32", "float16", "int8"], scope="session") +def dtype(request): + return request.param + + [email protected](params=[range(25)], scope="session") +def fuzz_seed(request): + return request.param # Explicitly specify a target, as this test is looking at the @@ -191,8 +198,14 @@ def test_vulkan_bool_load(): tvm_tirx.testing.run_with_gpu_lock(run_and_check) -vulkan_parameter_impl = tvm_tirx.testing.parameter("push_constants", "ubo") -vulkan_parameter_dtype = tvm_tirx.testing.parameter("int32", "float32", "int64") [email protected](params=["push_constants", "ubo"], scope="session") +def vulkan_parameter_impl(request): + return request.param + + [email protected](params=["int32", "float32", "int64"], scope="session") +def vulkan_parameter_dtype(request): + return request.param # Only run on vulkan because extremely large numbers of input diff --git a/tirx/tests/python/backend/vulkan/test_module_export.py b/tirx/tests/python/backend/vulkan/test_module_export.py index d24c545212..4e9e7fbe58 100644 --- a/tirx/tests/python/backend/vulkan/test_module_export.py +++ b/tirx/tests/python/backend/vulkan/test_module_export.py @@ -39,9 +39,9 @@ def test_export_load_with_fallback(monkeypatch, tmp_path): with T.launch_thread("threadIdx.x", 32) as i_1: B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0 - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") host_lib = tvm_tirx.tirx.compile(Module, target="vulkan") - monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK") + monkeypatch.delenv("TVM_TIRX_COMPILE_FORCE_FALLBACK") lib_path = str(tmp_path / "lib.so") host_lib.export_library(lib_path) diff --git a/tirx/tests/python/codegen/gpu_common/test_allreduce.py b/tirx/tests/python/codegen/gpu_common/test_allreduce.py index 14bce0a2b4..3e78b432c7 100644 --- a/tirx/tests/python/codegen/gpu_common/test_allreduce.py +++ b/tirx/tests/python/codegen/gpu_common/test_allreduce.py @@ -73,7 +73,9 @@ def generate_param_sets(): yield (d1, d2, d3) -dims = tvm_tirx.testing.parameter(*generate_param_sets()) [email protected](params=[*generate_param_sets()], scope="session") +def dims(request): + return request.param @pytest.mark.parametrize( @@ -104,7 +106,9 @@ def test_allreduce_sum(dims, target): tvm_tirx.testing.run_with_gpu_lock(run_and_check) -define_metal_compile_callback = tvm_tirx.testing.parameter(True, False) [email protected](params=[True, False], scope="session") +def define_metal_compile_callback(request): + return request.param @pytest.fixture diff --git a/tirx/tests/python/codegen/test_argument_validation.py b/tirx/tests/python/codegen/test_argument_validation.py index c738fbccac..9b22c18932 100644 --- a/tirx/tests/python/codegen/test_argument_validation.py +++ b/tirx/tests/python/codegen/test_argument_validation.py @@ -31,8 +31,12 @@ import tvm_tirx.testing from tvm_tirx.script import tirx as T from tvm_tirx.testing import env + # Parameterize over both LLVM and C backends -codegen_target = tvm_tirx.testing.parameter("llvm", "c") [email protected](params=["llvm", "c"], scope="session") +def codegen_target(request): + return request.param + # ── Argument count errors ──────────────────────────────────── diff --git a/tirx/tests/python/codegen/test_assert.py b/tirx/tests/python/codegen/test_assert.py index 992878ddf4..7c01ce2ace 100644 --- a/tirx/tests/python/codegen/test_assert.py +++ b/tirx/tests/python/codegen/test_assert.py @@ -22,7 +22,10 @@ import tvm_tirx import tvm_tirx.testing from tvm_tirx.script import tirx as T -codegen_target = tvm_tirx.testing.parameter("llvm", "c") + [email protected](params=["llvm", "c"], scope="session") +def codegen_target(request): + return request.param def test_assert_runtime_error(codegen_target): diff --git a/tirx/tests/python/codegen/test_nvshmem.py b/tirx/tests/python/codegen/test_nvshmem.py index bde2e3afba..c1d02f6c0c 100644 --- a/tirx/tests/python/codegen/test_nvshmem.py +++ b/tirx/tests/python/codegen/test_nvshmem.py @@ -34,7 +34,7 @@ from tvm_tirx.testing import env @pytest.fixture def cuda_source(monkeypatch): # Source generation needs neither NVSHMEM headers nor a CUDA device. - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") def generate(function): target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}) diff --git a/tirx/tests/python/runtime/test_memory_scope.py b/tirx/tests/python/runtime/test_memory_scope.py index 042001e41a..9a38e68250 100644 --- a/tirx/tests/python/runtime/test_memory_scope.py +++ b/tirx/tests/python/runtime/test_memory_scope.py @@ -89,7 +89,7 @@ def test_runtime_tensor_allocation(): def test_runtime_only_registration(): - env = {**os.environ, "TVM_USE_RUNTIME_LIB": "1"} + env = {**os.environ, "TVM_TIRX_USE_RUNTIME_LIB": "1"} subprocess.run( [ sys.executable, @@ -134,8 +134,8 @@ def test_opencl_memory_scope_registration_load_order(load_order): pytest.skip("Compiler library is not available") env = { **os.environ, - "TVM_USE_RUNTIME_LIB": "1" if load_order == "runtime_only" else "0", - "TVM_DEVICE_BACKEND_AUTOLOAD": "0", + "TVM_TIRX_USE_RUNTIME_LIB": "1" if load_order == "runtime_only" else "0", + "TVM_TIRX_DEVICE_BACKEND_AUTOLOAD": "0", } subprocess.run( [ diff --git a/tirx/tests/python/runtime/test_module.py b/tirx/tests/python/runtime/test_module.py index 9d7b80a509..73083d7aa1 100644 --- a/tirx/tests/python/runtime/test_module.py +++ b/tirx/tests/python/runtime/test_module.py @@ -86,7 +86,7 @@ np.testing.assert_equal(b.numpy(), a.numpy() + 1) """, str(native_library), ], - env={**os.environ, "TVM_USE_RUNTIME_LIB": "1"}, + env={**os.environ, "TVM_TIRX_USE_RUNTIME_LIB": "1"}, check=True, ) diff --git a/tirx/tests/python/sym/test_const_int_bound.py b/tirx/tests/python/sym/test_const_int_bound.py index 5503850da4..6699ff8247 100644 --- a/tirx/tests/python/sym/test_const_int_bound.py +++ b/tirx/tests/python/sym/test_const_int_bound.py @@ -67,12 +67,17 @@ class BaseCompare: class TestDataType(BaseCompare): - test_case = tvm_tirx.testing.parameter( - TestCase(tvm_tirx.tirx.Var("x", "int64"), (NEG_INF, POS_INF)), - TestCase(tvm_tirx.tirx.Var("x", "int8"), (-128, 127)), - TestCase(tvm_tirx.tirx.Var("x", "uint8"), (0, 255)), - TestCase(tvm_tirx.tirx.Var("x", "int32"), (-(2**31), 2**31 - 1)), + @pytest.fixture( + params=[ + TestCase(tvm_tirx.tirx.Var("x", "int64"), (NEG_INF, POS_INF)), + TestCase(tvm_tirx.tirx.Var("x", "int8"), (-128, 127)), + TestCase(tvm_tirx.tirx.Var("x", "uint8"), (0, 255)), + TestCase(tvm_tirx.tirx.Var("x", "int32"), (-(2**31), 2**31 - 1)), + ], + scope="session", ) + def test_case(self, request): + return request.param def test_plain_var_non_negative_bound_requires_context(): @@ -89,23 +94,33 @@ class TestCastBound(BaseCompare): x = tvm_tirx.tirx.Var("x", "int8") tmod = tvm_tirx.tirx.truncmod - test_case = tvm_tirx.testing.parameter( - TestCase(tmod(x, 3).astype("uint32"), (0, 2)), - TestCase(tmod(x, 3).astype("float32").astype("int32"), (-2, 2)), + @pytest.fixture( + params=[ + TestCase(tmod(x, 3).astype("uint32"), (0, 2)), + TestCase(tmod(x, 3).astype("float32").astype("int32"), (-2, 2)), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestAddSubBound(BaseCompare): x = tvm_tirx.tirx.Var("x", "int64") y = tvm_tirx.tirx.Var("y", "int64") - test_case = tvm_tirx.testing.parameter( - TestCase(x + y, (NEG_INF, POS_INF)), - TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}), - TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}), - TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}), - TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}), + @pytest.fixture( + params=[ + TestCase(x + y, (NEG_INF, POS_INF)), + TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}), + TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}), + TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}), + TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}), + ], + scope="session", ) + def test_case(self, request): + return request.param @pytest.mark.xfail(reason="Not currently supported") @@ -133,7 +148,7 @@ class TestBoundsUsingReciprocals(BaseCompare): symmetric_bounds = {A: (1, 4095), B: (1, 4095), C: (2048, 2048)} asymmetric_bounds = {A: (1, 1024), B: (1, POS_INF), C: (2048, 2048)} - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase((A + B) * C - A * B, (2048, None), known_bounds=symmetric_bounds), TestCase((A + B) * C - B * A, (2048, None), known_bounds=symmetric_bounds), TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds), @@ -144,15 +159,24 @@ class TestBoundsUsingReciprocals(BaseCompare): TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestMulBound(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}), - TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}), - TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}), + @pytest.fixture( + params=[ + TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}), + TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}), + TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestTruncDivBound(BaseCompare): @@ -160,12 +184,17 @@ class TestTruncDivBound(BaseCompare): expr = tvm_tirx.tirx.truncdiv(x, y) - test_case = tvm_tirx.testing.parameter( - TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}), - TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}), - TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}), - TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}), + @pytest.fixture( + params=[ + TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}), + TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}), + TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}), + TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestTruncModBound(BaseCompare): @@ -173,11 +202,16 @@ class TestTruncModBound(BaseCompare): expr = tvm_tirx.tirx.truncmod(x, y) - test_case = tvm_tirx.testing.parameter( - TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}), - TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}), - TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}), + @pytest.fixture( + params=[ + TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}), + TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}), + TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestFloorDivBound(BaseCompare): @@ -185,23 +219,33 @@ class TestFloorDivBound(BaseCompare): ux = tvm_tirx.tirx.Var("x", "uint32") uy = tvm_tirx.tirx.Var("y", "uint32") - test_case = tvm_tirx.testing.parameter( - TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}), - TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}), - TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}), - TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}), - TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}), + @pytest.fixture( + params=[ + TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}), + TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}), + TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}), + TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}), + TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestFloorModBound(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}), - TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}), - TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}), + @pytest.fixture( + params=[ + TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}), + TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}), + TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestModBoundWithModularSet(BaseCompare): @@ -222,7 +266,7 @@ class TestModBoundWithModularSet(BaseCompare): n = tvm_tirx.tirx.Var("n", "int64") tmod = tvm_tirx.tirx.truncmod - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( # gcd(320, 256) = 64, base 255 -> residue 63: values {63, 127, 191, 255} TestCase((n * 320 + 255) % 256, (63, 255)), # coeff divides the divisor, base 0: multiples of 16 @@ -250,38 +294,57 @@ class TestModBoundWithModularSet(BaseCompare): TestCase(tmod(n, 256), (-255, 0), {n: (-1000, -300)}), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestMinMaxBound(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(tvm_tirx.tirx.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}), - TestCase(tvm_tirx.tirx.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}), - TestCase(tvm_tirx.tirx.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}), - TestCase(tvm_tirx.tirx.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}), + @pytest.fixture( + params=[ + TestCase(tvm_tirx.tirx.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}), + TestCase(tvm_tirx.tirx.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}), + TestCase(tvm_tirx.tirx.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}), + TestCase(tvm_tirx.tirx.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestSelectBound(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase( - tvm_tirx.tirx.Select(x > 1, (y < 0).astype("int32"), y + 1), - (0, 11), - {x: (-9, 11), y: (4, 10)}, - ), + @pytest.fixture( + params=[ + TestCase( + tvm_tirx.tirx.Select(x > 1, (y < 0).astype("int32"), y + 1), + (0, 11), + {x: (-9, 11), y: (4, 10)}, + ), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestShiftAndBound(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}), - TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}), - TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}), + @pytest.fixture( + params=[ + TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}), + TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}), + TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestMixIndexBound(BaseCompare): @@ -289,29 +352,45 @@ class TestMixIndexBound(BaseCompare): tdiv = tvm_tirx.tirx.truncdiv tmod = tvm_tirx.tirx.truncmod - test_case = tvm_tirx.testing.parameter( - TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}), - TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}), - TestCase( - tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)} - ), + @pytest.fixture( + params=[ + TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}), + TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}), + TestCase( + tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)} + ), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestLetBound(BaseCompare): x = tvm_tirx.tirx.Var("x", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(tvm_tirx.tirx.Let(x, 1, x + 1), (2, 2)), + + @pytest.fixture( + params=[ + TestCase(tvm_tirx.tirx.Let(x, 1, x + 1), (2, 2)), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestFloorModNegativeDivisor(BaseCompare): flm, fld = tvm_tirx.tirx.floormod, tvm_tirx.tirx.floordiv a, b = tvm_tirx.tirx.Var("a", "int32"), tvm_tirx.tirx.Var("b", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}), + @pytest.fixture( + params=[ + TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestDivModAssumeNoZeroDivisor(BaseCompare): @@ -322,38 +401,61 @@ class TestDivModAssumeNoZeroDivisor(BaseCompare): a, b = tvm_tirx.tirx.Var("a", "int32"), tvm_tirx.tirx.Var("b", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}), - TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}), + @pytest.fixture( + params=[ + TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}), + TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestMultipleCondition(BaseCompare): a = tvm_tirx.tirx.Var("a", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase( - a % 58 - 1, - (0, None), - known_bounds={a: (0, 128)}, - constraint=tvm_tirx.tirx.all(1 <= a % 58, a % 58 < 57), - ), + + @pytest.fixture( + params=[ + TestCase( + a % 58 - 1, + (0, None), + known_bounds={a: (0, 128)}, + constraint=tvm_tirx.tirx.all(1 <= a % 58, a % 58 < 57), + ), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestBroadcastBound(BaseCompare): a = tvm_tirx.tirx.Var("a", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase(tvm_tirx.ir.prim.vector_expr.Broadcast(a, 4), (0, 128), {a: (0, 128)}), + + @pytest.fixture( + params=[ + TestCase(tvm_tirx.ir.prim.vector_expr.Broadcast(a, 4), (0, 128), {a: (0, 128)}), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestRampBound(BaseCompare): a = tvm_tirx.tirx.Var("a", "int32") - test_case = tvm_tirx.testing.parameter( - TestCase( - tvm_tirx.ir.prim.vector_expr.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)} - ), + + @pytest.fixture( + params=[ + TestCase( + tvm_tirx.ir.prim.vector_expr.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)} + ), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestModularSetBound(BaseCompare): @@ -363,9 +465,14 @@ class TestModularSetBound(BaseCompare): expr = (bx * 2048 + tx * 16) % 7168 - test_case = tvm_tirx.testing.parameter( - TestCase(expr, (0, 7152), {bx: (0, 3584), tx: (0, 128)}), + @pytest.fixture( + params=[ + TestCase(expr, (0, 7152), {bx: (0, 3584), tx: (0, 128)}), + ], + scope="session", ) + def test_case(self, request): + return request.param if __name__ == "__main__": diff --git a/tirx/tests/python/sym/test_iter_affine_map.py b/tirx/tests/python/sym/test_iter_affine_map.py index 84ded92be0..8c2aecc47f 100644 --- a/tirx/tests/python/sym/test_iter_affine_map.py +++ b/tirx/tests/python/sym/test_iter_affine_map.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: F841 +import pytest + import tvm_tirx import tvm_tirx.testing from tvm_tirx.script import tirx as T @@ -1112,14 +1114,17 @@ class TestPadding: fld = tvm_tirx.tirx.floordiv flm = tvm_tirx.tirx.floormod - positive_test_case = tvm_tirx.testing.parameter( + _positive_test_case_params = ( # left padding only, offset divisible ({y: 192}, {fld(64 + y, 32): (6, 2, 1), flm(64 + y, 32): (32, 0, 1)}, "bijective"), # left padding only, offset non-divisible ({y: 176}, {fld(80 + y, 32): (6, 2, 1)}), ({y: 176}, {flm(fld(80 + y, 2), 16): (16, 0, 1), flm(80 + y, 2): (2, 0, 1)}), # right padding only, offset divisible - ({x: 5, y: 4}, {fld(x * 32 + y * 8, 16): (10, 0, 1), flm(x * 32 + y * 8, 16): (2, 0, 8)}), + ( + {x: 5, y: 4}, + {fld(x * 32 + y * 8, 16): (10, 0, 1), flm(x * 32 + y * 8, 16): (2, 0, 8)}, + ), # right padding only, offset non-divisible ({x: 26}, {fld(x, 15): (2, 0, 1)}), ({x: 26}, {flm(fld(x, 3), 5): (5, 0, 1), flm(x, 3): (3, 0, 1)}), @@ -1150,7 +1155,11 @@ class TestPadding: ), ) - negative_test_case = tvm_tirx.testing.parameter( + @pytest.fixture(params=_positive_test_case_params, scope="session") + def positive_test_case(self, request): + return request.param + + _negative_test_case_params = ( # left padding only, offset non-divisible ({y: 176}, {fld(80 + y, 32), flm(80 + y, 32)}), ({y: 176}, {fld(80 + y, 32), fld(80 + y, 4)}), @@ -1174,6 +1183,10 @@ class TestPadding: ({x: 255}, {fld(flm(x, 255), 16)}), ) + @pytest.fixture(params=_negative_test_case_params, scope="session") + def negative_test_case(self, request): + return request.param + def test_padding(self, positive_test_case): iter_extent, mapped_iterators, *args = positive_test_case check_level = args[0] if args else "surjective" diff --git a/tirx/tests/python/sym/test_rewrite_simplify.py b/tirx/tests/python/sym/test_rewrite_simplify.py index 326f960f0b..05c015c462 100644 --- a/tirx/tests/python/sym/test_rewrite_simplify.py +++ b/tirx/tests/python/sym/test_rewrite_simplify.py @@ -96,7 +96,8 @@ class TestVector(BaseCompare): x64 = tvm_tirx.tirx.Var("x", "int64") vx = tvm_tirx.tirx.Var("vx", "int32x2") vc = tvm_tirx.tirx.Var("vc", "bool") - test_case = tvm_tirx.testing.parameter( + + _test_case_params = ( # Add rules TestCase( tvm_tirx.ir.prim.vector_expr.Ramp(x, 1, 4) + tvm_tirx.ir.prim.vector_expr.Ramp(y, 2, 4), @@ -190,7 +191,9 @@ class TestVector(BaseCompare): tvm_tirx.ir.prim.vector_expr.Ramp(tdiv(x, 2), 2, tirx.vscale() * 5), ), TestCase( - tdiv(tvm_tirx.ir.prim.vector_expr.Ramp(x * 8 + 1, 1, 4), 8), x.astype("int32x4"), x >= 0 + tdiv(tvm_tirx.ir.prim.vector_expr.Ramp(x * 8 + 1, 1, 4), 8), + x.astype("int32x4"), + x >= 0, ), TestCase( tdiv(tvm_tirx.ir.prim.vector_expr.Ramp(x * 8 + 15, 1, 4), 8), @@ -501,6 +504,10 @@ class TestVector(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestSelect(BaseCompare): x, y, z = ( @@ -509,7 +516,7 @@ class TestSelect(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( # Add rules TestCase( tvm_tirx.tirx.Select(x < 0, y, 0) + tvm_tirx.tirx.Select(x < 0, 1, z), @@ -534,6 +541,10 @@ class TestSelect(BaseCompare): TestCase(tvm_tirx.tirx.Select(x > 0, y + 1, y + 1), y + 1), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestCancellation(BaseCompare): var_int8 = tirx.Var("var_int8", "int8") @@ -543,7 +554,7 @@ class TestCancellation(BaseCompare): var_uint32 = tirx.Var("var_uint32", "uint32") var_uint64 = tirx.Var("var_uint64", "uint64") - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(tirx.const(5, "int64") - tirx.const(5, "int64"), tirx.const(0, "int64")), TestCase(tirx.const(5, "uint8") - tirx.const(5, "uint8"), tirx.const(0, "uint8")), TestCase(var_int8 - var_int8, tirx.const(0, "int8")), @@ -574,6 +585,10 @@ class TestCancellation(BaseCompare): TestCase(tirx.NE(var_uint64, var_uint64), tirx.const(False, "bool")), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestAddIndex(BaseCompare): x, y, z = ( @@ -582,7 +597,7 @@ class TestAddIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(x + (y - x), y), TestCase(x - (y + 1) + (y + 1), x), TestCase((x - 10) + (10 - z), x - z), @@ -626,6 +641,10 @@ class TestAddIndex(BaseCompare): TestCase(fld(flm(x, 2) + 7, 2) + fld(x, 2), fld(x + 7, 2)), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestSubIndex(BaseCompare): x, y, z = ( @@ -634,7 +653,7 @@ class TestSubIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(x + y - y, x), TestCase(x + y - x, y), TestCase(x - (y + x), 0 - y), @@ -725,6 +744,10 @@ class TestSubIndex(BaseCompare): TestCase(fld(y - z, 3) * 6 - 2 * y, (0 - flm(y - z, 3) - z) * 2), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestMulIndex(BaseCompare): x, y, z = ( @@ -732,13 +755,19 @@ class TestMulIndex(BaseCompare): tvm_tirx.tirx.Var("y", "int32"), tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( - TestCase((x + 2) * 3, x * 3 + 6), - TestCase((x * 2) * 3, x * 6), - TestCase(tvm_tirx.tirx.min(x, y) * tvm_tirx.tirx.max(x, y), x * y), - TestCase(tvm_tirx.tirx.max(x, y) * tvm_tirx.tirx.min(x, y), x * y), - TestCase((x - y) * (-2), (y - x) * 2), + + @pytest.fixture( + params=[ + TestCase((x + 2) * 3, x * 3 + 6), + TestCase((x * 2) * 3, x * 6), + TestCase(tvm_tirx.tirx.min(x, y) * tvm_tirx.tirx.max(x, y), x * y), + TestCase(tvm_tirx.tirx.max(x, y) * tvm_tirx.tirx.min(x, y), x * y), + TestCase((x - y) * (-2), (y - x) * 2), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestDivIndex(BaseCompare): @@ -749,7 +778,7 @@ class TestDivIndex(BaseCompare): ) non_negative = [x >= 0, y >= 0, z >= 0] - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(tdiv(x, x), 1), TestCase(tdiv(tdiv(x, 2), 3), tdiv(x, 6)), TestCase(tdiv(tdiv(x, 2) + 1, 3), tdiv(x + 2, 6), non_negative), @@ -757,17 +786,25 @@ class TestDivIndex(BaseCompare): TestCase(tdiv(x * 4, 2), x * 2), TestCase(tdiv(x * 4 + y, 2), x * 2 + tdiv(y, 2), non_negative), TestCase( - tdiv(tvm_tirx.tirx.min(x * 6, y), 2), tvm_tirx.tirx.min(x * 3, tdiv(y, 2)), non_negative + tdiv(tvm_tirx.tirx.min(x * 6, y), 2), + tvm_tirx.tirx.min(x * 3, tdiv(y, 2)), + non_negative, ), TestCase( - tdiv(tvm_tirx.tirx.max(x * 6, y), 2), tvm_tirx.tirx.max(x * 3, tdiv(y, 2)), non_negative + tdiv(tvm_tirx.tirx.max(x * 6, y), 2), + tvm_tirx.tirx.max(x * 3, tdiv(y, 2)), + non_negative, ), TestCase(tdiv(y + x * 4, 2), tdiv(y, 2) + x * 2, non_negative), TestCase( - tdiv(tvm_tirx.tirx.min(y, x * 6), 2), tvm_tirx.tirx.min(tdiv(y, 2), x * 3), non_negative + tdiv(tvm_tirx.tirx.min(y, x * 6), 2), + tvm_tirx.tirx.min(tdiv(y, 2), x * 3), + non_negative, ), TestCase( - tdiv(tvm_tirx.tirx.max(y, x * 6), 2), tvm_tirx.tirx.max(tdiv(y, 2), x * 3), non_negative + tdiv(tvm_tirx.tirx.max(y, x * 6), 2), + tvm_tirx.tirx.max(tdiv(y, 2), x * 3), + non_negative, ), # 3-operands TestCase(tdiv(x * 6 + y + z, 2), x * 3 + tdiv(y + z, 2), non_negative), @@ -789,6 +826,10 @@ class TestDivIndex(BaseCompare): TestCase(tdiv(y + z * x, z), tdiv(y, z) + x, non_negative), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestFloordivIndex(BaseCompare): x, y, z = ( @@ -797,7 +838,7 @@ class TestFloordivIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(fld(fld(x, 2), 3), fld(x, 6)), TestCase(fld(fld(x, 2) + 1, 3), fld(x + 2, 6)), TestCase(fld(x - flm(x, 21), 21), fld(x, 21)), @@ -846,6 +887,10 @@ class TestFloordivIndex(BaseCompare): TestCase(fld(x * 128 + y * 4 + z, 512), fld(x, 4), [y >= 0, y < 32, z >= 0, z < 4]), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestModIndex(BaseCompare): x, y, nx, ny, z = ( @@ -856,7 +901,7 @@ class TestModIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( # TODO(Lunderberg): Loosen these preconditions. When there's # a single term whose factor is divisible by the denominator, # the sign of the argument doesn't matter. @@ -884,6 +929,10 @@ class TestModIndex(BaseCompare): TestCase(tmod(x + ny * (-10), -2), tmod(x, 2), [x >= 0, ny <= 0]), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestFloormodIndex(BaseCompare): x, y, z = ( @@ -892,7 +941,7 @@ class TestFloormodIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase(flm(x * 10, 2), 0), TestCase(flm(x * 9600, 6400), flm(x * 3200, 6400)), TestCase(flm(x * 10 + y, 2), flm(y, 2)), @@ -915,6 +964,10 @@ class TestFloormodIndex(BaseCompare): # TestCase(flm(x * 10 + 1 + y * 2 + 2, 2), 1), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestFloorModTwo(BaseCompare): """Special-case simplifications for FloorMod(expr,2) @@ -936,7 +989,8 @@ class TestFloorModTwo(BaseCompare): tvm_tirx.tirx.Var("y", "int32"), tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + + _test_case_params = ( # Removing offsets from floormod TestCase(flm(x, 2) + flm(x + 1, 2), 1), TestCase(flm(x + 1, 2) + flm(x, 2), 1), @@ -958,6 +1012,10 @@ class TestFloorModTwo(BaseCompare): TestCase(flm(x + 1, 2) * 8192, flm(x + 1, 2) * 8192, [x >= 0, x < 2]), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestFloorModPadded(BaseCompare): """Special-case simplifications for divisibility proof @@ -965,7 +1023,8 @@ class TestFloorModPadded(BaseCompare): """ x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") - test_case = tvm_tirx.testing.parameter( + + _test_case_params = ( TestCase(flm(x - flm(x, 9), 9), 0), TestCase(flm(x - flm(x, -9), 9), 0), TestCase(flm(x + flm(-x, 9), 9), 0), @@ -975,6 +1034,10 @@ class TestFloorModPadded(BaseCompare): TestCase(flm(x + flm(-x, y), y), 0), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + def test_uint_floormod_const_fold(): analyzer = tvm_tirx.sym.Analyzer() @@ -994,7 +1057,8 @@ class TestFloorModUnsigned(BaseCompare): x = tirx.Var("x", "uint32") y = tirx.Var("y", "uint32") z = tirx.Var("z", "int32") - test_case = tvm_tirx.testing.parameter( + + _test_case_params = ( TestCase( flm(x * tirx.const(8, "uint32"), tirx.const(4, "uint32")), tirx.const(0, "uint32") ), @@ -1005,15 +1069,22 @@ class TestFloorModUnsigned(BaseCompare): ), TestCase( flm( - x * tirx.const(32, "uint32") + y * tirx.const(8, "uint32"), tirx.const(8, "uint32") + x * tirx.const(32, "uint32") + y * tirx.const(8, "uint32"), + tirx.const(8, "uint32"), ), tirx.const(0, "uint32"), ), # Non-power-of-two divisor must NOT simplify (unsigned wraparound # changes residues mod c2). TestCase( - flm(x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), tirx.const(12, "uint32")), - flm(x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), tirx.const(12, "uint32")), + flm( + x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), + tirx.const(12, "uint32"), + ), + flm( + x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), + tirx.const(12, "uint32"), + ), ), # The pre-existing x * c1 % c2 -> 0 rule is pow2-guarded: c2=12 must # not fold to 0 (e.g. x=2^29 wraps: uint32(x*12)=2^31, and 2^31 % 12 = 8). @@ -1032,6 +1103,10 @@ class TestFloorModUnsigned(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestMinIndex(BaseCompare): x, y, z = ( @@ -1039,7 +1114,8 @@ class TestMinIndex(BaseCompare): tvm_tirx.tirx.Var("y", "int32"), tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + + _test_case_params = ( # const int bound TestCase(tvm_tirx.tirx.min(tmod(x, 2), tmod(y, 2) + 10), tmod(x, 2)), TestCase(tvm_tirx.tirx.min(flm(x, 2), flm(y, 2) + 10), flm(x, 2)), @@ -1078,7 +1154,8 @@ class TestMinIndex(BaseCompare): TestCase( tvm_tirx.tirx.min( tvm_tirx.tirx.min( - tvm_tirx.tirx.min(tvm_tirx.tirx.min(tvm_tirx.tirx.min(x, y), z), x * 2), z * 2 + tvm_tirx.tirx.min(tvm_tirx.tirx.min(tvm_tirx.tirx.min(x, y), z), x * 2), + z * 2, ), y, ), @@ -1134,7 +1211,8 @@ class TestMinIndex(BaseCompare): ), TestCase(tvm_tirx.tirx.min(tdiv(x, 10), tdiv(y, 10)), tdiv(tvm_tirx.tirx.min(x, y), 10)), TestCase( - tvm_tirx.tirx.min(tdiv(x, (-10)), tdiv(y, (-10))), tdiv(tvm_tirx.tirx.max(x, y), (-10)) + tvm_tirx.tirx.min(tdiv(x, (-10)), tdiv(y, (-10))), + tdiv(tvm_tirx.tirx.max(x, y), (-10)), ), # floor div TestCase(tvm_tirx.tirx.min(fld(x + 3, 4) * 4, x), x), @@ -1156,6 +1234,10 @@ class TestMinIndex(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestMaxIndex(BaseCompare): x, y, z = ( @@ -1164,10 +1246,11 @@ class TestMaxIndex(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( # Identical operands simplify even when they contain an opaque call. TestCase( - tvm_tirx.tirx.max(x + tirx.vscale() * 4, x + tirx.vscale() * 4), x + tirx.vscale() * 4 + tvm_tirx.tirx.max(x + tirx.vscale() * 4, x + tirx.vscale() * 4), + x + tirx.vscale() * 4, ), # const int bound TestCase(tvm_tirx.tirx.max(tmod(x, 2), tmod(y, 2) + 10), tmod(y, 2) + 10), @@ -1207,7 +1290,8 @@ class TestMaxIndex(BaseCompare): TestCase( tvm_tirx.tirx.max( tvm_tirx.tirx.max( - tvm_tirx.tirx.max(tvm_tirx.tirx.max(tvm_tirx.tirx.max(x, y), z), x * 2), z * 2 + tvm_tirx.tirx.max(tvm_tirx.tirx.max(tvm_tirx.tirx.max(x, y), z), x * 2), + z * 2, ), y, ), @@ -1251,7 +1335,8 @@ class TestMaxIndex(BaseCompare): # truc div TestCase(tvm_tirx.tirx.max(tdiv(x, 10), tdiv(y, 10)), tdiv(tvm_tirx.tirx.max(x, y), 10)), TestCase( - tvm_tirx.tirx.max(tdiv(x, (-10)), tdiv(y, (-10))), tdiv(tvm_tirx.tirx.min(x, y), (-10)) + tvm_tirx.tirx.max(tdiv(x, (-10)), tdiv(y, (-10))), + tdiv(tvm_tirx.tirx.min(x, y), (-10)), ), TestCase(tvm_tirx.tirx.max(tdiv(x + 3, 4) * 4, x), tdiv(x + 3, 4) * 4), # floordiv @@ -1264,6 +1349,10 @@ class TestMaxIndex(BaseCompare): TestCase(tvm_tirx.tirx.max(x, fld(x, 4) * 4), x), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestComparisons(BaseCompare): x, y, z = ( @@ -1272,7 +1361,7 @@ class TestComparisons(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( # const int bound TestCase((tmod(x, 2) + 10).equal(0), tvm_tirx.tirx.const(0, "bool")), TestCase(tvm_tirx.tirx.NE(tmod(x, 2) + 10, 0), tvm_tirx.tirx.const(1, "bool")), @@ -1453,7 +1542,9 @@ class TestComparisons(BaseCompare): TestCase(x + y >= -10, tvm_tirx.tirx.const(1, "bool"), [x >= 0, y >= -10]), TestCase(z - 5 <= y + 10, tvm_tirx.tirx.const(1, "bool"), [z <= 5, y >= -10]), TestCase( - tvm_tirx.tirx.all(x > -1, z <= x + 5), tvm_tirx.tirx.const(1, "bool"), [x >= 0, z <= 5] + tvm_tirx.tirx.all(x > -1, z <= x + 5), + tvm_tirx.tirx.const(1, "bool"), + [x >= 0, z <= 5], ), TestCase(x * y <= 0, tvm_tirx.tirx.const(1, "bool"), [x >= 0, y <= 0]), TestCase((x + 1) * (y - 1) < 0, tvm_tirx.tirx.const(1, "bool"), [x >= 0, y <= 0]), @@ -1462,6 +1553,10 @@ class TestComparisons(BaseCompare): TestCase(tmod(y - 1, 3) == 0, tmod(y + (-1), 3) == 0), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestLogical(BaseCompare): x, y, z = ( @@ -1470,7 +1565,7 @@ class TestLogical(BaseCompare): tvm_tirx.tirx.Var("z", "int32"), ) - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase( tvm_tirx.tirx.And(tvm_tirx.tirx.EQ(x, y), tvm_tirx.tirx.NE(x, y)), tvm_tirx.tirx.const(False, "bool"), @@ -1480,7 +1575,8 @@ class TestLogical(BaseCompare): tvm_tirx.tirx.const(False, "bool"), ), TestCase( - tvm_tirx.tirx.And(x > 1, tvm_tirx.tirx.Not(x > 1)), tvm_tirx.tirx.const(False, "bool") + tvm_tirx.tirx.And(x > 1, tvm_tirx.tirx.Not(x > 1)), + tvm_tirx.tirx.const(False, "bool"), ), TestCase(tvm_tirx.tirx.And(x <= y, y < x), tvm_tirx.tirx.const(False, "bool")), TestCase(tvm_tirx.tirx.And(y < x, x <= y), tvm_tirx.tirx.const(False, "bool")), @@ -1527,14 +1623,23 @@ class TestLogical(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestLet(BaseCompare): x, y = tvm_tirx.tirx.Var("x", "int32"), tvm_tirx.tirx.Var("y", "int32") z = tvm_tirx.tirx.Let(x, 1, x + 1) - test_case = tvm_tirx.testing.parameter( - TestCase(z + z, 4), + @pytest.fixture( + params=[ + TestCase(z + z, 4), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestCast(BaseCompare): @@ -1552,14 +1657,22 @@ class TestCast(BaseCompare): tvm_tirx.tirx.const(i, dtype1), ) - test_case = tvm_tirx.testing.parameter(*_generate_tests()) + @pytest.fixture(params=[*_generate_tests()], scope="session") + def test_case(self, request): + return request.param class TestShiftLeft(BaseCompare): z = tvm_tirx.tirx.LShift(tvm_tirx.tirx.const(1, "int32"), tvm_tirx.tirx.const(10, "int32")) - test_case = tvm_tirx.testing.parameter( - TestCase(z, tvm_tirx.tirx.const(1 << 10, "int32")), + + @pytest.fixture( + params=[ + TestCase(z, tvm_tirx.tirx.const(1 << 10, "int32")), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestBitwiseConstants(BaseCompare): @@ -1580,34 +1693,46 @@ class TestBitwiseConstants(BaseCompare): tvm_tirx.tirx.const(2, "int32"), ) - test_case = tvm_tirx.testing.parameter(*_generate_tests()) + @pytest.fixture(params=[*_generate_tests()], scope="session") + def test_case(self, request): + return request.param class TestDivZero(BaseCompare): ramp = tvm_tirx.ir.prim.vector_expr.Ramp(1, 1, 2) broadcast = tvm_tirx.ir.prim.vector_expr.Broadcast(0, 2) - test_case = tvm_tirx.testing.parameter( - TestCase(tvm_tirx.tirx.Div(ramp, broadcast), RuntimeError), - TestCase(tvm_tirx.tirx.Mod(ramp, broadcast), RuntimeError), - TestCase(tvm_tirx.tirx.FloorDiv(ramp, broadcast), RuntimeError), - TestCase(tvm_tirx.tirx.FloorMod(ramp, broadcast), RuntimeError), + @pytest.fixture( + params=[ + TestCase(tvm_tirx.tirx.Div(ramp, broadcast), RuntimeError), + TestCase(tvm_tirx.tirx.Mod(ramp, broadcast), RuntimeError), + TestCase(tvm_tirx.tirx.FloorDiv(ramp, broadcast), RuntimeError), + TestCase(tvm_tirx.tirx.FloorMod(ramp, broadcast), RuntimeError), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestSubBufferload(BaseCompare): buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([1], "float32")) load = tvm_tirx.ir.TensorLoad(buf, [0]) - test_case = tvm_tirx.testing.parameter( - TestCase(load - load, 0.0), + @pytest.fixture( + params=[ + TestCase(load - load, 0.0), + ], + scope="session", ) + def test_case(self, request): + return request.param class TestIfThenElse(BaseCompare): x = tvm_tirx.tirx.Var("x", "int32") - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase( tvm_tirx.tirx.if_then_else(x < 5, tvm_tirx.tirx.if_then_else(x > 1, 1, 0), 0), tvm_tirx.tirx.if_then_else( @@ -1620,9 +1745,13 @@ class TestIfThenElse(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + class TestCLZ(BaseCompare): - test_case = tvm_tirx.testing.parameter( + _test_case_params = ( TestCase( tvm_tirx.tirx.call_intrin("prim.clz", 0, ty="int32"), T.int32(32), @@ -1657,6 +1786,10 @@ class TestCLZ(BaseCompare): ), ) + @pytest.fixture(params=_test_case_params, scope="session") + def test_case(self, request): + return request.param + if __name__ == "__main__": tvm_tirx.testing.main() diff --git a/tirx/tests/python/sym/test_simplify.py b/tirx/tests/python/sym/test_simplify.py index dd0eebb802..e9b35b2c83 100644 --- a/tirx/tests/python/sym/test_simplify.py +++ b/tirx/tests/python/sym/test_simplify.py @@ -38,7 +38,7 @@ def test_simplify_reshape_flattened_index(): ) -dtype = tvm_tirx.testing.parameter( +_dtype_params = ( "uint8", "uint16", "uint32", @@ -53,6 +53,11 @@ dtype = tvm_tirx.testing.parameter( ) [email protected](params=_dtype_params, scope="session") +def dtype(request): + return request.param + + def test_can_prove_self_identity(dtype): ana = tvm_tirx.sym.Analyzer() diff --git a/tirx/tests/python/target/test_target.py b/tirx/tests/python/target/test_target.py index aa7db72d2c..edef313001 100644 --- a/tirx/tests/python/target/test_target.py +++ b/tirx/tests/python/target/test_target.py @@ -399,17 +399,14 @@ def test_target_from_device_vulkan(input_form): def run_and_check(): dev = tvm_tirx.vulkan() target = Target.from_device("vulkan" if input_form == "string" else dev) - f_get_target_property = tvm_ffi.get_global_func("device_api.vulkan.get_target_property") assert target.kind.name == "vulkan" assert target.attrs["max_threads_per_block"] == dev.max_threads_per_block assert int(target.attrs["max_shared_memory_per_block"]) == dev.max_shared_memory_per_block assert int(target.attrs["thread_warp_size"]) == dev.warp_size - assert target.attrs["supports_float16"] == f_get_target_property(dev, "supports_float16") - assert target.attrs["supports_int16"] == f_get_target_property(dev, "supports_int16") - assert target.attrs["supports_int8"] == f_get_target_property(dev, "supports_int8") - assert target.attrs["supports_16bit_buffer"] == f_get_target_property( - dev, "supports_16bit_buffer" - ) + assert target.attrs["supports_float16"] == dev.get_attr("supports_float16") + assert target.attrs["supports_int16"] == dev.get_attr("supports_int16") + assert target.attrs["supports_int8"] == dev.get_attr("supports_int8") + assert target.attrs["supports_16bit_buffer"] == dev.get_attr("supports_16bit_buffer") tvm_tirx.testing.run_with_gpu_lock(run_and_check) diff --git a/tirx/tests/python/tirx/base/test_index_map.py b/tirx/tests/python/tirx/base/test_index_map.py index 7136e12f82..bd97dc387d 100644 --- a/tirx/tests/python/tirx/base/test_index_map.py +++ b/tirx/tests/python/tirx/base/test_index_map.py @@ -127,110 +127,115 @@ def test_nonbijective_inverse_gives_error(): dynamic_N = tvm_tirx.tirx.Var("N", "int32") -padding_test_case = tvm_tirx.testing.parameter( - by_dict={ - "no_padding": dict( - forward=lambda i: [i // 4, i % 4], - inverse=lambda i, j: [4 * i + j], - pre_shape=[16], - post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_ffi.convert(False), +_padding_test_cases = { + "no_padding": dict( + forward=lambda i: [i // 4, i % 4], + inverse=lambda i, j: [4 * i + j], + pre_shape=[16], + post_shape=[T.int32(4), T.int32(4)], + padding=lambda i, j: tvm_ffi.convert(False), + ), + "right_padding": dict( + forward=lambda i: [i // 4, i % 4], + inverse=lambda i, j: [4 * i + j], + pre_shape=[15], + post_shape=[T.int32(4), T.int32(4)], + padding=lambda i, j: tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), + ), + "left_padding": dict( + forward=lambda i: [(i + 1) // 4, (i + 1) % 4], + inverse=lambda i, j: [4 * i + j - 1], + pre_shape=[15], + post_shape=[T.int32(4), T.int32(4)], + padding=lambda i, j: tvm_tirx.tirx.And(i == 0, j < 1), + ), + "left_and_right_padding": dict( + forward=lambda i: [(i + 1) // 4, (i + 1) % 4], + inverse=lambda i, j: [4 * i + j - 1], + pre_shape=[14], + post_shape=[T.int32(4), T.int32(4)], + padding=lambda i, j: tvm_tirx.tirx.Or( + tvm_tirx.tirx.And(i == 0, j < 1), + tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), ), - "right_padding": dict( - forward=lambda i: [i // 4, i % 4], - inverse=lambda i, j: [4 * i + j], - pre_shape=[15], - post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), + ), + "dynamic_size": dict( + forward=lambda i: [i // 4, i % 4], + inverse=lambda i, j: [4 * i + j], + pre_shape=[dynamic_N], + post_shape=[(dynamic_N - dynamic_N % (-4)) // 4, T.int32(4)], + padding=lambda i, j: tvm_tirx.tirx.And( + dynamic_N % (-4) != 0, + tvm_tirx.tirx.And(i == dynamic_N // 4, j >= dynamic_N % 4), ), - "left_padding": dict( - forward=lambda i: [(i + 1) // 4, (i + 1) % 4], - inverse=lambda i, j: [4 * i + j - 1], - pre_shape=[15], - post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_tirx.tirx.And(i == 0, j < 1), - ), - "left_and_right_padding": dict( - forward=lambda i: [(i + 1) // 4, (i + 1) % 4], - inverse=lambda i, j: [4 * i + j - 1], - pre_shape=[14], - post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_tirx.tirx.Or( - tvm_tirx.tirx.And(i == 0, j < 1), - tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), + ), + "2d_padding": dict( + forward=lambda i, j: [(i + 1) // 4, (j + 5) // 8, (i + 1) % 4, (j + 5) % 8], + inverse=lambda i_outer, j_outer, i_inner, j_inner: [ + 4 * i_outer + i_inner - 1, + 8 * j_outer + j_inner - 5, + ], + pre_shape=[14, 31], + post_shape=[ + T.int32(4), # ceildiv(left_pad + i.extent, 4) = ceildiv(1 + 14, 4) = 4 + T.int32(5), # ceildiv(left_pad + j.extent, 8) = ceildiv(5 + 31, 8) = 5 + T.int32(4), # Range of iter%4 + T.int32(8), # Range of iter%8 + ], + padding=lambda i_outer, j_outer, i_inner, j_inner: tvm_tirx.tirx.Or( + tvm_tirx.tirx.Or( + tvm_tirx.tirx.And(i_outer == 0, i_inner < 1), + tvm_tirx.tirx.And(i_outer == 3, tvm_ffi.convert(3) == i_inner), ), - ), - "dynamic_size": dict( - forward=lambda i: [i // 4, i % 4], - inverse=lambda i, j: [4 * i + j], - pre_shape=[dynamic_N], - post_shape=[(dynamic_N - dynamic_N % (-4)) // 4, T.int32(4)], - padding=lambda i, j: tvm_tirx.tirx.And( - dynamic_N % (-4) != 0, - tvm_tirx.tirx.And(i == dynamic_N // 4, j >= dynamic_N % 4), + tvm_tirx.tirx.Or( + tvm_tirx.tirx.And(j_outer == 0, j_inner < 5), + tvm_tirx.tirx.And(j_outer == 4, j_inner >= 4), ), ), - "2d_padding": dict( - forward=lambda i, j: [(i + 1) // 4, (j + 5) // 8, (i + 1) % 4, (j + 5) % 8], - inverse=lambda i_outer, j_outer, i_inner, j_inner: [ - 4 * i_outer + i_inner - 1, - 8 * j_outer + j_inner - 5, - ], - pre_shape=[14, 31], - post_shape=[ - T.int32(4), # ceildiv(left_pad + i.extent, 4) = ceildiv(1 + 14, 4) = 4 - T.int32(5), # ceildiv(left_pad + j.extent, 8) = ceildiv(5 + 31, 8) = 5 - T.int32(4), # Range of iter%4 - T.int32(8), # Range of iter%8 - ], - padding=lambda i_outer, j_outer, i_inner, j_inner: tvm_tirx.tirx.Or( - tvm_tirx.tirx.Or( - tvm_tirx.tirx.And(i_outer == 0, i_inner < 1), - tvm_tirx.tirx.And(i_outer == 3, tvm_ffi.convert(3) == i_inner), - ), - tvm_tirx.tirx.Or( - tvm_tirx.tirx.And(j_outer == 0, j_inner < 5), - tvm_tirx.tirx.And(j_outer == 4, j_inner >= 4), - ), - ), - ), - "multiple_right_padding": dict( - forward=lambda i: [i // 32, (i // 4) % 8, i % 4], - inverse=lambda i, j, k: [32 * i + 4 * j + k], - pre_shape=[116], - post_shape=[T.int32(4), T.int32(8), T.int32(4)], - padding=lambda i, j, k: tvm_tirx.tirx.And(i == 3, 4 * j + k >= 20), - ), - "multiple_right_padding_transpose": dict( - forward=lambda i: [(i // 4) % 8, i // 32, i % 4], - inverse=lambda j, i, k: [32 * i + 4 * j + k], - pre_shape=[116], - post_shape=[T.int32(8), T.int32(4), T.int32(4)], - padding=lambda j, i, k: tvm_tirx.tirx.And(i == 3, 4 * j + k >= 20), - ), - "multiple_left_padding": dict( - forward=lambda i: [(i + 5) // 32, ((i + 5) // 4) % 8, (i + 5) % 4], - inverse=lambda i, j, k: [32 * i + 4 * j + k - 5], - pre_shape=[123], - post_shape=[T.int32(4), T.int32(8), T.int32(4)], - padding=lambda i, j, k: tvm_tirx.tirx.And(i == 0, j * 4 + k < 5), - ), - "multiple_left_padding_with_transpose": dict( - forward=lambda i: [((i + 5) // 4) % 8, (i + 5) // 32, (i + 5) % 4], - inverse=lambda j, i, k: [32 * i + 4 * j + k - 5], - pre_shape=[123], - post_shape=[T.int32(8), T.int32(4), T.int32(4)], - padding=lambda j, i, k: tvm_tirx.tirx.And(i == 0, j * 4 + k < 5), - ), - "outer_loop_extent_one": dict( - forward=lambda i: [i // 4, i % 4], - inverse=lambda i, j: [i * 4 + j], - pre_shape=[3], - post_shape=[T.int32(1), T.int32(4)], - padding=lambda i, j: tvm_ffi.convert(3) == j, - ), - } + ), + "multiple_right_padding": dict( + forward=lambda i: [i // 32, (i // 4) % 8, i % 4], + inverse=lambda i, j, k: [32 * i + 4 * j + k], + pre_shape=[116], + post_shape=[T.int32(4), T.int32(8), T.int32(4)], + padding=lambda i, j, k: tvm_tirx.tirx.And(i == 3, 4 * j + k >= 20), + ), + "multiple_right_padding_transpose": dict( + forward=lambda i: [(i // 4) % 8, i // 32, i % 4], + inverse=lambda j, i, k: [32 * i + 4 * j + k], + pre_shape=[116], + post_shape=[T.int32(8), T.int32(4), T.int32(4)], + padding=lambda j, i, k: tvm_tirx.tirx.And(i == 3, 4 * j + k >= 20), + ), + "multiple_left_padding": dict( + forward=lambda i: [(i + 5) // 32, ((i + 5) // 4) % 8, (i + 5) % 4], + inverse=lambda i, j, k: [32 * i + 4 * j + k - 5], + pre_shape=[123], + post_shape=[T.int32(4), T.int32(8), T.int32(4)], + padding=lambda i, j, k: tvm_tirx.tirx.And(i == 0, j * 4 + k < 5), + ), + "multiple_left_padding_with_transpose": dict( + forward=lambda i: [((i + 5) // 4) % 8, (i + 5) // 32, (i + 5) % 4], + inverse=lambda j, i, k: [32 * i + 4 * j + k - 5], + pre_shape=[123], + post_shape=[T.int32(8), T.int32(4), T.int32(4)], + padding=lambda j, i, k: tvm_tirx.tirx.And(i == 0, j * 4 + k < 5), + ), + "outer_loop_extent_one": dict( + forward=lambda i: [i // 4, i % 4], + inverse=lambda i, j: [i * 4 + j], + pre_shape=[3], + post_shape=[T.int32(1), T.int32(4)], + padding=lambda i, j: tvm_ffi.convert(3) == j, + ), +} + + [email protected]( + params=_padding_test_cases.values(), ids=_padding_test_cases.keys(), scope="session" ) +def padding_test_case(request): + return request.param def test_nonsurjective_inverse(padding_test_case): diff --git a/tirx/tests/python/tirx/base/test_intrin.py b/tirx/tests/python/tirx/base/test_intrin.py index faef661b39..b2fdf2d95d 100644 --- a/tirx/tests/python/tirx/base/test_intrin.py +++ b/tirx/tests/python/tirx/base/test_intrin.py @@ -266,7 +266,9 @@ def test_ldexp(): ) -dtype = tvm_tirx.testing.parameter("int32", "int64") [email protected](params=["int32", "int64"], scope="session") +def dtype(request): + return request.param @pytest.mark.parametrize( diff --git a/tirx/tests/python/tirx/codegen/test_cuda.py b/tirx/tests/python/tirx/codegen/test_cuda.py index 299ec0f355..652afbe8eb 100644 --- a/tirx/tests/python/tirx/codegen/test_cuda.py +++ b/tirx/tests/python/tirx/codegen/test_cuda.py @@ -236,7 +236,7 @@ def test_typed_buffer_bindings(monkeypatch, alignment): B[0] = alias[0] storage[0] = T.uint32(0) - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}) source = tvm_ffi.get_global_func("target.build.cuda")( tvm_tirx.IRModule.from_expr(main), target @@ -278,7 +278,7 @@ def test_vector_pointer_preserves_packed_offset(monkeypatch): ) lowered = tvm_tirx.tirx.transform.LowerIntrin()(tvm_tirx.IRModule.from_expr(func)) - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") source = tvm_ffi.get_global_func("target.build.cuda")(lowered, target).inspect_source() call = next(line.strip() for line in source.splitlines() if line.strip().startswith("consume(")) diff --git a/tirx/tests/python/tirx/iket/test_orchestration.py b/tirx/tests/python/tirx/iket/test_orchestration.py index 0d8f057e97..29852d9b0e 100644 --- a/tirx/tests/python/tirx/iket/test_orchestration.py +++ b/tirx/tests/python/tirx/iket/test_orchestration.py @@ -54,8 +54,8 @@ def _fake_profile_process(captured, *, artifact_count: int = 1): def test_public_namespace_and_old_bench_path_is_absent(): - pytest.importorskip("triton") # tvm.tirx.bench pulls in the triton harness - import tvm_tirx.tirx.bench as bench + pytest.importorskip("triton") # tvm_tirx.testing.bench pulls in the triton harness + import tvm_tirx.testing.bench as bench assert iket.__all__ == [ "IketProfileError", @@ -147,7 +147,7 @@ def test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey assert captured["cwd"] == cwd assert captured["timeout"] == 12.5 assert captured["env"]["IKET_TEST_ENV"] == "present" - assert captured["env"]["TVM_IKET_INJECTED_CHILD_ENABLE"] == "1" + assert captured["env"]["TVM_TIRX_IKET_INJECTED_CHILD_ENABLE"] == "1" assert os.environ["IKET_TEST_ENV"] == "inherited" assert result.output_dir == target assert result.command == (sys.executable, "workload.py") diff --git a/tirx/tests/python/tirx/iket/test_profiler.py b/tirx/tests/python/tirx/iket/test_profiler.py index c63a44afcc..dca8cbebd2 100644 --- a/tirx/tests/python/tirx/iket/test_profiler.py +++ b/tirx/tests/python/tirx/iket/test_profiler.py @@ -46,7 +46,7 @@ ORACLE_PATH = Path(__file__).parent / "oracle" / "iket_official_cutlass_4_6_0_or def _source_codegen_without_cuda(monkeypatch): # Metadata and source checks do not require an installed CUDA compiler. if not env.has_cuda(): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") @T.function @@ -448,7 +448,7 @@ def test_regular_lowering_strips_annotations_and_tokens(): def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch, config_env): monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False) monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False) - monkeypatch.setenv("TVM_IKET_INJECTED_CHILD_ENABLE", "1") + monkeypatch.setenv("TVM_TIRX_IKET_INJECTED_CHILD_ENABLE", "1") monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False) monkeypatch.setenv("CUDA_INJECTION64_PATH", "/verified/libsmodel_injection.so") monkeypatch.setenv(config_env, "/verified/config.json") @@ -461,14 +461,14 @@ def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch, co @pytest.mark.parametrize( "missing_env", - ("CUDA_INJECTION64_PATH", "config", "TVM_IKET_INJECTED_CHILD_ENABLE"), + ("CUDA_INJECTION64_PATH", "config", "TVM_TIRX_IKET_INJECTED_CHILD_ENABLE"), ) @pytest.mark.parametrize("config_env", ["IKET_INJECTION_CONFIG", "SMODEL_INJECTION_CONFIG"]) def test_injected_child_auto_enable_remains_fail_closed(monkeypatch, missing_env, config_env): monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False) monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False) values = { - "TVM_IKET_INJECTED_CHILD_ENABLE": "1", + "TVM_TIRX_IKET_INJECTED_CHILD_ENABLE": "1", "CUDA_INJECTION64_PATH": "/verified/libsmodel_injection.so", config_env: "/verified/config.json", } @@ -737,7 +737,7 @@ def test_extended_payload_compile_only_architectures(arch, capfd): not _has_packaged_nvdisasm(), reason="need the pinned nvidia-cuda-nvdisasm distribution" ) def test_native_and_extended_payload_placeholder_sass(monkeypatch, tmp_path): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") native_no_payload_source = _cuda_source(_compile(push_pop_kernel)) native_no_payload_sass = _nvrtc_disassemble( native_no_payload_source, tmp_path / "native-no-payload" diff --git a/tirx/tests/python/tirx/test_axis_enum.py b/tirx/tests/python/tirx/test_axis_enum.py index a885f7f9f0..712848c802 100644 --- a/tirx/tests/python/tirx/test_axis_enum.py +++ b/tirx/tests/python/tirx/test_axis_enum.py @@ -59,7 +59,7 @@ def test_axis_unknown_name_and_json_roundtrip(): def test_axis_runtime_only_import(): env = os.environ.copy() - env["TVM_USE_RUNTIME_LIB"] = "1" + env["TVM_TIRX_USE_RUNTIME_LIB"] = "1" subprocess.run( [ sys.executable, diff --git a/tirx/tests/python/tirx/test_bench_utils.py b/tirx/tests/python/tirx/test_bench_utils.py index e2af67dc2c..61bac30bc9 100644 --- a/tirx/tests/python/tirx/test_bench_utils.py +++ b/tirx/tests/python/tirx/test_bench_utils.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Tests for tvm_tirx.tirx.bench utilities.""" +"""Tests for tvm_tirx.testing.bench utilities.""" import importlib import inspect @@ -23,12 +23,12 @@ from types import SimpleNamespace import pytest import torch -pytest.importorskip("triton") # tvm.tirx.bench imports triton.profiler +pytest.importorskip("triton") # tvm_tirx.testing.bench imports triton.profiler from tvm_tirx.testing import env -from tvm_tirx.tirx.bench import DistributedBenchContext, bench +from tvm_tirx.testing.bench import DistributedBenchContext, bench -bench_module = importlib.import_module("tvm_tirx.tirx.bench") +bench_module = importlib.import_module("tvm_tirx.testing.bench") class _FakeStream: diff --git a/tirx/tests/python/tirx/test_compilation_pipeline.py b/tirx/tests/python/tirx/test_compilation_pipeline.py index 889d4733e2..872db11121 100644 --- a/tirx/tests/python/tirx/test_compilation_pipeline.py +++ b/tirx/tests/python/tirx/test_compilation_pipeline.py @@ -49,9 +49,9 @@ def test_explicit_tile_pass_order_and_diagnostic(monkeypatch, capfd, print_after from tvm_tirx.tirx.compilation_pipeline import tirx_pipeline if print_after_dispatch is None: - monkeypatch.delenv("TVM_PRINT_AFTER_TIRX_DISPATCH_OPS", raising=False) + monkeypatch.delenv("TVM_TIRX_PRINT_AFTER_DISPATCH_OPS", raising=False) else: - monkeypatch.setenv("TVM_PRINT_AFTER_TIRX_DISPATCH_OPS", print_after_dispatch) + monkeypatch.setenv("TVM_TIRX_PRINT_AFTER_DISPATCH_OPS", print_after_dispatch) @T.function def kernel(A: T.Tensor((2, 4), "float32", layout="default")): diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py index 6f5b42603c..68934898d1 100644 --- a/tirx/tests/python/tirx/test_cuda_backend_config.py +++ b/tirx/tests/python/tirx/test_cuda_backend_config.py @@ -360,7 +360,7 @@ def test_serialized_config_replays_in_fresh_process(tmp_path, monkeypatch, fallb x = T.cuda.thread_idx("x") B[x] = A[x] * T.float32(0.5) - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1" if fallback else "0") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1" if fallback else "0") config = { "cuda": { "arch": env.cuda_arch(0), @@ -398,7 +398,7 @@ np.testing.assert_array_equal(b.numpy(), x * np.float32(0.5)) } ) script = script.replace("LOAD_TARGET", repr(load_context)) - environment = dict(os.environ, TVM_COMPILE_FORCE_FALLBACK="0") + environment = dict(os.environ, TVM_TIRX_COMPILE_FORCE_FALLBACK="0") result = subprocess.run( [sys.executable, "-c", script, library], env=environment, text=True, capture_output=True ) @@ -491,7 +491,7 @@ def test_device_helpers_are_cloned_per_configuration(): @pytest.mark.skipif(not tvm_tirx.cuda().exist, reason="requires CUDA") @pytest.mark.parametrize("fallback", [False, True]) def test_cuda_host_embeds_independent_binaries_in_plain_cpp(tmp_path, monkeypatch, fallback): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1" if fallback else "0") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1" if fallback else "0") _require_compiler("nvrtc") _require_compiler("nvcc") import tvm_ffi.cpp @@ -638,7 +638,7 @@ else: """ path = tmp_path / "generic_cuda.py" path.write_text(script) - environment = dict(os.environ, CUDA_VISIBLE_DEVICES="", TVM_TEST_TARGETS="llvm;cuda") + environment = dict(os.environ, CUDA_VISIBLE_DEVICES="") result = subprocess.run( [sys.executable, str(path)], env=environment, text=True, capture_output=True ) @@ -679,7 +679,7 @@ assert "arch" not in tvm_tirx.target.Target("cuda").attrs def test_source_fallback_does_not_invoke_compiler(monkeypatch): from tvm_tirx.support import nvcc - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") def unavailable(*args, **kwargs): raise AssertionError("source fallback must not require a CUDA compiler") @@ -703,7 +703,7 @@ def test_source_fallback_does_not_invoke_compiler(monkeypatch): def test_identical_entries_share_one_module(monkeypatch): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") config = {"cuda": {"arch": "sm_100a", "nvrtc": ["--ftz=false"]}} @T.function @@ -723,7 +723,7 @@ def test_identical_entries_share_one_module(monkeypatch): def test_build_snapshots_config_before_passes(monkeypatch): - monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") + monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1") config = {"cuda": {"arch": "sm_100a", "nvrtc": ["--ftz=false"]}} @T.function diff --git a/tirx/tests/python/tirx/test_thread_scope.py b/tirx/tests/python/tirx/test_thread_scope.py index a205e84b00..c8fd97f6ee 100644 --- a/tirx/tests/python/tirx/test_thread_scope.py +++ b/tirx/tests/python/tirx/test_thread_scope.py @@ -132,7 +132,7 @@ def test_cross_backend_scopes_are_rejected_at_each_boundary(): def test_runtime_only_import_and_removed_interfaces(): env = os.environ.copy() - env["TVM_USE_RUNTIME_LIB"] = "1" + env["TVM_TIRX_USE_RUNTIME_LIB"] = "1" subprocess.run( [ sys.executable, diff --git a/tirx/tests/python/tirx/test_tile_namespace.py b/tirx/tests/python/tirx/test_tile_namespace.py index fbffa563c0..48997db995 100644 --- a/tirx/tests/python/tirx/test_tile_namespace.py +++ b/tirx/tests/python/tirx/test_tile_namespace.py @@ -56,7 +56,7 @@ assert "tvm_tirx.backend.cuda.axis" not in sys.modules assert "tvm_tirx.backend.trn.axis" not in sys.modules """, ], - env={**os.environ, "TVM_DEVICE_BACKEND_AUTOLOAD": "0"}, + env={**os.environ, "TVM_TIRX_DEVICE_BACKEND_AUTOLOAD": "0"}, capture_output=True, text=True, ) diff --git a/tirx/tests/python/tirx/transform/test_fp8_legalize.py b/tirx/tests/python/tirx/transform/test_fp8_legalize.py index f0126a8421..0f9ad490dd 100644 --- a/tirx/tests/python/tirx/transform/test_fp8_legalize.py +++ b/tirx/tests/python/tirx/transform/test_fp8_legalize.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import pytest + import tvm_tirx import tvm_tirx.script import tvm_tirx.testing @@ -202,8 +204,14 @@ def get_after_storage_legalize(dtype: str, promote_dtype: str): return After -dtype = tvm_tirx.testing.parameter("float8_e4m3fn", "float8_e5m2") -promote_dtype = tvm_tirx.testing.parameter("float16", "float32") [email protected](params=["float8_e4m3fn", "float8_e5m2"], scope="session") +def dtype(request): + return request.param + + [email protected](params=["float16", "float32"], scope="session") +def promote_dtype(request): + return request.param def test_fp8_compute_legalize(dtype, promote_dtype): diff --git a/tirx/tests/python/tirx/utils.py b/tirx/tests/python/tirx/utils.py deleted file mode 100644 index 13a83393a9..0000000000 --- a/tirx/tests/python/tirx/utils.py +++ /dev/null @@ -1,16 +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.
