## 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.
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.
This PR cleans up the python API to make things more consistent
with existing python array api and torch.
Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device
API updates:
- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func
* [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
[REFACTOR] Phase out getattr based attribute handling
This PR phases out getattar based attribute handling as they are slower
and introduces extra code path.
This does mean that if an Object is not explicitly registered
in python side, we will no longer be able to access the field by name.
Likely this is also desirable as we would like to enable faster use that
updates the python end and do not rely on these behavior.
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.
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 introduces the `CommentDoc` for comments printing and `DocStringDoc` for docstring printing. It enables to add free comments and docstring as `stmt` in printing, e.g.
```python
# comment 1
# comment 2
"""
docstring 1
docstring 2
"""
```
The free here means to not be bound to any `stmt`, but acts as a single `stmt`, similar to `ExprStmtDoc` for `ExprDoc`.
This PR also introduces an example for the `CommentDoc`, as follow up of #13819.
In the old printer, we always print a `# with T.block("root"):`, when there is an implicit root block skipped when printing. For example,
```
@T.prim_func
def main():
# with T.block("root"):
a = T.alloc_buffer((128, 128))
for i, j in T.grid(128, 128):
with T.block(""):
...
```
We bring this syntax reminder back in this PR.
In addition, we introduce a field of `ir_usage` and `print_headers` into the printer configuration, to support the printing of headers for `IRModule` and `PrimFunc`. For example,
```python
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module():
@T.prim_func
def func():
...
```
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(...)`
This PR adds ObjectPath to LiteralDoc to allow integer/float/string/...
literals to have their own object path. This is a final preparation
towards structural error rendering when SEqual fails.
This PR introduces support for TIR fragment printing.
Fragment printing makes it possible to print TIR fragments in the text
format consistency with TVMScript PrimFunc/IRModule printing.
This PR still preserves the legacy ReprPrinter format by introducing an
API `LegacyTIRPrint` for TIR PrimExpr. This method is used in
AutoScheduler and TIR CSE for full backward compatibility.
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 adds boolean operators to OperationDoc. This is needed by the TIR expression printing because it has `tir::And` and `tir::Or`.
Tracking issue: #11912
This PR:
- Adds an entry point for the TVMScript Unified Printer
- Adds a helper object class `RootNodeContainer` to provide an injection point for the actual printer implementation to add specialized logic on the root node to print.
Tracking issue: https://github.com/apache/tvm/issues/11912
This PR:
- Adds VarTable for the new TVMScript Printer
Compared to the prototype version, this:
- Removes unnecessary public methods.
- GetObjectName
- GetUniqueName
- Add Frame parameter for `Define` methods. VarTable will add callback to Frame to remove variable when Frame exits.
- Changes DocFactory from `ExprDoc(ObjectPath)` to `ExprDoc()` to simplify var definition.
Tracking issue: https://github.com/apache/tvm/issues/11912
This PR:
- Implement Frame for the TVMScript Unified Printer
Compared to the prototype version, this:
- Removes the dependency of VarTable (SymbolTable) from Frame
- Adds a callback array to the Frame base class so that VarTable can add callback to clean variable when Frame goes out scope
Tracking issue: https://github.com/apache/tvm/issues/11912
This adds an ability to print a "diagnostic marker" based on a given ObjectPath. For example, say we are printing a fragment of TIR like
```
for i in T.serial(10):
a[i] = 5
```
and we would like bring the user's attention to the bound of the loop:
```
for i in T.serial(10):
^^
a[i] = 5
```
In this case we would give the doc printer an object path that represents this loop bound, i.e. something like `path_to_underline=ObjectPath.root().attr("extent")`
Tracking issue: https://github.com/apache/tvm/issues/11912
This PR:
- Add the source_paths attribute to Doc base class.
- Add the corresponding Python binding for it.
This PR is depended by multiple tasks, including the diagnostic output in DocPrinter, VarTable and IRDocisifer.
Tracking issue: https://github.com/apache/tvm/issues/11912
Co-authored-by: Greg Bonik <gbonik@octoml.ai>
This PR addes:
- All ExprDoc subclasses
- Their Python bindings
- Support of ExprDoc in PythonDocPrinter
- Unit tests for ExprDoc in PythonDocPrinter
Tracking issue: https://github.com/apache/tvm/issues/11912
Bohan Hou