## Motivation and context
A TIRx pointer carries two pieces of information that later lowering
needs: the pointee element type and the storage scope. Both must survive
when a pointer-producing expression is assigned to a Python name and
then used as the backing storage of a buffer.
A concrete example is accessing an mbarrier in another CTA through
distributed shared memory:
```python
ptr_ty = PointerType(PrimType("uint64"), "shared")
remote_ptr = T.reinterpret(
ptr_ty,
T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0)),
)
remote_mbar = T.decl_buffer(
[1], "uint64", data=remote_ptr, scope="shared"
)
```
`map_shared_rank` returns the raw `uint64` address produced by PTX
`mapa`, and `reinterpret` gives that address the intended
`PointerType(uint64, shared)`. Because `decl_buffer(data=...)` requires
a pointer `Var`, assigning the expression to `remote_ptr` should create
an immutable typed pointer binding.
Before this PR, an unannotated assignment such as `remote_ptr = <pointer
expression>` followed the same parser path as a numeric assignment. That
path allocates a mutable local scalar and therefore cannot represent a
`PointerType`. The pointer expression could not be carried as a
correctly typed `Var` into `decl_buffer` and CUDA lowering.
This PR makes an unannotated pointer-valued assignment emit a TIRx
`Bind`. The bound `Var` has exactly the type of the right-hand side,
including its element type and storage scope. Pointer bindings are
immutable, so reassignment in the same scope is diagnosed; shadowing a
name supplied through `extra_vars` remains valid. Numeric assignments
keep their existing mutable-local behavior.
## Type propagation fixes
The parser fix exposed several other boundaries where pointer type
information must remain consistent:
| Boundary | Previous behavior | Behavior after this PR |
| --- | --- | --- |
| Unannotated pointer assignment | Tried to materialize the value as a
local scalar | Emits an immutable `Bind` with the RHS `PointerType` |
| `address_of(buffer)` / `buffer.ptr_to(...)` | Reused the raw backing
pointer type | Returns a pointer to `buffer.dtype` while preserving the
backing pointer storage scope |
| `tvm_access_ptr` / `ptr_byte_offset` | Accepted strings or annotation
expressions, but not a `PrimType` object directly | Accepts `PrimType`
and produces the corresponding typed pointer |
| Printed `T.ptx.mapa` call | The printer emits all intrinsic attributes
positionally, but the Python helper required keyword-only arguments |
Accepts the canonical printed form so pointer code round-trips through
TVMScript |
The `address_of` distinction matters for typed views over byte-addressed
storage. For example, if a `float32` buffer is backed by a `uint8*`
allocation in `shared.dyn`, the address of a buffer element must be
`PointerType(float32, shared.dyn)`, not `PointerType(uint8,
shared.dyn)`.
With these changes, the DSMEM example above round-trips through
TVMScript and CUDA codegen declares the remote buffer pointer as
`uint64_t*`.
## TMA dtype normalization
This PR also contains a small, separate type-representation fix in TMA
descriptor construction. `TmaPlan.elem_dtype` is a string consumed by
the host-side `runtime.cuTensorMapEncodeTiled` packed call, but
`_assemble_plan` stored `g_buf.dtype`, which is a `PrimType`. Converting
it with `str(g_buf.dtype)` ensures that the generated packed-call
argument is `StringImm("float16")` rather than an IR type object. This
does not change the pointer-binding semantics described above.
## Testing
- Verify that an unannotated pointer expression creates a `Bind` whose
`Var` type matches the RHS type.
- Verify that pointer reassignment is rejected while shadowing an
`extra_vars` name is allowed.
- Verify parser/printer structural round-tripping for the pointer
binding and canonical `T.ptx.mapa` call.
- Verify that `address_of` uses the logical buffer element type and
preserves the storage scope for byte-backed buffer views.
- Verify that `tvm_access_ptr` and `ptr_byte_offset` accept `PrimType`
inputs.
- Compile the DSMEM `map_shared_rank` example through the CUDA TIRx
pipeline and check for a typed `uint64_t*` remote buffer pointer.
- Verify that the TMA host initialization passes the descriptor dtype as
a `StringImm`.
## Motivation and context
The TVMScript parser already tracks Python AST locations for
diagnostics, but TIRx statements and expression results emitted through
`IRBuilder` did not retain those locations. After parsing, a direct
intrinsic call, an inlined helper body, or a `TilePrimitiveCall`
therefore could not be traced back to the source range that produced it.
Inline expansion also needs more than a single flat location. The
generated IR should retain both the caller location and the
helper-definition location, while ordinary nested AST evaluation within
one source should not accumulate redundant enclosing spans.
## Changes
- Add an active source-span stack to `IRBuilder`, with scoped push/pop
support.
- Make the parser activate the current AST source range while visiting
statements and evaluating expressions.
- Attach the active span to emitted TIRx statements and to expression
results that do not already carry an explicit span.
- Normalize nested spans from the same source to the innermost relevant
range.
- Preserve cross-source inline expansion history as a `SequentialSpan`,
ordered from the call site to the expanded definition.
- Reuse the same source-coordinate calculation for diagnostics and IR
spans so their line and column conventions remain consistent.
Source spans remain diagnostic metadata: functions parsed from different
source locations keep the same structural hash and remain structurally
equal.
## Testing
- Verify exact parser source coordinates against diagnostic coordinates.
- Verify spans on direct intrinsic calls and `TilePrimitiveCall` nodes.
- Verify that inline expansion produces a `SequentialSpan` containing
caller and callee ranges.
- Verify direct `IRBuilder.with_source_span` behavior.
- Verify that source spans do not affect structural identity.
- Run changed-files pre-commit checks, including clang-format.
Focused result: 9 tests passed.
This PR adds an autoload mechanism for out-of-tree backends, simplifies
TVM's Python library loading, and removes `TVMError` in favor of native
Python errors.
## Autoload out-of-tree backends
Out-of-tree packages can register an autoload callable under the
`tvm.backends` entry-point group (mirroring torch's device-backend
autoload). At `import tvm` startup each entry point is discovered and
its callable invoked once, after the core runtime and the `tvm`
namespace are fully initialized, so an extension can register
ops/targets/funcs or load extra libraries.
```toml
[project.entry-points."tvm.backends"]
tvm_foo = "tvm_foo:_autoload"
```
A failing extension is caught and surfaced via `warnings.warn` so it
cannot break `import tvm`. Autoload can be disabled with
`TVM_DEVICE_BACKEND_AUTOLOAD=0`.
## Simplify library loading
The library-loading path in `base.py` is consolidated around a single
`_LOADED_LIBS` dict (basename to ctypes handle) so downstream and
autoloaded extensions can skip already-loaded libraries; the per-backend
runtime DSO list is folded into `load_backend_libs`. Accumulated cruft
is removed: the Python-3.9 check, the readline shim, the `_FFI_MODE`
ctypes check, the `base.__version__` re-export, and `py_str` (call sites
inline `.decode("utf-8")`).
## Remove TVMError in favor of native Python errors
`TVMError` added a layer atop `RuntimeError` that downstream code had to
import and learn. It is removed; the registered FFI error kinds
(`InternalError`, `RPCError`, `OpError`, `DiagnosticError`,
`ScheduleError`) now subclass `RuntimeError` directly while staying
registered, so the FFI keeps throwing the right kinds. All `TVMError`
imports, `except`/`raise`/`isinstance` uses, and
`pytest.raises(tvm.TVMError)` sites move to the `RuntimeError` builtin.
## 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
Rename `LetStmtNode`/`LetStmt` to `BindNode`/`Bind` and remove the
`body` field.
The variable defined by `Bind(var, value)` is now visible in all
subsequent
statements within the same enclosing scope, rather than being scoped to
a nested body.
This flattens deeply nested let-chains into sequential
`SeqStmt([Bind(...), Bind(...), ...])`,
making the IR easier to read, transform, and analyze.
## Key Changes
- **New `BindNode`**: `{var, value}` — no body field. Variable scope is
the enclosing
statement's body (For, IfThenElse, AllocBuffer, etc.)
- **ScopeStack pattern**: Passes that need scope-aware cleanup
(ConvertSSA, CSE,
tir_visitor_with_path) use `ScopeStack` instead of manual save/restore
or RAII wrappers
- **All passes migrated**: 89 files updated across codegen backends, TIR
transforms,
S-TIR transforms, analyses, TVMScript printer/parser/ir_builder
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.
* support continue and break in tvmscript
* fix black format
* fix pylint issue
* Update tests/python/tvmscript/test_tvmscript_syntax_sugar.py
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
* add printer/parser test, fix lint
* Fit to latest ffi update
* Skip i386 numpy-related test
* Introduce AnnotateIrregularLoop before any lowering loop expansions.
---------
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
### **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
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.
* [TVMScript] Enable T.macro decorateing class method
This PR refactors the implementation of `T.macro`, so that the `self` argument can be passed through the TVMScript parser. Then we can decroate the class methods with `T.macro`.
* update test
* [TVMScript] Avoid segfault from invalid TVMScript
Prior to this commit, after the `DiagnosticContext` prints its error,
it overwrites the `DiagnosticRenderer` with a NULL renderer. If a
second call to `DiagnosticContext::Render` occurs, it will segfault.
This appears to be intended to prevent double-printing of error
messages, but double-printing error messages is much worse than a
segfault.
In addition, `DiagnosticContext::Render` should only be called once.
There's a common pattern in the parser where it will wrap exceptions
in `DiagnosticError`, but re-raise exceptions that are already a
`DiagnosticError`. This requires every such location to include
`except DiagnosticError: raise`, and can easily be missed.
This PR makes two changes: First, the `DiagnosticRenderer` is updated
to have a no-op callback rather than a NULL callback. Second, the
re-raising of `DiagnosticError` is moved to `Parser.report_error`, so
that it does not need to be handled separately at several independent
locations in the TVMScript parser.
Prior to this commit, exceptions raised during the parsing of
TVMScript would be caught and replaced with a new exception. While
this does allow the TVMScript location of the error to be included in
the exception, it also removes the stack trace of the original error.
This commit updates the `Parser.report_error` function to provide the
original stack trace alongside the updated exception object.
* [TVMScript] Support starred indices in for-loop
An extension of https://github.com/apache/tvm/pull/15404, which
allowed starred expressions in the rhs of
assignments (e.g. `T.decl_buffer(shape=[*dim, 128])`), this PR also
enables starred expressions in the lhs of
assignments (e.g. `for *spatial,reduction in T.grid(*A.shape)`).
* Fix single-argument indices
* Updated test case for T.grid()
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.
This PR implements the privacy annotation proposal. Namely, the @R.function decorator now has an optional private attribute. If a function is marked as private, then it will not have a global symbol attached to it and thus will not be externally accessible. By default, functions are not private, so the parser does insert a global symbol for them.
* [TVMScript] Cherry-pick module.other_func syntax from unity
* [TIR][TVMScript] Added format/parsing of subroutine calls
Similar to `module.relax_func(args)` syntax used when parsing Relax
functions, this allows `module.tir_func(args)` to be used when parsing
TIR PrimFuncs.
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.
* [Bugfix][TVMScript] Handle LetStmt for `var1 = var2` expressions
Usually, when using TVMScript to represent a `PrimFunc` variable
definition `var_name = expr` defines `LetStmt` with a variable named
`var_name` bound to the expression `expr`. However, prior to this
commit, if `expr` is a `tir::Var`, the TVMScript parser would instead
silently omit the `LetStmt`, and rename all instances of that variable
to `var_name`.
The root cause was in the `VarTable.exist` check, which erroneously
returned False in all cases. This was due to a `value is v` check,
which checked if the value was the same as the stack of
maybe-shadowing values that share the same name. Replacing the
'value is v` check with a `value in v` check resolves this issue.
This bug dates to the initial implementation of the new TVMScript
parser in https://github.com/apache/tvm/pull/12496.
* Avoid implicit `PrimExpr.__bool__` from `if value in value_stack`
* Use T.meta_var where variable renaming is required.
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
This PR enables context-aware parsing for TVMScript. It means that the parser has full control of the statements in the specific context/namespace.
For example, we can override the global var `__call__` method in Relax function to make sure to generate Relax Calls instead of Relay Calls.
This PR tries to fix the crash of parser when the old value of a var is an array but the new value is not. For example:
```python
from tvm.script import tir as T
def func_wrapper(shape, dtype):
@T.prim_func
def test_case():
a = T.alloc_buffer(shape, dtype=dtype)
return test_case
if __name__ == "__main__":
a = np.zeros((10, 10), dtype="int8")
print(func_wrapper((256, 256), dtype="int8").script())
```
In the above code, there are two assignment to var 'a'. In the global scope, its value is a numpy array. But it is a Buffer in the prim function. There is a table named 'name2value' to track the value of vars like 'a' here.
When the parser wants to update its value, it will compare the value between the new and the old assignment. Here the problem comes. When we use '==' to compare an array with a value, the result is an array too, which can not be used as a condition of a if stmt directly. So, the code above will emit an error:
```shell
error: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
--> /workspace/code_newest/tvm/private_test/test_meta_programming.py:16:9
|
16 | a = T.alloc_buffer(shape, dtype=dtype)
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
```
This PR fixes this by change "==" to "is".
Co-authored-by: lightzhan-intellif <zhan.liang@intellif.com>
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>
Hongyi Jin