This pr switches CI to the Ubuntu 24.04 (noble) images. Bump ci_tag in
ci/jenkins/docker-images.ini to 20260629-192919-24bbfd2e -- the images
built from #19893 (ci_cpu/ci_arm/ci_wasm/ci_gpu on Ubuntu 24.04, whose
default g++ is gcc-13, giving full C++20 support).
Also, this pr re-enables Z3 (AUTO). #19828 temporarily set USE_Z3=OFF
(in CMakeLists.txt and the pyproject wheel build) to dodge a z3-static
build failure. The CI image now ships z3-static (#19835), so this
restores USE_Z3=AUTO: the Z3-backed Analyzer proving is enabled when
z3-static is available and silently skipped otherwise.
While Z3 stayed disabled, PrimExprNode::ty became a method, leaving two
stale field accesses in z3_prover.cc's IsZ3SupportedExpr (only compiled
under TVM_USE_Z3). Fixed expr->ty -> expr->ty().
Verification:
- The Ubuntu 24.04 images (#19893) built successfully for ci_cpu/ci_arm/
ci_wasm/ci_gpu (the GPU image includes ROCm 6.4.4 and the CUDA 24.04
base).
- Re-enabling Z3 was validated with a build-only wheel run: all four
wheels (Linux x86_64/aarch64 manylinux_2_28, macOS arm64, Windows) build
green with z3-static compiled and linked, confirming the earlier
z3-static link concern is resolved on the current toolchain.
## Rationale
Analyzer constraint scopes continue to provide loop-positive facts to
constant-bound and rewrite proofs. During IntSet relaxation, however, a
scoped domain constraint is a refinement only for a variable explicitly
present in the relaxation map. Applying it to an unmapped variable
reinterprets a free parameter as a relaxation domain and can let a
loop-local symbol survive recursive interval evaluation.
## Changes
- Apply scoped IntSet constraints only to variables already present in
the relaxation map; unmapped variables remain free parameters.
- Remove the finite-bound restoration fallback and its stronger
parametric-bound contract.
- Add a direct compact-buffer regression that prevents a loop-local
variable from escaping into a function-scope allocation extent.
This PR updates `IRMutatorWithAnalyzer` and `IRVisitorWithAnalyzer` to
add the constraint `loop_extent > 0` while visiting a `ForNode` body.
If execution reaches the loop body, the loop must have at least one
iteration, so the body context can safely assume the extent is positive.
The constraint is scoped only to the loop body, leaving the loop header
expressions (`min`, `extent`, `step`) outside of this assumption.
While validating this change, it exposed an issue in `IntSetAnalyzer`:
one-sided constraints such as `m > 0` could cause existing parametric
bounds like `m - 1` to be recursively relaxed to `+inf`. This caused
`DomainTouched` to lose finite symbolic bounds. The PR fixes this by
preserving existing parametric bounds when recursive interval relaxation
would otherwise replace them with infinity.
## Rationale
`SizeVar` encodes nonnegativity in runtime subtype identity, which is
fragile under cloning and remapping. Symbolic integer values should use
one `Var` representation, with nonnegative facts recorded in the
analyzer at the use sites that establish them.
## Changes
- Remove `SizeVar` from the C++, Python, TE, TVMScript, FFI, visitor,
and serialization surfaces, and migrate callers to `Var`.
- Preserve the existing Relax constraint ownership model and use
`MarkGlobalNonNegValue` as the canonical path for global nonnegative
facts.
- Preserve `T.handle()` as the normal opaque-handle form. An optional
dtype constructs a typed pointer, with `T.handle("void")` reserved for
an explicit pointer-to-void.
## Summary
- Make `PrimExpr` a typed C++ view over `Expr` values whose
`ExprNode::ty` is `PrimType`, instead of using a separate runtime node
class as the proof of primitive-ness.
- Use the shared `ir::Call` node for Relax, TIRX, and primitive-valued
calls, while keeping primitive-only APIs explicit at their semantic
boundaries.
- Keep Python on the general `Expr` surface for primitive-typed values
so `isinstance` behavior does not imply a nominal primitive-expression
subclass.
## Design Rationale
The main advantage of this change is that common expression nodes such
as `Call` can be unified without specializing each one to `PrimType`. A
single `ir::Call` can represent a Relax tensor call, a Relax scalar
call, or a primitive-valued intrinsic call; the result type stored in
`ExprNode::ty` determines whether that particular value can be viewed as
`PrimExpr`.
This keeps the IR node hierarchy focused on expression structure rather
than result-type categories. Nodes that are intrinsically primitive,
such as integer and floating-point literals or TIRX primitive operators,
still have strongly typed C++ APIs and data structures. General nodes
whose result type may vary, such as `Call`, remain general `Expr` nodes
and are narrowed to `PrimExpr` only where primitive-only semantics are
required.
The PR also keeps the compatibility surface practical: C++
primitive-only APIs continue to accept `PrimExpr`, Python exposes a
compatibility predicate for checking the primitive typed category, and
visitors/printers use one natural `Call` path rather than duplicating
Relax and primitive call handling. Missing expression types are
represented explicitly with `Type::Missing()` so constructors can leave
type inference to later analysis without relying on nullable `Type`
values.
In the past we have been using `DataType` in PrimExpr.dtype field to
check type information for PrimExpr while still having BaseExpr.ty for
richer type information. DataType is also used both in runtime and
compiler. This PR streamlines the boundary:
- PrimExpr.ty now carries PrimType that replaces original use of
`DataType`
- Runtime use will now favor DLPack DLDataType, removing one layer of
indirection.
- Constants attributes where values are usually runtime values, will use
`DLDataType`
- DataType will be phased out after this PR
We also brings up helper functions in PrimType, but also limits them to
a more concise set so the functions do not grow with the data type codes
in DLPack.
This is a major refactor that changes the IR primitive. It helps to
bring possible future benefits:
- Unified type mechanism through Expr.ty
- Possibility of carry future Type nodes
Migration Guide:
- Use `PrimType` when code reasons about compiler expression types,
tensor element compiler types, or constructs a `PrimExpr`/compiler type.
- Use existing source types such as `expr.ty()`, `ExprOp.expr_ty()`, or
TE tensor element `dtype` where possible instead of rebuilding a type
from dtype text.
- Use raw `DLDataType` for runtime constants, ABI paths, dtype-valued
attrs, and storage/runtime helper logic.
- Prefer direct `PrimType` equality, `MatchesCode(...)`,
`MatchesElementType(...)`, and `WithCode(...)` over local wrappers or
string dtype checks.
Performance:
Using Object type instead of DLDataType would indeed bring some
performance impact to the IR. We have done the following performance
optimizations:
- Make sure most of the outputs reuse one of the PrimType from inputs
- Cache a thread local PrimType based on input so we don't repeatly
realloc
We did benchmarks show that rewrite simplify operation stays within
+-10% overhead of original one. Which merits the refactor given the
benefit the unfication brings
Copying an Analyzer handle shares the same mutable AnalyzerObj, so a
pass had no way to snapshot accumulated facts (variable bounds, modular
sets, rewrite/canonical bindings, integer-set domains, literal
constraints, transitive comparisons) and keep exploring without mutating
the original.
This pr adds AnalyzerObj::Clone(), which allocates a fresh AnalyzerObj
and copies each sub-analyzer's persistent state through a new
per-sub-analyzer CopyFrom. Parent back-pointers are re-established by
the fresh constructor rather than copied, and per-query/recursion
scratch state is left default. Exposed to Python as Analyzer.clone().
This PR fixes#19825, which restricts the rewrites
```
floormod(x * c1 + y, c2) -> floormod(x * floormod(c1, c2) + y, c2)
```
and
```
floormod(x + y * c1, c2) -> floormod(x + y * floormod(c1, c2), c2).
```
While algebraically valid in isolation, these transformations rewrite
only the `floormod` side of a matching `floordiv`/`floormod` pair. As a
result, the two expressions no longer share a visible fused index
expression, causing `DetectIterMap` to reject otherwise bijective splits
such as:
```
lane = flat % 128
reg = flat // 128
```
where both expressions originate from the same fused index.
### Context
Per the suggestion in #19825, the two rewrites are guarded with `c1 % c2
== 0` rather than dropped outright. The multiplied term is still
eliminated when it is a multiple of the divisor (e.g. `(x*10 + y) % 2 ->
y % 2`), which is safe for `DetectIterMap`; only the
coefficient-shrinking case (`c1` not a multiple of `c2`) is disabled.
Both operand orderings are covered, and the PR adds a rewrite-simplify
regression plus an end-to-end `DetectIterMap` regression test.
## Summary
This PR adds a Z3 SMT solver backend to `tvm::arith::Analyzer` for
stronger integer arithmetic proving.
The integration is guarded by `USE_Z3`, which defaults to `AUTO`. In the
default mode, TVM enables Z3 when the static Z3 development artifacts
are available and otherwise builds the conservative stub implementation.
When Z3 is enabled, `Analyzer::CanProve` runs the existing TVM
arithmetic analysis path first, then falls back to Z3 only when the
existing analyzers cannot prove the predicate and the requested strength
is `kSymbolicBound`. Z3 is linked statically from the PyPI `z3-static`
package, so `libtvm` does not need a runtime `libz3` dependency.
## Features
- Z3 build support through `USE_Z3`, defaulting to `AUTO`.
- A new `arith::Z3Prover` sub-analyzer owned by `arith::Analyzer`.
- SMT-LIB2 export for debugging and external solver reproduction.
- Python debug/config APIs: `Analyzer.get_smtlib2`,
`Analyzer.set_z3_timeout_ms`, `Analyzer.set_z3_rlimit`, and
`Analyzer.get_z3_stats`.
- C++ APIs for proving, binding, constraints, stats, model inspection,
and satisfying-value counting.
- Scalar integer, unsigned integer, and boolean expression translation
to Z3.
- Support for arithmetic, comparisons, boolean operators, `min`, `max`,
`select`, `if_then_else`, `let`, casts, truncated division/modulo, floor
division/modulo, and selected bitwise/shift operations.
- Deterministic solver control using Z3 `rlimit`, with `random_seed`
fixed to `42`.
- Thread-local Z3 context sharing to reduce initialization overhead
while keeping thread safety.
- A disabled-mode stub implementation that returns conservative results
when Z3 is not built.
## Implementation Notes
- The real and stub implementations live in `src/arith/z3_prover.cc`,
selected by the `TVM_USE_Z3` macro from
`cmake/modules/contrib/Z3.cmake`.
- `cmake/modules/contrib/Z3.cmake` first resolves the PIC static `libz3`
layout provided by `z3-static` using its `z3_static.get_cmake_dir()`
helper, then falls back to a custom `Z3_DIR` or `CMAKE_PREFIX_PATH`
installation.
- `USE_Z3=ON` requires Z3 to be found, while `USE_Z3=AUTO` allows source
builds and CI jobs without Z3 artifacts to continue with the stub.
- The Z3 fallback is exception-safe and gated behind `kSymbolicBound`,
so the common `kDefault` path does not pay solver cost.
- TVM `Div` and `Mod` are translated with truncating helpers rather than
Z3's Euclidean operators to stay sound for negative dividends.
- Shift handling relies on Z3's native bit-vector semantics and does not
add hard assertions to the shared solver.
## References
The implementation is based on the Z3 analyzer integration used in
TileLang's TVM fork, with the upstream port kept scoped to TVM's
arithmetic analyzer.
- [tile-ai/tilelang#1367](https://github.com/tile-ai/tilelang/pull/1367)
- [tile-ai/tilelang#1458](https://github.com/tile-ai/tilelang/pull/1458)
- [tile-ai/tilelang#2216](https://github.com/tile-ai/tilelang/pull/2216)
- [TileLang/tvm#22](https://github.com/TileLang/tvm/pull/22)
- [TileLang/tvm#24](https://github.com/TileLang/tvm/pull/24)
- [Original TileLang TVM
commit](https://github.com/tile-ai/tvm/commit/e633295de994a89668d7a9930dbbd455af3efc66)
---------
Signed-off-by: Ubospica <ubospica@gmail.com>
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.
This PR makes `arith::Analyzer` a first-class tvm-ffi object.
The implementation splits the previous concrete `Analyzer` class into:
- `AnalyzerObj`, the mutable object node that owns analyzer state,
sub-analyzers, caches, and bindings
- `Analyzer`, a reference-counted `ObjectRef` handle that can be passed
across the tvm-ffi boundary
This allows Python and C++ to share the same analyzer instance, so
bindings, constraints, and cached facts can persist across FFI calls.
Public APIs that accept an analyzer now use `const arith::Analyzer&`,
while internal helper APIs that only borrow the object continue to use
`AnalyzerObj*`.
---------
Co-authored-by: Ubospica <ubospica@gmail.com>
## Summary
`CanonicalSimplifier::Impl::VisitExpr_(LTNode)` Case 2 rewrites
S + xn < 0 ⇔ S/d + xn // d < 0 where d = gcd(scales)
The Case 1 derivation only works when `xn ≥ 0`. With `scale = -1` the
equivalence becomes `≤` rather than `<`, and the rewrite silently
strengthens the predicate by dropping the boundary `S/d == xn // d`.
After CSE/inlining, a comparison such as `2*(tx%4) < 16*warp +
(tx%32)//4` (where `row` and `col` are independent projections of the
same lane id) reaches canonical_simplify with the divided projection on
the LHS (scale = -1), and Case 2 folds it to a plain `0 < warp_id` —
zeroing every thread that should have written `val` in warp 0. The same
path also folds other configurations (e.g. `0 < (tx%32) - 8*warp`) all
the way to `False`.
The fix gates Case 2 with `extra->args[0]->scale == 1`. The original
target shape (`yn % m` with positive scale and `lower_factor=1`, plus
the `scale = +1 / lower_factor > 1` generalization) is unchanged;
truly-always-true comparisons still fold to `True`.
## Test plan
- New regression test `test_simplify_le_negative_scale_extra` in
`tests/python/arith/test_arith_canonical_simplify.py` — asserts on
simplified `PrimExpr`, no GPU required; pre-fix fails, post-fix passes.
It also pins the buggy `scale = -1` shapes to their unsimplified form,
confirms the `scale = +1` Case 2 path still optimizes, and re-asserts
the truly-always-true variant still folds to `True`.
- Existing `test_simplify_le` (the original Case 2 target with `scale =
+1`) still passes.
- `tests/python/arith/test_arith_canonical_simplify.py` — 16 passed.
- Full `tests/python/arith/` — 932 passed (1 pre-existing flaky
random-seed failure in `test_arith_solve_linear_equations.py` unrelated
to this change, passes on rerun).
## Problem
`Analyzer::Bind` could hang indefinitely (>300s, ~200% CPU, no GPU work)
while binding a small expression for one variable. The root cause is
general and lives in `src/arith/int_set.cc`.
Diagnosis: 100% of the time is spent in `arith::Analyzer::Bind` →
`IntSetAnalyzer` → `IntervalSetEvaluator`, evaluating a **5-node** bound
expression. A counter showed **>2^20 `VisitExpr` calls at recursion
depth 67** with no end in sight.
## Root cause
`IntervalSetEvaluator::VisitExpr_(VarNode)` relaxes a variable's bounds
by recursively evaluating **both** the `min` and `max` sub-expressions
of its mapped interval. For diamond-shaped variable dependency chains
(`a → {b, c}`, `b → {d, e}`, …) the shared sub-expressions are
re-expanded along every path, so cost is **O(2^depth)** in the length of
the dependency chain — bounded only by `dom_map_.size()` (~67
interdependent vars in the failing case).
## Fix
Memoize the fully-relaxed interval **per variable** (`relax_memo_`) and
break cyclic dependencies with an in-progress set
(`relax_in_progress_`). A variable's relaxed interval is deterministic
for a given evaluator instance (`dom_map_`/`dom_constraints_` are
fixed), so memoizing collapses the diamonds to linear cost. Short chains
— the common case, which never reached the old `recur_depth_ >=
dom_map_.size()` cutoff — are unaffected, so the change is
behavior-preserving outside the pathological case.
## Tests
New regression tests in `tests/python/arith/test_arith_intset.py`:
- `test_relax_deep_variable_dependency_chain` — a 64-deep diamond
(`O(2^64)` without the fix; verified to hang on a clean build), also
asserting the relaxed result is correct (`x0 → [-n, 100+n]`).
- `test_relax_cyclic_variable_dependency` — a cyclic `x↔y` dependency
must terminate.
## Verification
- `tests/python/arith/test_arith_intset.py` — 20 passed (the deep-chain
test completes instantly).
- Full `tests/python/arith/` — 933 passed (1 pre-existing flaky
random-seed failure in `test_arith_solve_linear_equations.py` unrelated
to this change, passes on rerun).
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
## 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
Phase out `src/arith/scalable_expression.{h,cc}`. The arith layer no
longer attempts to prove anything about scalable vectors — proofs that
depended on `Target::Current()` are removed. Scalable vectors remain a
first-class concept; arith just doesn't reason about their lengths.
## Use-site summary
Only 16 call sites total across 7 symbols (9 live, 7 proof-related).
| Symbol | Live callers (kept) | Proof callers (deleted) | New home |
|---|---|---|---|
| `ExtractVscaleFactor` | 4 × `arith/rewrite_simplify.cc` + 2 ×
`tirx/ir/expr.cc` | — | file-local in each |
| `IsVScaleCall` | 1 × `tirx/op/op.cc` + 1 ×
`tirx/transform/vectorize_loop.cc` | — | inline at use sites |
| `ContainsVscaleCall` | 4 × `arith/rewrite_simplify.cc` + 1 ×
`s_tir/schedule/ir_comparator.cc` | — | inline at use sites |
| `TargetHasVLA` | 2 × `tirx/transform/vectorize_loop.cc` | analyzer.cc
+ const_int_bound.cc | local in vectorize_loop.cc |
| `GetVScaleValues` | 1 × `target/llvm/codegen_aarch64.cc` | analyzer.cc
+ const_int_bound.cc | inlined at codegen_aarch64 |
| `CanProveVscaleExpressionFromKnownValues` | — | analyzer.cc | DELETE |
| `SubstituteVScaleWithKnownValue` | — | internal only | DELETE |
## Changes (6 commits)
1. Move `ExtractVscaleFactor` to file-local anonymous-namespace helpers
in `rewrite_simplify.cc` and `tirx/ir/expr.cc`. Function is small;
per-file duplication is cleaner than a shared header.
2. Inline `IsVScaleCall` / `ContainsVscaleCall` / `TargetHasVLA` at call
sites (1-3 line predicates, anonymous-namespace per consumer `.cc`).
3. Drop the scalable-vector proof scaffolding from `arith/analyzer.cc`
(substitution-proof loop) and `arith/const_int_bound.cc` (vscale
branch). `vscale()` calls fall back to `Everything()` — no special bound
narrowing.
4. Delete `scalable_expression.{h,cc}`. Inline the `GetVScaleValues`
body at `codegen_aarch64.cc` (computes `max_val = vector_width / 8`
floor-rounded to a power of two for the LLVM `vscale_range` attribute).
5. Mark `pytest.mark.xfail` on 19 tests that relied on the deleted
substitution-proof loop.
6. `pre-commit` line-length cleanup.
## Compatibility / intentional regression
This is a hard break for any consumer of the deleted symbols. They were
already in a private header (`src/arith/scalable_expression.h`, not
under `include/`).
19 tests that proved vscale-bearing inequalities on SVE / RVV are
xfailed. The proofs were target-dependent and the new policy is that
arith does not attempt them.
## Summary
This PR cleans up technical debt in the TIR simplification machinery via
two commits:
**Commit 1: Phase out ControlFlowGraph and NarrowPredicateExpression**
- Remove `ControlFlowGraph` (~2360 lines) from `src/tirx/analysis/` —
used only in
non-default config paths that are no longer maintained
- Remove `NarrowPredicateExpression` from `src/arith/` — sole non-test
caller was `ControlFlowGraph`
- Remove gated config fields `propagate_knowns_to_prove_conditional` and
`propagate_knowns_to_simplify_expressions` from `SimplifyConfig`
- Remove `use_dataflow_analysis` from `RemoveNoOpConfig`
- Delete the associated test files and test cases that tested the
now-removed paths
- ~3800 lines deleted
**Commit 2: Rename Simplify → StmtSimplify**
- Rename `src/tirx/transform/simplify.{h,cc}` → `stmt_simplify.{h,cc}`
- Rename C++ identifiers: `Simplify` → `StmtSimplify`, `SimplifyConfig`
→ `StmtSimplifyConfig`
- Rename FFI keys: `"tirx.Simplify"` → `"tirx.StmtSimplify"`,
`"tirx.transform.Simplify"` → `"tirx.transform.StmtSimplify"`
- Update Python wrappers and all call sites (~40 files)
- Clarifies that this pass operates on statements (distinct from
expression-level `arith::Analyzer::Simplify()`)
## Test plan
- [x] `tests/python/tirx-transform/test_tir_transform_simplify.py` — 52
tests pass
- [x] `tests/python/tirx-transform/test_tir_transform_remove_no_op.py` —
18 pass, 5 xfail
- [x] `tests/python/arith/` — full arith test suite passes
- [x] `tests/python/tirx-transform/` — full suite: 315 passed, 8
xfailed, 1 xpassed (pre-existing vectorize failure unrelated to this
change)
- [x] `pre-commit run --all-files` — all hooks pass
## 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.
The C++ Analyzer::Bind() already supports allow_override, but the FFI
bridge always used the default (false). This change threads the optional
argument through the FFI layer and the Python wrapper so callers can
rebind variables without triggering an error.
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`
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 removes unused `from tvm import te` imports from 25 test files
across the codebase, continuing the ongoing TE → TVMScript migration
cleanup.
Changes:
- Remove unused `from tvm import te` from 24 test files in
tir-transform/, s_tir/transform/, codegen/, arith/,
all-platform-minimal-test/, and testing/
- Replace `te.var("x")` with `tvm.tir.Var("x", "int32")` in
`test_s_tir_transform_decorate_device_scope.py` (the only file where te
was actually used)
- Clean up stale `tvm.tir.ir_builder` comment references in
`test_tir_transform_convert_ssa.py`
This PR phases out legacy target string format in favor of the json
style format that is more well formed. It also simplfies our overall
code in handling multiple formats.
* Refactor division simplification in RewriteSimplifier and add corresponding test
This commit removes the specific case for rewriting division by a constant float in the RewriteSimplifier. Additionally, a new test is introduced to verify the behavior of float division simplification, ensuring that the division is correctly handled without the previous rewrite logic.
* test fix
* test fix
* cifix
* fix
* Enhance ConstIntBoundAnalyzer and IntervalSet with modular set analysis
- Added modular set analysis to ConstIntBoundAnalyzer for tighter bounds when min_value equals max_value.
- Introduced ComputeGCD function to calculate the GCD of two integers.
- Updated Combine functions in IntervalSet to accept operation nodes for better type handling.
- Enhanced tests for modular set bounds in both const integer bounds and interval sets.
* replace gcd compute with ZeroAwareGCD
* doc op node
* replace Compute GCD with ZeroAwareGCD
* add example
* test fix
* test fix
* lint fix
This commit adds a new IsBound method to the ConstIntBoundAnalyzer class
that allows checking whether a variable is bound to a range. The method
is exposed through both the C++ and Python APIs.
Changes:
- Add IsBound method to ConstIntBoundAnalyzer C++ class
- Expose IsBound through FFI as is_bound_const_int
- Add Python binding with proper type hints
- Add test case to verify the functionality
- Improve type annotations throughout the analyzer module
* [ARITH] Canonicalize mul-coefficient to rhs
This PR updates the rewrite simplify logic to canonicalize mul-coefficient to rhs.
This change is consistent with rest of the code base and allows better simplification
of more cases. A test case of floormod with linear offset is added.
Co-authored-by: Ghosts381937 <she.xing88@gmail.com>
* Fix the grad testcase caused by the changed simplification behavior
---------
Co-authored-by: Ghosts381937 <she.xing88@gmail.com>
Fix a bug in canonical simplification of less-than expressions where
the algorithm incorrectly assumed variables could have negative values
when simplifying expressions of the form `ax + b < c`.
The previous implementation checked if `-d < xn < d` before simplifying,
but this was incorrect when variables are constrained to non-negative
ranges. For example, with constraints `0 < x, y < 2` and expression
`2x + y < 8`, the algorithm would incorrectly check if `-2 < y < 2`
and then simplify to `x < 4`. However, when x=4 and y=-1, we get
2*4 + (-1) = 7 < 8, which satisfies the original constraint but
violates the intended variable bounds.
The fix changes the range check to `0 <= xn < d`, ensuring that
simplification only occurs when variables are properly bounded
from below at zero.
Co-authored-by: FeiyangChen <92138383+smallscientist1@users.noreply.github.com>
This PR removes legacy tir::any which was used to represent unknown
shape in relay. As we move toward first class symbolic shape, we no longer
need the ? shape in the system.
If an iter has been normalized with a lower bound, and then try to normalize with
a new lower bound, the iter_min need to be updated only when the new lower bound
is smaller than the original one.
Co-authored-by: liujiaqiang <liujiaqiang@kuaishou.com>
* 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
* [TIR][Analyzer] Simplify `x==x` expressions for all dtypes
Prior to this commit, there was no rule to simplify `x == x` into
`True`. In some cases, despite not having an explicit rewrite rule in
`RewriteSimplifier`, the `RewriteSimplifier::CanProve` function would
check if `x-x` simplifies to zero, relying on the rewrite rules used
for `tir::Sub`. However, the rule to rewrite `x-x` into zero was only
enabled for `int32`, `int64`, and floating-point types, so relying on
this behavior was inconsistent.
This commit updates the rewrite rules for both `tir::EQ` and
`tir::Sub` to check for simplification of `x-x` or `x==x`, regardless
of the datatype. This change preserves the fast-path for index
data-types, in which `int32` and `int64` expressions may be simplified
without checking for side effects. For all other dtypes, the
cancellation only applies when evaluating `x` has no side effects.
* Add comment about simplifications of NaN/Inf
* [UnitTests] Use tvm.ir.assert_structural_equal whenever possible
Prior to commit, many unit tests were implemented as `assert
tvm.ir.structural_equal(output, expected)`. While this is correct, it
doesn't provide much information when the test fails. The
`tvm.ir.assert_structural_equal` method performs the equivalent check,
but displays the exact location where a mismatch occurs.
This commit replaces all use of `assert tvm.ir.structural_equal` with
`tvm.ir.assert_structural_equal`.
* fix unit tests
This commit adds a scalable `arm_cpu` conv2d NHWC schedule for fp32 which generates SME instructions by using the tensor intrinsics introduced in #16921.
Alongside the SME schedule, the logic of the TE schedule `schedule_conv2d_gemm_native()` for both non-scalable and scalable vector implementations has also been translated into the new TIR schedule. This means that the TE compute definition `compute_conv2d_NHWC_hybrid()` is now compatible with both the original TE schedules (e.g. `schedule_conv2d_NHWC_hybrid()`) and the newly introduced TIR schedule `schedule_conv2d_NHWC_hybrid_TIR()`. The corresponding TOPI test has been extended to reflect that.
This commit adds an `arm_cpu` conv2d NHWC schedule which generates SVE instructions by extending the hybrid GeMM approach implemented in #16106 to use scalable expressions as splitting factors.
Various vscale-related fixes needed to implement the schedule are also included, such as:
- adding vscale bounds in the `ConstIntBoundAnalyzer` and `IntervalSetEvaluator`
- simplifying `MinNode` and `MaxNode` that have scalable expression operands in `RewriteSimplifier`, which would appear when defining the shape of a buffer padded to be a multiple of vscale and in its respective buffer access indices (e.g. `C_1 = T.Buffer((1024 * (T.vscale() * 16 + 256 - 16 % T.vscale() * 16),), data=C)` instead of `C_1 = T.Buffer((1024 * (T.max(255, T.vscale() * 16 + 255 - 16 % T.vscale() * 16) + 1),), data=C)`)
The correctness of the new schedule is checked using a TOPI test, while the presence of generated SVE instructions is verified by a codegen_aarch64 test. The new rewrite_simplify rules are also covered by additional test cases.
This commit adds support for splitting via the compile-time unknown
constant `vscale`. Two main changes are introduced; they are described
below.
The split scheduling primitive has a new parameter disable_predication
that allows the user to avoid introducing a block-level predicate when
splitting with a factor of `vscale`. This feature is useful when schedule
writers know that the loop they're splitting is a factor of the scalable
vector length for their target. Otherwise, a predicate must be introduced
due to the nature of `vscale`.
CanProve has been extended to prove expressions that use multiple
instances of `vscale`. Known possible scalar values of the `vscale`
intrinsic are iterated over and substituted into the expression. If
the expression holds true for each possible value, we can conclude the
expression true. Currently only support for an SVE target has been
added, but it is possible to extend to other targets as/when needed. If
the analyzer becomes more powerful in the future and is able to deal
with multiple instances of a symbolic value in an expression, this
feature can be removed.
---------
Co-authored-by: Elen Kalda <elen.kalda@arm.com>
Co-authored-by: Neil Hickey <neil.hickey@arm.com>
Tianqi Chen