Skip to content

[Script][Tests] Fix dialect redirect module re-execution and stray category-less tirx.intrin_test op#19731

Merged
tqchen merged 2 commits into
apache:mainfrom
tlopex:fix-tirx-test-isolation
Jun 11, 2026
Merged

[Script][Tests] Fix dialect redirect module re-execution and stray category-less tirx.intrin_test op#19731
tqchen merged 2 commits into
apache:mainfrom
tlopex:fix-tirx-test-isolation

Conversation

@tlopex

@tlopex tlopex commented Jun 11, 2026

Copy link
Copy Markdown
Member

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.

  1. Remove unused tirx.intrin_test op registration

test_s_tir_transform_lower_match_buffer.py registered a dummy op:

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.

tlopex added 2 commits June 11, 2026 02:23
_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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread python/tvm/script/__init__.py
Comment thread python/tvm/script/__init__.py
@tlopex

tlopex commented Jun 11, 2026

Copy link
Copy Markdown
Member Author

@tvm-bot rerun

@tqchen tqchen merged commit 07e6034 into apache:main Jun 11, 2026
9 of 10 checks passed
MasterJH5574 pushed a commit to MasterJH5574/tvm that referenced this pull request Jun 15, 2026
…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)
MasterJH5574 pushed a commit to MasterJH5574/tvm that referenced this pull request Jun 15, 2026
…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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants