This PR fixes two independent test-isolation issues that only surface
when certain test files run together in one pytest session.
1. Fix `_DialectRedirectFinder` duplicate module execution
`_DialectRedirectFinder.find_spec` used to pre-register the redirect
target module under the legacy alias name before returning the alias
spec.
This interacts badly with CPython import logic: when the requested
module name is already in `sys.modules`, CPython may ignore the returned
alias spec and reuse the target module's original spec instead. As a
result, the target source can be executed again under the canonical
module name, creating a duplicate module object.
This caused patches on aliased modules to silently miss the module
object used by existing code. For example,
`unittest.mock.patch("tvm.tirx.script.builder.buffer_store")` patched
the duplicate module, while the tirx parser still held references to the
original one, so `test_scalar_assign_error_not_swallowed` failed with
`DID NOT RAISE`.
This pr removes the pre-registration and let the import machinery
register the alias normally. Since the alias spec is now used,
`_AliasLoader.exec_module` also restores the canonical `__spec__` and
`__loader__` to avoid stale alias metadata on the loaded module.
2. Remove unused `tirx.intrin_test` op registration
`test_s_tir_transform_lower_match_buffer.py` registered a dummy op:
```python
tvm.ir.register_op_attr("tirx.intrin_test", "")
```
This was a leftover from the old TVMScript parser and is no longer
needed. The modern tirx parser eagerly evaluates `intrin_test(...)`
calls into `T.evaluate(0)`, so this op never appears in parsed IR.
The only remaining effect was adding a category-less `tirx.intrin_test`
entry to the global op registry, which could break
`test_registered_tirx_ops_have_exactly_one_category` depending on test
import order.
This pr removes the unused registration.
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.
## Summary
Restructure TVMScript to be dialect-agnostic at the script-core layer
while letting each extension dialect (TIRX, Relax) own its own
per-dialect script subtree. IR is below script in the dependency
stack and is NOT a peer dialect — its script handlers stay in the
shared core.
This PR folds together two coupled refactors that were initially
opened as separate PRs (#19478 and the original #19479); they
share rename / relocation surface so they ship as one cohesive
change.
## What this PR does
### Per-dialect script subtree (originally #19479)
- Moves per-dialect printer + builder from
`src/script/{printer,ir_builder}/{tirx,relax}/` to
`src/{tirx,relax}/script/{printer,builder}/`.
- Tightens `src/script/*.cc` CMake glob to the dialect-free core.
- Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry
(`script.ir_builder.decl_function.<type-key>`); removes
cross-dialect includes from the shared core.
- Adds `tvm.script.register_dialect` API + `__getattr__` + a
`sys.meta_path` finder for Python-side dialect discovery.
In-tree dialects (tirx, relax) registered centrally in
`python/tvm/__init__.py`.
- Drops the obsolete static re-export shims at
`python/tvm/script/{parser,ir_builder}/{tirx,relax}/`.
### Dialect-agnostic printer config (originally #19478)
- Relocates `include/tvm/ir/script_printer.h` →
`include/tvm/script/printer/config.h` next to the rest of the
printer's public surface. The header is not IR-specific.
- Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for
clarity (the macro registers Script as the kRepr callback +
per-type vtable dispatch). Aligns with the `TVM_REGISTER_*`
family.
- Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`,
`relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor
of a generic `ffi::Map<String, Any> extra_config` keyed by
`"<dialect>.<knob>"`. Each call site reads via the templated
accessor `config->GetExtraConfig<T>("...", default)`.
- Promotes `std::string` config fields to `ffi::String`.
After this lands, the script-printer core knows nothing specific
about any dialect — new dialects plug in via the registry pattern
with zero core edits. Public Python API surface unchanged.
## Why
This refactor reshapes each backend into a self-contained
`src/target/<X>/`
cluster (with optional `src/target/<X>/llvm/` for LLVM-dependent
codegen) and
introduces a per-backend fallback module that absorbs the cross-compile
role
cleanly — without `target/opt/` stubs, without `DeviceSourceModuleNode`,
and
without leaking a synthetic `kind()` to consumers.
## High-level principles
- **One directory per backend.** All codegen-side files for backend
`<X>` live
under `src/target/<X>/`. Optional `src/target/<X>/llvm/` subdir for
files
that require `USE_LLVM` at build time. Backend grouping wins over
build-dependency grouping (the latter being upstream's `target/llvm/`).
- **Plugin-only runtime modules.** `src/runtime/<X>/<X>_module.h` is
deleted.
The runtime's real `<X>ModuleNode` is reachable only via the FFI
registry
(`ffi.Module.create.<kind>`, `ffi.Module.load_from_bytes.<kind>`). No
C++ API surface other than the static registrations.
- **Per-backend fallback module for cross-compile.** Each `<X>` gets a
`<X>FallbackModuleNode` in `src/target/<X>/<X>_fallback_module.{h,cc}`.
Same `kind()` as the real module. Codegen-time only — never reachable
via
load. `GetFunction` errors with a backend-specific "runtime not linked"
message; `InspectSource` works.
- **Codegen-side wrapper does the fallback selection.** Codegen calls
`<X>ModuleCreateWithFallback(...)`, which tries
`ffi.Module.create.<kind>`
via the registry; on miss, falls through to `<X>FallbackModuleCreate`
(plain C++; reachable directly from the fallback header). When
`USE_<X>=ON`
is in effect, the registry hit returns the real module; when
`USE_<X>=OFF`, the fallback is what codegen gets. No CMake `if/else`
gating; fallback always compiled.
## Specific changes
### New per-backend directories (codegen + fallback)
- `src/target/cuda/` — `codegen_cuda.cc` + `intrin_rule_cuda.cc` +
fallback module pair + `llvm/codegen_nvptx.cc`
- `src/target/rocm/` — fallback module pair + `llvm/codegen_amdgpu.cc` +
`llvm/intrin_rule_rocm.cc`
- `src/target/hexagon/` — fallback module pair +
`llvm/codegen_hexagon.cc` + `llvm/intrin_rule_hexagon.cc`
- `src/target/metal/` — `codegen_metal.cc` + `intrin_rule_metal.cc` +
fallback module pair
- `src/target/vulkan/` — `build_vulkan.cc` + the rest of `target/spirv/`
absorbed + fallback module pair
- `src/target/opencl/` — `codegen_opencl.cc` + `intrin_rule_opencl.cc` +
fallback module pair
- `src/target/webgpu/` — `codegen_webgpu.cc` + fallback module pair (was
`WebGPUSourceModuleNode`, renamed)
### New fallback classes
`<X>FallbackModuleNode` for X in {`CUDA`, `ROCm`, `Hexagon`, `Metal`,
`Vulkan`, `OpenCL`, `WebGPU`}. Each:
- `kind()` matches the real backend
- Stores `(code or smap, fmt, fmap, source)` — no driver/runtime calls
- `GetFunction` errors with backend-specific "runtime not linked"
message
- `InspectSource` works fully
- `SaveToBytes` byte-identical to real
## Problem
On Python 3.14, any use of TVMScript raises a `TypeError` before the
module body is even parsed:
```
TypeError: unsupported format string passed to type.__format__
```
The traceback points to
`python/tvm/script/parser/core/diagnostics.py:120`:
```python
raise TypeError(f"Source for {obj:!r} not found")
```
## Root Cause
`{obj:!r}` is an invalid f-string expression. The `:` introduces a
`format_spec`, so `!r` is passed to `type.__format__` as a format string
— which it does not support.
The intended syntax for a `repr()` conversion is `{obj!r}` (no colon).
Python 3.14 re-implemented f-string parsing under [PEP
701](https://peps.python.org/pep-0701/) and now strictly validates
format specs, surfacing this latent bug. Python 3.10–3.13 silently
passed the invalid spec to `__format__` and happened not to raise in
most code paths, so the bug went unnoticed.
## Fix
```diff
- raise TypeError(f"Source for {obj:!r} not found")
+ raise TypeError(f"Source for {obj!r} not found")
```
One character change. Valid across all Python versions >= 3.6.
## Testing
Verified on Python 3.14.2 (darwin/arm64):
- TVMScript `ir_module` + `prim_func` parses and compiles correctly
after the fix
- Full TVMScript test suite: **628 passed, 1 xfailed** (the 1 failure in
`test_tvmscript_roundtrip.py::test_roundtrip[relax_symbolic_size_var]`
is pre-existing and unrelated to this change)
Add handling for `doc.keyword` nodes in `ExprEvaluator._visit` to ensure
expressions (e.g. `BoolOp`) in keyword arguments are processed with
correct evaluation methods.
Fix#18972 . For more details, please refer to this issue.
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 do a rebuild of TIR Common Subexpression Elimination (CSE) using
a two-phase architecture:
- **Phase 1 — CSEPlanner**: Read-only visitor that builds a scope tree
and expression DAG. Computes a plan (InsertBeforeTable + ExprRemapTable)
in a single pass using shallower-first processing with repr propagation
— no cascade loop needed.
- **Phase 2 — CSERewriter**: Mechanical mutator that inserts
`Bind(cse_var, expr)` statements and substitutes expressions per the
plan.
Key improvements over the old implementation:
- **Simpler architecture**: Two clean classes (planner + rewriter)
instead of interleaved analysis/mutation
- **No cascade loop**: Shallower-first processing with repr propagation
resolves all CSE opportunities in one plan + one rewrite
- **Incremental DAG construction**: Expression depth, children, and
consumed counts computed during bottom-up scan — no separate traversals
- **No single-use bindings**: Consumed count tracking avoids introducing
bindings that would only be used once
- **Unified insertion via VisitStmt**: SeqStmt flattening handles all
insertion contexts uniformly
Other changes:
- Rename `CommonSubexprElimTIR` → `CommonSubexprElim`, remove
`enable_cse_tir` and `identify_equiv_terms` params
- Move old CSE tools (used by cache_index) to
`cache_index_helpers.{cc,h}`
- Remove unused `arith.detect_common_subexpr` API
- Add `T.bind` as lowercase alias for `T.Bind`
## Summary
- Remove `body` field from `AllocBufferNode` and `DeclBufferNode`,
making them flat statements consistent with `Bind`
- Buffer scope extends to end of enclosing scope via flat `SeqStmt`
semantics
- 60 files changed across core IR, codegen backends, transforms, script
IR builder, and tests
## Test plan
- All existing test suites pass (tir-transform, tir-base, tvmscript,
s_tir, codegen, C++)
## 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
## 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.)
With `from __future__ import annotations`, Python stores annotations as
strings
and does not capture annotation-only variables in `__closure__`. This
broke
TVMScript when buffer shapes/dtypes referenced closure variables.
Fix: wrap `extra_vars` in a `collections.ChainMap` with snapshots of all
live
caller-frame locals (from `inspect.stack()`) as fallback layers in both
`tir/entry.py` (`prim_func`) and `ir/entry.py` (`ir_module`). The
`ir_module`
function now also captures `outer_stack = inspect.stack()` at its entry
point,
mirroring the existing pattern in `prim_func`. Lookup falls back to
frame locals
only on cache miss, preserving existing behavior for non-PEP-563 code.
Add `tests/python/tvmscript/test_tvmscript_pep563_closure.py` (requires
`from __future__ import annotations` at the top) covering closure
variables in
buffer shapes, dtypes, nested scopes, ir_module, and mixed
annotation+body use.
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 phases out alloc const node in the TIR.
This node was oroginally introduced to include embedded weights into the
allocation. However, the presence of the particular IR couples the data
allocation and the weight placement, which is not as desirable especialy
when weights get large. A better approach is to have extra annotation on
the allocation and store weights separately either outside module or as
part of module/function attribute.
As a result, we phases out this node which can help us to simplify code
logic in the codebase.
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
## Why
ONNX models use the Size operator to get total element count of a
tensor. Relax didn't have a native equivalent.
## How
- Adds R.size(tensor) operator that returns the total number of elements
in a tensor as a scalar int64
Introduces the below features over texture annotation
- Lowering, codegen and runtime for texture.
- image2d_array_t support - Added depth dimension allows more
allocations using texture instead of falling back to buffer when the
texture limits exceeds.
- A comprehensive set of schedules for Adreno textures.
- Texture packing of arbitrary types up to 128 bit (FP16-NCHW8c,
INT8-NCHW16c ...etc.).
- A clBufferDescriptor debug dump controlled by cmake options.
- Pipeline definition for adreno target.
While covering these features the below interfaces or passes or enhanced
which need a review.
- alloc_tensor: VDevice information is passed across these API's. The
way of texture allocation is ```alloc_storage``` allocates buffer/image
objects as requested followed by alloc_tensor being a view of any scope.
This takes care of optimum utilization backing memory across different
image objects or scopes.
- Constants Saving: Handled by adding memory scope section in
executable. This introduces a new header magic to retain the backward
compatibility.
- Static Memory Planing: Mostly port from Relay static memory planner
with mixed mode allocator.
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Sanjay <sanjs@qti.qualcomm.com>
This PR supports NVRTC as an alternative to NVCC for faster, device-side
JIT compilation of CUDA kernels, in favor of the PR
[https://github.com/apache/tvm-ffi/pull/283](https://github.com/apache/tvm-ffi/pull/283).
It enhances the CUDA compilation backend by:
- Adding Python NVRTC support using cuda-python bindings
- Removing legacy C++ NVRTC fallback in favor of a Python-first approach
- Keeping nvcc as the default compiler with fatbin output (no behavior
change for existing users)
Users can choose the compilation backend using an environment variable
`TVM_CUDA_COMPILE_MODE`, choosing from "nvcc" and "nvrtc". For example,
`TVM_CUDA_COMPILE_MODE=nvrtc python3 your_program.py`
Here is a short benchmark of the compilation speed of kernels in
`test_target_codegen_cuda.py`.
### NVCC vs NVRTC Compilation Time Comparison (Python-side Call)
| Test Case | Code Size | NVCC Time (ms) | NVRTC Time (ms) | Speedup |
| :--- | :--- | :--- | :--- | :--- |
| `test_crossthread_reduction1` | 1945 B | 241.27 | 51.23 | **4.7x** |
| `test_cuda_bf16_vectorize_add` | 3760 B | 342.72 | 44.50 | **7.7x** |
| `test_cuda_const_float_to_half` | 12394 B | 272.85 | 31.99 | **8.5x**
|
| `test_cuda_device_func_call` | 975 B | 215.58 | 21.47 | **10.0x** |
| `test_cuda_float_const_hex_format` | 685 B | 217.39 | 20.52 |
**10.6x** |
| `test_cuda_floordiv_with_vectorization` | 1050 B | 213.88 | 23.32 |
**9.2x** |
| `test_cuda_inf_nan` | 673 B | 214.33 | 24.94 | **8.6x** |
| `test_cuda_tensormap` | 755 B | 213.91 | 20.74 | **10.3x** |
| `test_cuda_thread_sync_inside_condition` | 1007 B | 213.43 | 28.29 |
**7.5x** |
| `test_cuda_vectorize_add` | 908 B | 226.81 | 40.39 | **5.6x** |
| `test_cuda_vectorize_load` | 734 B | 217.25 | 24.02 | **9.0x** |
| `test_device_host_call_same_func` | 924 B | 216.03 | 21.21 | **10.2x**
|
| `test_vectorized_intrin1` | 847 B | 226.15 | 26.34 | **8.6x** |
### NVSHMEM Support
Currently, NVSHMEM is **not** supported via NVRTC.
- Fallback Behavior: When NVSHMEM is required, the compilation pipeline
will automatically fall back to NVCC, even if `TVM_CUDA_COMPILE_MODE` is
set to nvrtc.
- Future Roadmap: Support for NVRTC with NVSHMEM is planned for
follow-up PRs.
## Related Issue
closes https://github.com/apache/tvm/issues/18344
## Why
When a `T.macro` containing a block was called multiple times in a TIR
function, all expanded blocks had the same name, causing a "Duplicated
block name" error in meta_schedule.
## How
Implemented automatic block name suffixing during macro expansion
An initial change to add `ForNode::step`.
- Add `Optional<PrimExpr>` typed step attribute to ForNode. Then add
minimal codes for
- Roundtrip support for TIR tvmscript grammar
- Correctness of TIR lowering pipeline:
- Canonicalize the loop in default pipeline
- Ensure the original `ForNode::step` is not dropped by mutations on
`ForNode`.
- CodeGen support for non-zero min and non-trivial step.
- TODOs in the future (hopefully)
- For **all transformations and analysis tools**, make adaptions to
non-consecutive loop iteration indices
- Correctness of TensorIR schedule and MetaSchedule
---------
Co-authored-by: baoxinqi <bao.xinqi@intellif.com>
Implement proper parsing and evaluation of chained comparison operators
(e.g., `0 < i < 128`) in TVMScript. The sequence comparisons are now
correctly expanded to their logical equivalents (e.g., `(0 < i and i < 128)`).
Changes:
- Updated expression evaluator to handle sequence comparisons correctly
- Added test case to verify sequence comparison functionality
Add support for conditional expressions in TVMScript
This PR adds support for conditional expressions in TVMScript parser,
which allows developers to use Python-style conditional expressions
```python
@T.prim_func
def func(A: T.buffer((128, 128), "float32")):
for i, j in T.grid(128, 128):
A[i, j] = i if i < j else j
@T.prim_func
def expected(A: T.buffer((128, 128), "float32")):
for i, j in T.grid(128, 128):
A[i, j] = T.if_then_else(i < j, i, j)
```
This PR allows calling Python functions directly from Relax IR,
where integration between Relax computations and Python/PyTorch
operations can be supported.
### Usage Example
```python
@I.ir_module
class MyModule(BasePyModule):
@I.pyfunc
def pytorch_add(self, x, y):
return x + y
@R.function
def compute(x: R.Tensor((5,), "float32"), y: R.Tensor((5,), "float32")) -> R.Tensor((5,), "float32"):
result = R.call_py_func("pytorch_add", (x, y), out_sinfo=R.Tensor((5,), "float32"))
return result
```
* 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>
This PR cleans up the python API to make things more consistent
with existing python array api and torch.
Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device
API updates:
- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func
This PR Updates the NDArray => Tensor.
Both tensor and ndarray are commonly used terms.
Because the term Tensor is getting more common in the context of ML,
we do the rename to stay more aligned with torch.Tensor and DLTensor.
### **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
[REFACTOR] Phase out getattr based attribute handling
This PR phases out getattar based attribute handling as they are slower
and introduces extra code path.
This does mean that if an Object is not explicitly registered
in python side, we will no longer be able to access the field by name.
Likely this is also desirable as we would like to enable faster use that
updates the python end and do not rely on these behavior.
Shushi Hong