From 7082e35719b92a05c4e663d28207876ebf51e9e8 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Thu, 11 Jun 2026 02:23:30 -0400 Subject: [PATCH 1/2] [Script] Fix dialect redirect finder re-executing target modules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _DialectRedirectFinder.find_spec pre-registered the redirect target in sys.modules under the legacy alias name before returning its alias spec. CPython's _bootstrap._find_spec, upon seeing the requested name already in sys.modules after a finder returns, discards the returned spec in favor of module.__spec__ — the target's original SourceFileLoader spec, whose spec.name is the canonical module path. _load_unlocked then created a fresh module, re-executed the target's source into it, and replaced the canonical sys.modules entry with the duplicate. Any code holding the original module object (e.g. the tirx parser's `builder as T` global) then diverged from sys.modules, silently defeating unittest.mock.patch on the canonical path. This surfaced as tests/python/tirx/test_parser_printer.py::test_scalar_assign_error_not_swallowed failing with DID-NOT-RAISE whenever an earlier test imported tvm.s_tir.tensor_intrin (whose arm_cpu module imports tvm.script.ir_builder.tirx in statement form). Drop the pre-registration so the import machinery loads the alias spec itself, and make _AliasLoader.exec_module restore the canonical __spec__/__loader__ that module_from_spec unconditionally stamps onto the returned module (a stale alias __spec__.parent otherwise raises "__package__ != __spec__.parent" DeprecationWarnings on relative imports inside the aliased module). --- python/tvm/script/__init__.py | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py index 99c7aec2bf05..940739df076f 100644 --- a/python/tvm/script/__init__.py +++ b/python/tvm/script/__init__.py @@ -154,9 +154,11 @@ class _DialectRedirectFinder: When the import machinery asks for a module whose full name starts with ``tvm.script.`` (or ``tvm.script.parser.``, etc.) and that dialect is in ``_DIALECT_REGISTRY``, :meth:`find_spec` imports the - real target module (e.g. ``tvm.tirx.script.parser.entry``) and registers - it in ``sys.modules`` under the legacy name, so all subsequent imports and - attribute walks resolve correctly without going through the redirect again. + real target module (e.g. ``tvm.tirx.script.parser.entry``) and returns an + alias spec whose loader hands back that module, so the import machinery + registers it in ``sys.modules`` under the legacy name and all subsequent + imports and attribute walks resolve without going through the redirect + again. """ @classmethod @@ -164,9 +166,15 @@ def find_spec(cls, fullname, path, target=None): redirected = _redirect_target(fullname) if redirected is None: return None - # Resolve the target module and alias it under the legacy name. + # Resolve the target module and return an alias spec for it. Do NOT + # pre-register ``sys.modules[fullname]`` here: if ``fullname`` is in + # ``sys.modules`` when find_spec returns, ``_bootstrap._find_spec`` + # discards the returned spec in favor of ``module.__spec__`` — the + # target's original SourceFileLoader spec — and ``_load_unlocked`` + # then RE-EXECUTES the target module and replaces its canonical + # ``sys.modules`` entry with the duplicate. The machinery registers + # the alias under ``fullname`` itself when loading the spec below. module = importlib.import_module(redirected) - sys.modules[fullname] = module return importlib.util.spec_from_loader(fullname, _AliasLoader(module)) @@ -175,13 +183,24 @@ class _AliasLoader: def __init__(self, module): self._module = module + # ``module_from_spec`` unconditionally stamps the alias spec onto the + # module returned by ``create_module``; capture the canonical values + # so ``exec_module`` can restore them. + self._spec = getattr(module, "__spec__", None) + self._loader = getattr(module, "__loader__", None) def create_module(self, spec): return self._module def exec_module(self, module): - # Module is already populated by the redirect target. - return None + # Module is already populated by the redirect target; just restore + # the canonical ``__spec__``/``__loader__`` that the import machinery + # overwrote with the alias spec (a stale alias ``__spec__.parent`` + # breaks relative imports inside the module). + if self._spec is not None: + module.__spec__ = self._spec + if self._loader is not None: + module.__loader__ = self._loader # Install the redirect finder once. Re-importing tvm.script (e.g. during a From 100d95a89bb3434ced683978a98a9c431f7bf984 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Thu, 11 Jun 2026 02:37:47 -0400 Subject: [PATCH 2/2] [Tests][TIR] Drop vestigial tirx.intrin_test op registration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_s_tir_transform_lower_match_buffer.py registered a dummy op via @tvm.ir.register_op_attr("tirx.intrin_test", "") — a leftover from the legacy TVMScript parser, mechanically renamed from tir.intrin_test during the tirx namespace bring-up. The modern parser evaluates the intrin_test(...) calls eagerly (they collapse to T.evaluate(0)), so the op never appears in any parsed IR; the registration's only effect is creating a category-less "tirx.intrin_test" entry in the global op registry. That entry breaks the exactly-one-category invariant checked by tests/python/tirx/test_op_namespace_cleanup.py::test_registered_tirx_ops_have_exactly_one_category whenever this file is imported earlier in the same pytest session, e.g. pytest tests/python/s_tir/transform/ \ tests/python/tirx/test_op_namespace_cleanup.py Keep intrin_test as a plain Python helper and drop the registration. --- .../transform/test_s_tir_transform_lower_match_buffer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py index 42f0c98d0568..39cda0d4aa90 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py @@ -64,7 +64,10 @@ def transformed_buffer_load_store(a: T.handle, c: T.handle) -> None: A[i * 4 + ii, j, k * 2 + kk] += C[i * 4 + ii, k * 2 + kk] -@tvm.ir.register_op_attr("tirx.intrin_test", "") +# Dummy intrinsic whose arguments exercise match_buffer fields. TVMScript +# evaluates the call eagerly (to 0), so it must NOT be registered as an op: +# registering "tirx.intrin_test" only leaves a category-less op in the tirx +# registry, breaking the exactly-one-category invariant for later tests. def intrin_test(data, elem_offset, stride_0, stride_1, shape_0, shape_1): return 0