This PR introduces Relax AnyType as the primary top/base type spelling,
replacing the previous ObjectType naming for the type that represents
any Relax value.
Changes:
- Add AnyType/AnyTypeNode with relax.AnyType registration and keep
ObjectType/R.Object compatibility aliases.
- Update Relax type analysis, type visitors, opaque function defaults,
and script printer/parser handling to use AnyType/R.Any.
- Migrate affected Python/C++ call sites, docs, and focused tests to the
new spelling.
Validation:
- cmake --build build --parallel 16
- Focused Relax/TVMScript pytest: 704 passed, 1 xfailed
- pre_commit run --files <changed files>
## 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
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
Python callers should reach the canonical tvm-ffi structural helpers
directly instead of going through a TVM-side redirect layer. This makes
the public tvm.ir bindings exact aliases of the tvm_ffi APIs and exposes
get_first_structural_mismatch from tvm.ir.
Main changes:
- Import structural_equal, get_first_structural_mismatch, and
structural_hash directly from tvm_ffi
- Remove the pure wrappers from tvm.ir.base while keeping
assert_structural_equal's TVM-specific formatting
- Update mismatch tests and add identity coverage for the direct
bindings
## 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
## Summary
This PR introduces `AllocBufferNode`/`AllocBuffer` as a single TIR
statement that both allocates memory and declares a buffer into scope.
This replaces the previous pattern of `Allocate(var, dtype, shape, cond,
DeclBuffer(buf, body))` with the simpler `AllocBuffer(buf, body)`.
### Main changes
- **New IR node** `AllocBufferNode` with fields `{buffer, annotations,
body}` — same semantics as `DeclBuffer` but also allocates memory
- **TVMScript**: `T.alloc_buffer(shape, dtype, scope)` now emits
`AllocBuffer` directly (statement-level allocation).
`T.sblock_alloc_buffer(...)` for SBlock-level buffer allocation (full
parameter set)
- **All codegen backends** (C, CUDA, Metal, OpenCL, WebGPU, LLVM, NVPTX,
AMDGPU, SPIR-V) updated to handle `AllocBufferNode`
- **All TIR transforms** (storage_rewrite, flatten_buffer,
vectorize_loop, lower_warp_memory, etc.) updated
- **All S-TIR transforms** (compact_buffer_region, merge_shared_memory,
inject_double_buffer, etc.) updated
- **Removed `AllocateNode`** entirely — `AllocBuffer` is now the sole
allocation primitive
- **Removed `AllocDescriptor`** from merge_shared_memory_allocations —
uses `Buffer` objects directly
- **Added `AllocBuffer::ConstantAllocationSize()`** inline helper method
### Design rationale
The old `Allocate + DeclBuffer` pair was a historical artifact:
`AllocateNode` stored raw fields (`buffer_var`, `dtype`, `extents`,
`condition`) separate from the `Buffer` object, requiring pattern
matching (`IsAllocateDeclBufferPattern`) to reconstruct the buffer
association. `AllocBuffer` unifies this into a single node with a proper
`Buffer` reference, simplifying codegen backends and transform passes.
225 files changed, ~3500 insertions/deletions (net near-zero, mostly
mechanical migration).
## Test plan
- [x] All TIR base tests pass
- [x] All TIR transform tests pass
- [x] TVMScript roundtrip tests pass
- [x] S-TIR transform tests pass
- [x] Codegen tests pass
- [x] All-platform minimal tests pass
- [x] C++ functor tests pass
- [x] Pre-commit clean (clang-format, ruff, etc.)
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 renames tir.Block to SBlock. This clearly indicate the
scheduable property of the block and is a prereq for followup stir
passes refactor.
Main changes:
- Data structure change from Block to SBlock
- Syntax change from T.block to T.sblock
This PR modernizes the FFI foundation of the project and introduce
a new minimal and lightweight module [tvm ffi](https://github.com/apache/tvm/tree/refactor-s3/ffi)
based on our lessons in the past few years. It implements a modern
version of the [Unified Packed and Object RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0097-unify-packed-and-object.md)
that unifies the packed function call and object systems.
Summary of the change:
- A dedicated clean Any/AnyView that can store strong and weak
references of items
- Function(previously PackedFunc) system built on top of the Any/AnyView
- A minimal C API that backs the overall calls. We are stabilizing the
API with a goal to bring clean, stable FFI conventions for both compiled
and registered code
- A rewrite of core python binding and generated code based on the module
- Update existing code and test cases to the new module
- Latest dlpack support
The new module brings many benefits thanks to the cleaner design,
to name a few:
- Any can support both POD types(int) and object types.
- Containers (e.g. Array) can now also contain Any value, e.g. now
`Array<int>` is supported, no need for boxed types
- Error handling now upgrades to object-based, allowing cleaner
traceback across languages
- Map now preserves insertion orders
- Path toward isolated stabilize minimum core ABI/API foundation module
- Type traits based design that cleanly defines how values interact
with Any system
- Automatic conversion of different types based on traits if needed
Because FFI upgrade is at heart of the project, the change touches every
component of the system. Importantly, this is an upgrade of the ABI so the
change is not backward compatible. The code compiled under the old
FFI won't work under the new one. We did provide example ABI translation
(e.g. LegacyTVMArgValueToFFIAny) functions for compatibility.
The PR tries to leave files in their old places while creating redirections.
The goal is to have the first milestone landed and infrastructure in place,
so we can do further refactors to complete features and cleanup legacy code
as trackable PRs. As of now, python binding and compiled code are under the
new convention while RPC and some other bindings still relies on legacy ABI
translation. We will work on upgrades in the coming PRs, including areas such
as reflection, phasing out legacy redirections etc.
Prior to this commit, the Relax well-formed checker validated
arguments provided to Relax functions, but did not validate arguments
provided to `R.call_tir`. As a result, incorrect arguments from Relax
to TIR would not be checked until runtime, if at all.
This commit updates the well-formed checker to verify that
`R.call_tir` has received the correct arguments, and has the correct
output shape specified in the `out_sinfo` parameter.
Initial implementation performed the validation as part of
`FNormalize`, to maximize coverage of this check. This increased
end-to-end compilation time by ~10%, and so the check was requested to
be restricted to the well-formed checker. Expensive operator-specific
validation is now performed in the new `FValidate` attribute.
* Revert "Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool" (#17252)"
This reverts commit 11be832620.
* [FFI] Re-introduce the boxed primitive values
Initially introduced in https://github.com/apache/tvm/pull/16183,
these changes were reverted in
https://github.com/apache/tvm/pull/17252 due to performance
degredation in some Relax models. This could occur when a model
contained a large number of calls to `"vm.builtin.tuple_getitem"`,
which may occur when model weights are provided as a tuple.
This PR re-applies the changes from
https://github.com/apache/tvm/pull/16183, but with the performance
degredation resolved. The root cause was unnecessary type-checking
when converting from an untyped `tvm::ArrayNode*` to the typed
`tvm::Array<T>`, in the case where `T` is `ObjectRef`.
* Correct typo from T to U
* [Container] Support non-nullable types in Array::Map
Prior to this commit, the `Array::Map` member function could only be
applied to nullable object types. This was due to the internal use of
`U()` as the default value for initializing the output `ArrayNode`, where
`U` is the return type of the mapping function. This default
constructor is only available for nullable types, and would result in
a compile-time failure for non-nullable types.
This commit replaces `U()` with `ObjectRef()` in `Array::Map`,
removing this limitation. Since all items in the output array are
overwritten before returning to the calling scope, initializing the
output array with `ObjectRef()` does not violate type safety.
* [FFI] Separate runtime types from IR types for int/float/bool
Prior to this commit, `int`, `float`, and `bool` arguments from Python
were converted to `IntImm`, `FloatImm`, and `Bool`. These are
subtypes of `PrimExpr`, and should only be used at compile-time. By
automatically applying this conversion as part of the FFI, these types
are required to be present whenever a primitive is converted to a
`tvm::ObjectRef`.
This can become especially fragile for an end-user when storing
objects into a TVM container. Because TVM containers require all
contents to be `ObjectRef` subclasses, an automatic conversion may be
applied on storing into a container, resulting in an unexpected type
being retrieved from the container. For example, this currently
occurs in Relax when extracting a `R.Prim` from a `R.Tuple`.
This commit introduces a `Box<T>` type for storage of boxed primitives
at runtime, distinct from the IR types.
* Primitive arguments provided to a PackedFunc that requires an
`ObjectRef` will be converted to the corresponding boxed type.
(e.g. Passing a Python `int` to a C++ function accepting `ObjectRef`
produces a `Box<int64_t>`.
* Boxed primitives provided to a PackedFunc that requires an unboxed
primitive will be converted to the corresponding primitive.
* PackedFunc return values of `ObjectRef` are converted to the
corresponding primitive, if present. (e.g. If a `tuple_getitem`
with static return type `ObjectRef` returns a `Box<int64_t>`, it
will be unwrapped to a python `int`.)
Together, these three rules provide backwards compatibility for
existing PackedFunc definitions, while avoiding exposing the user to
any container-induced type conversions betweeen primitive types and
`ObjectRef`.
* Fix unit test failure after merge
* Fix breakage in new unit test
* 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
Tianqi Chen