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.
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 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 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.
When processing matmul with transpose (permute_dims operator), the current gradient pass will introduce additional permute_dims ops. That's because the gradient visit all binding from last one to the first one, and matmul and permute_dims are regarded as two separate ops. E.g.
Forward is:
```
out = matmul(a, transpose(b))
```
Then backward is:
```
grad_a = matmul(grad_out, transpose(transpose(b)))
grad_b = transpose(matmul(transpose(a), grad_out))
```
This PR introduces a new pass, GradientSimplifier, and enhances the Gradient pass to simplify these patterns. The example above will be simplified to
```
grad_a = matmul(grad_out, b)
grad_b = matmul(transpose(grad_out), a)
```
This PR adds these features to the gradient system:
- Checkpointing for gradient pass
- `tvm.relax.testing.nn.checkpoint`
- `tvm.relax.op.grad.start_checkpoint` and `tvm.relax.op.grad.start_checkpoint`
- Support in the Gradient pass
- Fix several minor problems in op_gradient
### Introduction
This PR introduces the high-level reverse-mode automatic differentiation pass `Gradient` for Relax. It's the core component when we are trying training or fine-tuning in Relax IR.
Before upstreaming, this work is actively iterated and maintained in many forks like [mlc](https://github.com/mlc-ai/relax) and [relax-training](https://github.com/ACMClass-TVM-20/relax-training). Now it reaches a relatively stable version and it's time for us to upstream this important work to the unity branch.
The Python side API:
- `Gradient(func_name: str, require_grads: Optional[Union[Var, List[Var]]] = None, target_index: int = 0) -> tvm.ir.transform.Pass`
It will transform the given funcion in the IRModule, and adds a new function that calculates the gradient with regard to the function's output.
### Examples
```
@I.ir_module
class Module:
@R.function
def main(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
# use R.sum to reduce the tensor to a scalar
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
R.output(lv2)
return lv2
After = relax.transform.Gradient("main")(Module)
```
Then the transformed module `After` will be
```
@I.ir_module
class After:
@R.function
def main(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
R.output(lv2)
return lv2
@R.function
def main_adjoint(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tuple(
R.Tensor((), dtype="float32"),
R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")),
):
with R.dataflow():
# original bindings
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
# bindings w.r.t. intermediate variables
lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32")
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(
lv2_adjoint, (3, 3)
)
# bindings w.r.t. parameters
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
R.output(lv2, x_adjoint, y_adjoint)
# return value: (orig_return_values, tuple(adjoints))
return (lv2, (x_adjoint, y_adjoint))
```
Here we specify the target function `main` by its name.
We let the `require_grads` be default value (`None`) so it will calculate all inputs' adjoints (`x_adjoint`, `y_adjoint`) and return them.
We let the `target_index` be default value `0` so it will take the unique return value `lv2` as the target (the scalar we start to differentiate and propagating adjoints) of AD.
Co-authored-by: Yixin Dong <ubospica@gmail.com>
Tianqi Chen