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.
## Summary
Restructure TVMScript to be dialect-agnostic at the script-core layer
while letting each extension dialect (TIRX, Relax) own its own
per-dialect script subtree. IR is below script in the dependency
stack and is NOT a peer dialect — its script handlers stay in the
shared core.
This PR folds together two coupled refactors that were initially
opened as separate PRs (#19478 and the original #19479); they
share rename / relocation surface so they ship as one cohesive
change.
## What this PR does
### Per-dialect script subtree (originally #19479)
- Moves per-dialect printer + builder from
`src/script/{printer,ir_builder}/{tirx,relax}/` to
`src/{tirx,relax}/script/{printer,builder}/`.
- Tightens `src/script/*.cc` CMake glob to the dialect-free core.
- Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry
(`script.ir_builder.decl_function.<type-key>`); removes
cross-dialect includes from the shared core.
- Adds `tvm.script.register_dialect` API + `__getattr__` + a
`sys.meta_path` finder for Python-side dialect discovery.
In-tree dialects (tirx, relax) registered centrally in
`python/tvm/__init__.py`.
- Drops the obsolete static re-export shims at
`python/tvm/script/{parser,ir_builder}/{tirx,relax}/`.
### Dialect-agnostic printer config (originally #19478)
- Relocates `include/tvm/ir/script_printer.h` →
`include/tvm/script/printer/config.h` next to the rest of the
printer's public surface. The header is not IR-specific.
- Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for
clarity (the macro registers Script as the kRepr callback +
per-type vtable dispatch). Aligns with the `TVM_REGISTER_*`
family.
- Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`,
`relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor
of a generic `ffi::Map<String, Any> extra_config` keyed by
`"<dialect>.<knob>"`. Each call site reads via the templated
accessor `config->GetExtraConfig<T>("...", default)`.
- Promotes `std::string` config fields to `ffi::String`.
After this lands, the script-printer core knows nothing specific
about any dialect — new dialects plug in via the registry pattern
with zero core edits. Public Python API surface unchanged.
Prior to this PR, `python/tvm/script/__init__.py` imports both tir and relax
submodules. This leads to the phenomenum that when people does
```python
from tvm.script import tir as T
```
, the relax submodule will be implicitly visited by `__init__.py` as well.
Since TIR does not rely on Relax, it is good not to import both of them
at the same time. (This can prevent cyclic imports sometimes.)
This PR does this decoupling by introducing two files
* `python/tvm/script/relax.py`
* `python/tvm/script/tir.py`
and removing the imports from `python/tvm/script/__init__.py` and
`python/tvm/script/parser/__init__.py`. With this change, we force people to
manually do `from tvm.script import tir` and `from tvm.script import relax`
to use TVMScript parser, which is right our conventional way.
This PR refactors the TVMScript printer and includes the following
changes:
- Consolidate the logics of VarTable into IRDocsifier
- Decouple TracedObject into Object and ObjectPath for less syntactic
noise
- Restructure the folder to ensure logics and consistency
Some tests removed because the APIs do not exist any more due to the
consolidation.
This PR introduces some minor restructuring of the `python/tvm/script`
folder structure to make it more convenient for future upstreaming.
Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
Shushi Hong