Replace TVM's `Diagnostic` / `DiagnosticContext` machinery with the
tvm-ffi
`visit_error_context` mechanism. Validators throw an `ffi::Error` seeded
with the
offending node; leaf pass executors (`ModulePass` / relax `Function` /
`DataflowBlock`) catch and rethrow `EnrichPassErrorWithContext`, which
appends the
failing pass name and a TVMScript-rendered, underlined source location.
`relax.analysis.well_formed` now throws on the first violation; a new
`check_well_formed` returns a bool, and all C++/Python/test callers are
routed
accordingly. `include/tvm/ir/diagnostic.h` and `src/ir/diagnostic.cc`
are deleted.
The enrichment renders with `num_context_lines=10` so a small function
shows in
full with no skipped-lines marker, while a large module stays bounded.
The TVMScript parser diagnostics
(`python/tvm/script/parser/core/diagnostics.py`)
stay self-contained pure-Python with no `DiagnosticContext` dependency,
and
restore multi-line source rendering: a diagnostic whose offending AST
node spans
multiple source lines now renders every spanned line with its gutter
line number
and an underline covering the span. `tvm.error.DiagnosticError` (used by
the
TVMScript parser) is retained.
A rendered end-to-end enriched-error example is posted as a comment
below.
## 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.
### **Overview**
This PR implements native Python function support in TVM Relax through
the `@I.pyfunc` decorator and `BasePyModule`, which enable seamless
integration between TVM's compilation pipeline and Python/PyTorch runtime
environments. This enhancement allows users to write Python functions
directly in TVMScript that can interoperate with Relax and TIR functions
that provides enhanced debugging capabilities and leveraging existing
PyTorch operator libraries.
### **Key Features**
**TVMScript Parser Enhancement**
- `@I.pyfunc` decorator: Marks Python functions for integration into IRModules
- Dual storage format: Stores both raw string representation (for TVMScript
printing) and captured PackedFunc (for runtime execution)
- ExternFunc representation: Each Python function is represented as an
ExternFunc node with attributes storing source code and runtime wrapper
**Complete BasePyModule Implementation**
- DLPack-based tensor conversion: Seamless conversion between PyTorch
tensors and TVM NDArrays
- Cross-function interoperability: Python functions can call Relax/TIR
functions and vice versa
- JIT compilation: Delays compilation until module instantiation for flexible
late-stage modifications
- Dynamic function registration: Supports runtime addition of Python functions
### Future Work
- TVMScript printer for IRModules with Python functions: Print IRModules
in proper format with high-level operator mapping from Relax ops to PyTorch
ops, handling symbolic shapes
- R.call_py_func primitive: Introduce Relax primitive to invoke corresponding
PackedFunc of specified Python functions at runtime
* [Analysis] Allow calls to GlobalVar in @R.function
Prior to this commit, the post-parsing well-formed check performed by
TVMScript allowed a call to `GlobalVar` in a `@R.function`, but only
if it occurred within the context of a `@I.ir_module`. If
`@R.function` appeared on its own, calls to a `GlobalVar` would be
treated as calls to an undefined function.
* Use approrpirate well-formed checks TIR/Relax functions
* Lint fix
* Import order fix
* Check well-formedness in the parser
* Correct packed funcs in NN frontend
* Support the check_well_formed optional argument to I.ir_module
* Also check well-formedness in TIR
* Enable normalization for individual Relax functions and PrimFuncs
* Use the error raised by the TIR well-formed checker for the message
* Fix tvmscript test failures
* Whitespace
* Fix errors in verify_well_formed test
* Include a more helpful error message
* Fix TIR test failures
* Address well-formed failures in test_tir_specialize
* Correct well-formedness error in test_tir_analysis_oob
* Correct further well-formedness failures
* Remove __tvm_meta__ from test case to avoid parsing error
* Avoid circular import in entryy.py
* Formatting fixes
* lint fix
* Add pylint exceptions
* Fix whitespace
* Fix more failed test cases
* Catch inappropriate use of decl_function instead of segfaulting
* Fix test_lower.py
* Mark purity in test_relax_2d_buffer_allocation.py
* Mark purity in test_dma_builtin.py
* Remove __tvm_meta___ from test_tir_usmp_analysis_extract_bufferinfo.py
* Suppress well-formed check in test_tir_transform_convert_blocks_to_opaque.py
* Remove __tvm_meta__ in test_tir_usmp_algo.py
* Remove __tvm_meta__ from more USMP tests
* Fix incorrect var in test_tir_transform_storage_flatten.py
* Remove all remaining instances of __tvm_meta__
* Fix purity error in test_dataflow_pattern.py
* Fix purity error in test_ast_printer
* Fix test_arith_domain_touched example
* Okay to set check_well_formed to True in test_tir_analysis_identify_mcmcpy
* Define variable in test_tir_analysis_oob
* Typo fix
* Add explanatory comment to test case
* Define the undefined vars in test_tir_transform_common_subexpr_elim
* Exception no longer necessary in test_tir_transform_inject_rolling_buffer
* Remove unnecessary check exemption in test_tir_transform_convert_ssa
* Avoid checking exemption in test_inject_ptx_ldg32
* Note special case in test_distributed_transform_propagate_sharding
* Exempt well-formed error in dlight/test_benchmark
* Exempt well-formedness errors in test_ethosu/, mostly uninitialized vars
* Whitespace
* Include non-CUDA GPUs in IsScheduledOnGPU
* Fix thread binding bug by changing thread binding var dtype
* Include overrides in test_runtime_builtin_paged_attention_kv_cache.py
* add exemptions in test_ethosu/test_replace_conv2d
* Add more ethosu exemptions
* More exemptions for ethosu tests
* Remove unused reference
* Indicate purity in test_transform_rewrite_cuda_graph
* Indicate purity in test_transform_normalize
* Reorder MergeSharedMemoryAllocations in GPU codegen
* Add target parameter for FP8StorageLegalize and FP8ComputeLegalize
* Don't re-import Target in tvm/tir/transform/transform.py
As a background info---the script parser works by visiting a "statement"
(or top-level expression) at a time. The expression parts of the state-
ment are evaluated, and then the IR corresponding to the statement is
constructed if necessary.
In TIR, macro calls can only occur at the statement level, and they don't
produce any values. This means that the statement visitor (visit_expr_stmt)
can see these calls directly in its node parameter. At this point it could
simply visit the body of the macro instead, which is the basis of the
existing implementation.
In other dialects there may be a need for macros to produce values. This
means that macro calls can occur in the middle of complex expressions.
As a result, these calls will not be present at the statement level, and
the TIR approach by intercepting them in visit_expr_stmt will no longer
work. Instead, these macros delay the visiting of the macro body to the
evaluation time. A macro is represented by an ScriptMacro (TIRMacro in
the current implementation) object (created via macro decorator). When the
evaluator evaluates an expression with a macro call, it will call the
macro object (since macro calls use function call syntax). It is in the
macro object's __call__ function where the macro parsing picks up. The
remaining issue was to pass the Parser object to the __call__ function.
This is done by injecting it into the global dictionary under a reserved
name.
It turns out that the same approach also works for TIR, and the macro
processing can be generalized, leaving only language-specific details to
the language-specific language macro objects.
* [TIR] Implement TIR macros
This patch introduces two new symbols: `T.macro` and `T.insert`.
`T.macro` is a decorator that, when applied to a function, turns the
body of that function into a piece of TIR that can be inserted via
`T.insert` into a PrimFunc.
For example:
```python
@T.macro
def copy_backwards(dst, src, size):
with T.block("backwards"):
for i in T.serial(size):
ai = T.axis.remap("S", [i])
T.reads(src[0:size])
T.writes(dst[0:size])
dst[ai] = src[size - ai - 1]
@T.prim_func
def foo_int32(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")):
T.insert(copy_backwards, A, B, 128)
@T.prim_func
def foo_int8(A: T.Buffer((128,), "int8"), B: T.Buffer((128,), "int8")):
T.insert(copy_backwards, A, B, 128)
```
The above will generate two PrimFuncs that do the same backwards copy,
but applied to buffers with different data types.
Semantics:
- Function that is decorated with @T.macro can have any parameters that
follow Python syntax, i.e. positional, keyword, etc. Type annotations
are not required, but are allowed.
- The arguments to `T.insert` are macro name followed by the argument
list.
For `T.insert(arg1, arg2, arg3, ...)`, the values are substituted into
the body of the macro as in the call `arg1(arg2, arg3, ...)`.
The body with the substituted values is then inserted at the point
where the `T.insert` is located.
* Fix linter
* Fix linter again
One linter suggested something that the other didn't like...
* Get rid of T.insert, apply macro via function-call syntax
* Store closure vars in TIRMacro
* ast.parse always returns ast.Module, hence doc is doc.Module
* Simplify `expand_macro`, capture environment variables
* Implement macro hygiene
* Fix linter
* Make T.macro work same as T.macro()
The previous commit inadvertently made T.macro (without parentheses)
illegal, only abbreviated form allowed was T.macro(). Restore T.macro
as a valid decorator use.
* Edit comment: insertion -> expansion
* Add import pytest
* One more typo...
* Remove stale testcase
Consider the case below:
```python
dtype = "float32"
@T.prim_func:
def f(
A: T.Buffer((1, ), dtype),
B: T.Buffer((1, ), dtype),
):
...
```
The variable `dtype` only appears in the type annotation of the function
being parsed. In this case, the python interpreter will evaluate the
annotation first before invoking the decorator, and thus if `dtype`
doesn't appear in the function body, it will not be considered as being
captured by the function itself. As a result, `inspect` module will be
unable to supply the value of `dtype` during parsing, leading to
failure.
This PR fixes the bug by maintaining a copy of function annotations
that are already parsed. Whenever expression evaluation fails during
parsing, it falls back to using the copy that is evaluated by python
interpreter.
Multi-line strings might make less sense to be printed out by default,
as they could be LLVM snippets, CUDA source code and anything hard to
comprehend but easy to mess up with the TVMScript itself. Therefore,
this PR is introduced to print them as metadata by default.
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>
Tianqi Chen