## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
## Background
The `tvm::ir` layer previously had a reverse dependency on
`tvm::script`, injected via the `TVM_OBJECT_ENABLE_SCRIPT_PRINTER()`
macro that added a `Script()` member method to IR node types (IRModule,
PrimExpr, Buffer, PrimFunc, Stmt). This violated the intended one-way
dependency: `script` should depend on `ir`, never the other way around.
Additionally, `PrinterConfigNode` accumulated dialect-specific fields
(`tir_prefix`, `tir_import_module`, `tirx_prefix`, `relax_prefix`) that
created leakage between the generic printer infrastructure and dialect
internals.
## Changes
This PR restores the clean dependency direction and encapsulates dialect
config properly, in 5 commits:
1. **Lift TVMScript entry point into `script/printer/printer.h`**: New
header `include/tvm/script/printer/printer.h` introduces:
- `tvm::Script()` free function replacing `TVMScriptPrinter::Script()`
static method
- `TVMScriptPrinter` class with vtable (`NodeFunctor<std::string(...)>`)
- `TVM_REGISTER_SCRIPT_AS_REPR` macro for registering per-type repr
callbacks
2. **Drop `TVM_OBJECT_ENABLE_SCRIPT_PRINTER` macro**: Remove the macro
from all IR headers (`ir/expr.h`, `ir/module.h`, `tirx/buffer.h`,
`tirx/function.h`, `tirx/stmt.h`), eliminating the reverse `ir` →
`script` dependency. All call sites of `.Script()` member methods
updated to use `tvm::Script()`.
3. **Move dialect-specific `PrinterConfig` fields to `extra_config`**:
Remove `tir_prefix`, `tir_import_module`, `tirx_prefix`, `relax_prefix`
from `PrinterConfigNode`. Dialect internals now read their config via
`GetExtraConfig<T>(key, fallback)` with dotted keys (e.g.,
`"tirx.prefix"`). `buffer_dtype` is kept as a top-level field alongside
`int_dtype`/`float_dtype` since it is a shared scalar-literal default,
not a dialect-specific knob.
4. **Python: drop dialect kwargs, expose `extra_config`**: Update
`PrinterConfig`, `Scriptable.script()`, `Scriptable.show()`,
`Scriptable._relax_script()`, and `BasePyModule.script()` to use
`extra_config: dict | None = None` instead of individual dialect kwargs.
The tirx auto-switch logic is preserved.
5. **Fix transitive include breakage**: Explicitly add direct includes
for `config.h` and `node_functor.h` where headers previously relied on
transitive paths through `expr.h`/`module.h`.
## Testing
- C++ unit tests: 118/118 pass
- TVMScript printer tests: 771 passed, 1 skipped, 1 xfailed
- TIR namespace tests
(`tests/python/tirx/test_printer_tir_namespaces.py`): 13/13 pass
- Relax AST printer tests: 24/24 pass
- Minimal platform tests: 37/37 pass
- Pre-commit (ASF headers, ruff, clang-format): all clean
## Summary
This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.
The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.
## Main Changes
- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.
## Validation
- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
- `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
- `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
- `664 passed, 25 skipped, 9 xfailed, 1 xpassed`
## Local CI Notes
Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
## 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.
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
* [FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This PR establishes tvm_ffi as a standalone python module.
The ffi is structured as a minimal pip module that can be
directly install by path or url.
examples/get_started provided a minimal example.
This is a major change as we are decoupling tvm_ffi as a
separate package, users need to install tvm_ffi separately.
Thanks to its minimal dependency, tvm_ffi can be easily installed
even just from the source by pip install ./ffi
This change would enable future improvement for library plugins
to have lightweight dependencies by just working on top of
the tvm_ffi, while the main compiler toolchain and runtime
can be layered on top.
* [FFI] Improve traceback setups
This PR improves traceback related setups
This PR formalizes the namespace for all object registered so
we do not have object that sits on root namespace
Also fixes the Visitor style in TensorMapNode
This PR phases out tvm._ffi redirections in favor of new FFI
new functions are now called via tvm.ffi.
We also enabled limited API support for python 3.12+
so the compiled binary can be forward compatible to future
python versions.
* [Unity][TVMScript] Optionally hide StructInfo that can be inferred
By default, TVMScript prints the struct info of every variable being
bound, which can become quite verbose. This commit adds the
configuration option `show_inferable_type_annotations`, which
determines whether struct info annotations are shown in cases where
they can be inferred.
The `show_inferable_type_annotations` option defaults to `True`,
preserving the current default behavior.
* Rename show_inferable_type_annotations to show_all_struct_info
* Add unit test for round-trip of opaque function
Prior to this commit, the default behavior of the `black_format`
argument in TVMScript printing has changed back and forth, based on
conflicting user preferences. This commit allows the default to be
specified by each using using the `TVM_BLACK_FORMAT` environment
variable. If unspecified in a `obj.show()` method call, this
environment variable is used to determine the default.
The new feature of printing object addresses is introduced in #15579.
It's an excellent feature for debugging, but it should not be turned on
by default.
* [TVMScript] Optionally output the address as part of variable names
When debugging IRModule transformations, it can be useful to
know the exact C++ address of an object. For example, to determine
whether an undefined TIR variable was caused by erroneous insertion of
a new variable, or from failing to remove a previous variable. While
TVMScript does de-duplicate all names within a single print statement,
there isn't a convenient way to identify which variable in TVMScript
corresponds to a specific variable in C++ logging statements.
This commit adds `show_object_address` to the `PrinterConfig`. If
false (the default), no change is made to the TVMScript variable
names. If true, the address of the C++ object is appended to the
variable name. For example, printing a `tir.Var('my_name','int64')`
as `"my_name_0x1234abcd"`.
* Updated unit test to target TVM main
The `private=True` branch is only on the unity branch.
The current cross-function calls in TVMScript will cause PyLint warnings,
since the GlobalVar will be marked as undefined vars, e.g.:
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
gv0 = R.call_tir(tir_func, x, R.Tensor((128,), dtype="float32")) # <= `tir_func` is not defined in Python syntax.
return gv0
```
This PR changes the behavior into `TestModule.tir_func` instead of direct `tir_func`
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
cls = TestModule # Use `cls` to refer the current Module
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
return gv0
```
NOTE: It's a breaking change, the old style is deprecated.
Additionally, this PR contains the following minor fixes:
- mark `R.function` as staticmethod as what we do for `T.prim_func`
- make `I`, `R`, `T`, `cls` be the builtin keywords for the printer
- define names for functions, modules to prevent naming conflict
- checking the var names is valid via regex expression
- fix typos
Remove relax prefix for now
This PR cleans up relax prefix in printer for now.
While these setups are useful and do not cause any
technical debts in the codebase. We remove it given requests.
They can be added back to unity branch and later as part
of upstream
This PR adds the functionality to roundtrip metadata during printing.
Users may turn on the flag below to allow the printer to dump metadata to screen.
```python
ir_node.show(show_meta=True, ...)
```
This PR makes the following changes over TVMScript printer:
1. Apply the new TVMScript printer to report and render the schedule error, and keeps effect same as old `AsTVMScriptWithDiagnostic`.
2. Introduce more annotating and underlining interfaces for TVMScript ptrinter
- `obj_to_annotate`: add comments to the final `StmtDoc` from `ObjectRef`
- `obj_to_underline`: underline the final `Doc` from `ObjectRef`
- `path_to_annotate`: add comments to the final `StmtDoc` from `ObjectPath`
3. Beautify the underline logic by introducing the `underline_exempted`, to exempt some useless underlines, e.g. indent, comments and doc string.
demo:
1. `*_to_annotate`: for function:
```python
@T.prim_func
def _func():
T.evaluate(0)
T.evaluate(1)
T.evaluate(2)
T.evaluate(3)
T.evaluate(4)
T.evaluate(5)
T.evaluate(6)
T.evaluate(7)
```
both the results of
```python
_func.script(
path_to_annotate={
ObjectPath.root().attr("body").attr("seq").array_index(1): "annotation 1",
ObjectPath.root().attr("body").attr("seq").array_index(3): "annotation 3",
ObjectPath.root().attr("body").attr("seq").array_index(5): "annotation 5",
ObjectPath.root().attr("body").attr("seq").array_index(7): "annotation 7",
}
)
```
and
```python
_func.script(
obj_to_annotate={
_func.body.seq[1]: "annotation 1",
_func.body.seq[3]: "annotation 3",
_func.body.seq[5]: "annotation 5",
_func.body.seq[7]: "annotation 7",
}
)
```
are
```python
# from tvm.script import tir as T
@T.prim_func
def main():
T.evaluate(0)
T.evaluate(1) # annotation 1
T.evaluate(2)
T.evaluate(3) # annotation 3
T.evaluate(4)
T.evaluate(5) # annotation 5
T.evaluate(6)
T.evaluate(7) # annotation 7
```
2. `obj_to_underline`: for function
```python
@T.prim_func
def func(a: T.int32, b: T.int32):
T.evaluate(a)
T.evaluate(b)
T.evaluate(a)
T.evaluate(b)
T.evaluate(a)
T.evaluate(b)
```
the result of `func.script(obj_to_underline=[func.params[0]])` is
```python
# from tvm.script import tir as T
@T.prim_func
def main(a: T.int32, b: T.int32):
T.evaluate(a)
^
T.evaluate(b)
T.evaluate(a)
^
T.evaluate(b)
T.evaluate(a)
^
T.evaluate(b)
```
This PR refactors the output of `assert_structural_equal`. Different from the directly printing mismatching nodes, in the old version, the improved one will print the whole scripts, with mismatching nodes underlined. And we print the `ObjectPath` to the mismatching nodes for further better debug. For example, we have following functions
```python
@T.prim_func
def func1(a: T.handle, b: T.handle):
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@T.prim_func
def func2(a: T.handle, b: T.handle):
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 256))
```
the log of `assert_structural_equal(func1, func2)` will be like
```python
ValueError: StructuralEqual check failed, caused by lhs at <root>.buffer_map[b].shape[1].value:
# from tvm.script import tir as T
@T.prim_func
def main(a: T.handle, b: T.handle):
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
^^^
T.evaluate(0)
and rhs at <root>.buffer_map[b].shape[1].value:
# from tvm.script import tir as T
@T.prim_func
def main(a: T.handle, b: T.handle):
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 256))
^^^
T.evaluate(0)
```
instead of
```python
ValueError: StructuralEqual check failed, caused by lhs:
128
and rhs:
256
```
which is not readable sometimes.
This PR introduces `PrinterConfig`, a systematic way to configure
TVMScript printer without having to set global flags.
This PR enables more customization of printer behavior. More
specifically, now any TVM’s object in python, as long as it
inherits from `Scriptable`, it automatically gains two methods:
- `.script(tir_prefix=...)`
- `.show(...)`
Tianqi Chen