[Script][Tests] Fix dialect redirect module re-execution and stray category-less tirx.intrin_test op#19731
Conversation
_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).
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.
There was a problem hiding this comment.
Code Review
This pull request updates the import redirect mechanism in python/tvm/script/__init__.py to avoid pre-registering the legacy module name in sys.modules during find_spec, returning an alias spec instead. It also updates _AliasLoader to capture and restore the canonical __spec__ and __loader__ attributes of the redirected module to prevent them from being overwritten. Additionally, it removes the @tvm.ir.register_op_attr decorator from a dummy intrinsic in the tests to avoid breaking registry invariants. The reviewer suggests also capturing and restoring the canonical __name__ and __package__ attributes in _AliasLoader to prevent mutating the canonical module's identity and package context.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@tvm-bot rerun |
…tegory-less tirx.intrin_test op (apache#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. (cherry picked from commit 07e6034)
…tegory-less tirx.intrin_test op (apache#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. (cherry picked from commit 07e6034)
This PR fixes two independent test-isolation issues that only surface when certain test files run together in one pytest session.
_DialectRedirectFinderduplicate module execution_DialectRedirectFinder.find_specused 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, sotest_scalar_assign_error_not_swallowedfailed withDID 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_modulealso restores the canonical__spec__and__loader__to avoid stale alias metadata on the loaded module.tirx.intrin_testop registrationtest_s_tir_transform_lower_match_buffer.pyregistered a dummy op:This was a leftover from the old TVMScript parser and is no longer needed. The modern tirx parser eagerly evaluates
intrin_test(...)calls intoT.evaluate(0), so this op never appears in parsed IR.The only remaining effect was adding a category-less
tirx.intrin_testentry to the global op registry, which could breaktest_registered_tirx_ops_have_exactly_one_categorydepending on test import order.This pr removes the unused registration.