This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
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
`ReplaceGlobalVars` was a public IR-layer API with only one in-tree C++
caller (`relax::AttachGlobalSymbol`). The mechanism used a NodeFunctor
vtable populated at static-init time by per-dialect `.cc` files in
relax and tirx, which made the IR layer logically depend on its
dialects even though the include graph did not show it.
Move the dispatch logic into the consumer as file-local mutators and
a private helper. Delete the public header, the IR-layer driver, both
per-dialect dispatch registrations, the `IRModule.replace_global_vars`
python method, and its dedicated test file. The behavior is still
covered by `tests/python/relax/test_transform_attach_global_symbol.py`
and by the pipelines that include the `AttachGlobalSymbol` pass.
## Summary
TVM-side cleanup that drops the `python/tvm/runtime/object.py` shim and
routes `tvm.runtime.Object` directly to `tvm_ffi.Object`. The
`tvm.runtime.Object` re-export is preserved (now a re-export of
`tvm_ffi.Object`) so external callers keep working.
The load-bearing `__object_repr__` install — which wires TVM IR objects
up to the rich C++ `ReprPrinter` registered through
`init_ffi_api("node", ...)` — moves into
`python/tvm/runtime/_ffi_node_api.py`.
That module is already imported as a side-effect-only module from
`python/tvm/runtime/__init__.py`, so the override fires at the right
time (after `init_ffi_api` registers the C++ printer).
`_ffi_node_api.AsRepr` itself is **kept**: `tvm_ffi`'s default repr is
primitive (`ClassName(ptr)`); TVM IR objects need the rich printer
registered via `init_ffi_api("node", ...)`. `AsRepr` is what bridges
that printer back into Python `repr(obj)` and is also the runtime-only
fallback when `libtvm.so` is unavailable.
The 7 in-tree importers of the deleted shim (plus one straggler in
`runtime/disco/session.py`) are switched to either
`from tvm.runtime import Object` or `from tvm_ffi import Object`,
depending on which pattern the file already uses.
## Test plan
- [x] `python -c "import tvm; print(repr(tvm.IRModule({})))"` produces
TVMScript-style output (rich repr preserved).
- [x] `pytest tests/python/all-platform-minimal-test/ -x` — 75 passed,
77 skipped (matches baseline).
- [x] `pytest tests/python/tirx-base/ -x` — 273 passed, 2 skipped.
- [x] `pre-commit run --files <changed files>` — all hooks pass.
- [ ] CI green.
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
* [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
This PR formalizes the namespace for all object registered so
we do not have object that sits on root namespace
Also fixes the Visitor style in TensorMapNode
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.
* cleanup relay c++
* [REFACTOR] Phase out relay c++ components
This PR phases out the relay C++ components and
simplifies the overall codegen runtime logic.
---------
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
This PR starts the step 0 to phase out relay from the current
development main branch. This PR focuses on the python
components of relay, autotvm, auto_scheduler. To make the change
manageable, we will also do followup steps on te.Schedule and
c++ components in followup PRs.
To continue support community members who depends on
legacy flows, the [v0.19.0](https://github.com/apache/tvm/tree/v0.19.0)
branch will continue contain these components.
As noted in [discussion on phasing out legacy components](https://discuss.tvm.apache.org/t/phasing-out-legacy-components/17703/30),
this would help us to do two purposes:
- By removing outdated or redundant elements, we can significantly
reduce complexity and improve maintainability.
- Unify our focus: Concentrating our efforts on the new unity flow
will allow for more efficient development and innovation.
It is also a good opportunity for us to revisit and reduce CI time.
The past relay legacy flow contains a lot of end to end tests that
requires hardware resources to run and causing long CI time.
Moving onwards, we can focus more on unit-tests that focuses
on structural equality and runs within seconds, while be mindful
about tests that requires hardware resources (by restricting them
to specific folders and CI nightly in some cases).
---
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
* [IR] Expose ReplaceGlobalVars utility in the Python API
This is a follow-up PR to https://github.com/apache/tvm/pull/17202,
which added a general utility to replace `GlobalVar` instances across
all TVM IR types. This PR exposes this new utility through the Python
API, and explicitly tests its functionality.
* Lint fix
* [IR] Add utility methods to IRModule
* `IRModule.clone`: Clone the module. While in C++, a module can be
copied using `IRModule::CopyOnWrite()`, copying a module in Python
required passing all members into the `IRModule` initializer. The
`IRModule.clone` method provides an easier way to copy an `IRModule`
from python.
* `IRModule.__delitem__`: Remove a function from the module. This
exposes the C++ method `IRModuleNode::Remove` for use in the python
API. This uses the python `del` keyword, similar to a native python
list. Similar to the existing `IRModule.__getitem__`, this can be
called with either a `GlobalVar` or a python string.
* `IRModule.__contains__`: Check if a function is in the module. This
allows the pythone keyword `in` to check if a module contains a
specific function. Similar to the existing `IRModule.__getitem__`,
this can be called either with a `GlobalVar` (`if gvar in mod`) or
with a python string (`if "function_name" in mod`).
* [Unity][CodeGen] RunCodegen based on externally-exposed functions
Prior to this commit, `relax.transform.RunCodegen` required a list of
entry functions for a module, defaulting to `"main"` if not specified.
The list of entry functions is duplicate information that could be
inferred from the module, and should not be required from the user.
This commit updates `RunCodegen` to treat all externally-exposed
functions as entry points, in the same manner as
`DeadCodeElimination`.
For backwards compatibility, the `entry_functions` argument is still
accepted, and is used to augment the list of externally-exposed
functions.
Prior to this PR, visiting order of the following for-loop is
non-deterministic:
```python
mod: tvm.ir.IRModule
for gv, func in mod.functions.items():
...
```
This is because `IRModule` stores those functions inside a hash map
`Map<GlobalVar, BaseFunc>`, which is based on pointer equality.
This behavior is usually innocent and harmless given many previous
workloads only have one "main" function as the primary entry point, e.g.
a single Relax Function and a bunch of TIR functions. However, it is not
true in LLM usecases where `prefill`, `decoding` are both equally
important entrypoints, and in those cases, it is possible that the
names of generated TIR functions from LegalizeOps are essentially
different from each run, which is, again, harmless in basic usecases,
but it does make debugging more challenging.
This PR corrects this behavior by sorting the functions alphabetically
by their names.
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 gradually migrates more pieces of the default printing to
TVMScript printer for TIR.
This PR gradually migrates more pieces of the default printing to
TVMScript printer for TIR. Details:
- Introduced a method `AsLegacyRepr` which preserves existing
`AsRepr` provided by `ReprPrinter`, so that the legacy behavior
could be 100% preserved.
- Introduced `Script` method to `IRModule`, `PrimFunc`, `tir.Stmt`,
`tir.PrimExpr`. The `script` method exists in python side before,
and this PR introduced them to C++ to be consistent.
- Replace TIR's `PrettyPrint` to `operator <<` that is provided by
the new `ReprPrinter`, which outputs in TVMScript format by default.
`PrettyPrint` on Relay is all preserved for backward compatibility.
* [TVMScript] Improvements tvm.script.highlight
- Automatically use "black" formatter if available.
- Allow overrides of pygmentize style based on environment variable
`TVM_PYGMENTIZE_STYLE`.
- Forwarded `black_format` argument from `show` method to `cprint`
* feat(ux): highlight tvm script
* resolve dependency
* refact(tvmscript): highlight fallback as plain text with warning; put ansi colors as default terminal style
* refact(tvmscript): make terminal style close to the default notebook style
* fix(pylint): disable=import-outside-toplevel
* fix(ci-dep): Pygments>=2.4.0 to support ansicolors w/o #
* refact: making Pygments versioning most robust and user-friendly
* fix: pylint var naming
* [TVMC][Relay] Introduce executor and runtime parameters
This introduces `executor` and `runtime` into the various entrypoints but also into `tvmc` as `--executor` and `--runtime`. This touchs a lot of files and I've tried to update anywhere as necessary.
Notable, executor code generators now accept the initial `IRModule` rather than creating
it themselves so it can be annotated once.
Validated the demo application continues to classify the tabby cat with
new CLI options.
* Correct Graph Executor Python API
* Refactor the type checker to use diagnostics
Although this patch is very large and seemingly disjoint the
fixes are required to get it working for the entire stack.
I started with first changing InferType to use the diagnostics,
these weren't yet in the pass manager so this required changes
to module and module pass. InferType wasn't actually written
correctly as a pass requring refactoring there, then in order
to add spans to AST it required turning on AnnotateSpans which
in term required changes to the parser, and module to make
it possible to use the errors. These changes to parse and module
required changes to diagnostics and InferType. Althought seemingly
disconnected there are hidden cycles between the components which
require simultaneous change in order to remove the old error
reporting.
A huge change due to this patch is that the module no longer
implicitly type checks functions which are added.
* Apply suggestions from code review
Co-authored-by: Robert Kimball <bobkimball@gmail.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
* Apply suggestions from code review
Co-authored-by: Tristan Konolige <tristan.konolige@gmail.com>
* Clean up parser
* CR feedback
* Apply Bobs suggestions
* Fix up Python interface for diagnostics
* Fix test_ir_parser and formatting
* Fix cpplint
* Fix lint
* Fix format
* More lint
* Fix format
* Kill dead doc comment
* Fix documentation comment
* Rebase fixups
* Add docs for type.h
* Fix parser.cc
* Fix unittests
* Fix black
* Skip previously typechecked functions
* fix ACL
* Fix numerous issues
* Add repr method
* Fix issue with Pytest, I am ready to cry
* Fix the rest of tests
* Kill dead code
* Fix dignostic tests
* Fix more tests
* fix more tests (#11)
* Fix diagnostic.py deinit bug
* Fix deinit issue
* Format
* Tweak disabling of override
* Format
* Fix BYOC
* Fix TensorArray stuff
* Fix PyTorch
* Format
* Format
Co-authored-by: Robert Kimball <bobkimball@gmail.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Tristan Konolige <tristan.konolige@gmail.com>
Co-authored-by: Cody Yu <comaniac0422@gmail.com>
Co-authored-by: Zhi <5145158+zhiics@users.noreply.github.com>
* [REFACTOR][TIR] Migrate most of low-level build to use the Pass Manager.
- SplitHostDevice
- ThreadSync
- BindDevice
- LowerThreadAllreduce
- Provide a temp fix for printing IRModule with PrimFunc before the formal text printer.
* Address comments, fix tests.
* Fix relay tests
* Explicit move
Tianqi Chen