This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 07e60343bc [Script][Tests] Fix dialect redirect module re-execution
and stray category-less tirx.intrin_test op (#19731)
07e60343bc is described below
commit 07e60343bc73587d6e4dd6c562e15589fabbf10f
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Jun 11 14:59:34 2026 -0400
[Script][Tests] Fix dialect redirect module re-execution and stray
category-less tirx.intrin_test op (#19731)
This PR fixes two independent test-isolation issues that only surface
when certain test files run together in one pytest session.
1. Fix `_DialectRedirectFinder` duplicate module execution
`_DialectRedirectFinder.find_spec` used to pre-register the redirect
target module under the legacy alias name before returning the alias
spec.
This interacts badly with CPython import logic: when the requested
module name is already in `sys.modules`, CPython may ignore the returned
alias spec and reuse the target module's original spec instead. As a
result, the target source can be executed again under the canonical
module name, creating a duplicate module object.
This caused patches on aliased modules to silently miss the module
object used by existing code. For example,
`unittest.mock.patch("tvm.tirx.script.builder.buffer_store")` patched
the duplicate module, while the tirx parser still held references to the
original one, so `test_scalar_assign_error_not_swallowed` failed with
`DID NOT RAISE`.
This pr removes the pre-registration and let the import machinery
register the alias normally. Since the alias spec is now used,
`_AliasLoader.exec_module` also restores the canonical `__spec__` and
`__loader__` to avoid stale alias metadata on the loaded module.
2. Remove unused `tirx.intrin_test` op registration
`test_s_tir_transform_lower_match_buffer.py` registered a dummy op:
```python
tvm.ir.register_op_attr("tirx.intrin_test", "")
```
This was a leftover from the old TVMScript parser and is no longer
needed. The modern tirx parser eagerly evaluates `intrin_test(...)`
calls into `T.evaluate(0)`, so this op never appears in parsed IR.
The only remaining effect was adding a category-less `tirx.intrin_test`
entry to the global op registry, which could break
`test_registered_tirx_ops_have_exactly_one_category` depending on test
import order.
This pr removes the unused registration.
---
python/tvm/script/__init__.py | 33 +++++++++++++++++-----
.../test_s_tir_transform_lower_match_buffer.py | 5 +++-
2 files changed, 30 insertions(+), 8 deletions(-)
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index 99c7aec2bf..940739df07 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.<dialect>`` (or ``tvm.script.parser.<dialect>``, 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 @@ class _DialectRedirectFinder:
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
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 b385ce7249..a174c771a2 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]
[email protected]_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
