## 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.
* [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 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 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 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:
- 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