## 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.
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().
## 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 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>
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 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
* [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.
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 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
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.
This is a follow-up to https://github.com/apache/tvm/pull/16588. Due
to an incorrect rebase, the version that was merged into `main` had
the tighter `ConstIntBounds` enabled by default, rather than having
them implemented in `RewriteSimplifier`, gated behind a feature flag.
* [ARITH] Enhance Canonical Simplify for LE
This PR enhances the canonical simplifier to support the following patterns:
x0 * s0 + x1 * s1 + ... + xn + c < 0, let d = gcd(s0, s1, ..., s{n-1}, c)
1. if can prove -d < xn < d, then we can simplify
the expression to x0 * (s0/d) + x1 * (s1/d) + ... + x{n-1} * (s{n-1}/d) < c/d,
e.g. `x * 8 + y < 16` where `y` \in [0, 8), we can simplify it to `x < 2`
2. if xn is in pattern of yn % m, where m % d == 0, convert it to yn // d % (m/d)
e.g. `x1 * 64 + (x2 * 8 + x3) % 64 < 120`, `x3` \in [0, 8), we can simplify it to
`x1 * 8 + (x2 * 8 + x3) // 8 % 8 < 15` ==> `x1 * 8 + x2 % 8 < 15`
* [Arith] Implement statistics counters for RewriteSimplifier
Previously, so long as `RewriteSimplifier` produces the same output,
unit tests of its behavior would pass. This could have severe
performance regressions, such as the one resolved in
https://github.com/apache/tvm/pull/14528, which caused the runtime of
two test to increase from ~1.5 seconds to ~10 minutes each.
This commit implements statistics counts in RewriteSimplifier, which
are exposed through both the C++ and Python APIs, and uses these to
guard against the known performance regression from
https://github.com/apache/tvm/pull/14528.
* lint fixes
* Updates based on review comments
* Consistent int64_t with kMaxRecurDepth
* Removed unused is_currently_visiting_
* Add missing \brief for RewriteSimplifierStatsNode
* Use int64_t in ControlFlowGraph for max simplification steps
This PR refactors and enhances DetectIterMap and IterMapSimplify
to enable symbolic shape simplification. Specifically, we add
a routine to combine multiple IterSplitExpr into one if they
come from the same source.
It is helpful to distinguish iterator from normal constants
in the simplification process. IterMapSimplify takes advantage
of these information.
This improvements is helpful to simplify the indices in flattened buffer
when there is symbolic shape involved and normal simplifier.
Also updated FlattenBuffer to take benefit of the enhanced simplifier.
Test cases are added.
Co-authored-by: Junru Shao <junrushao@apache.org>
This PR enhances CanProve to handle symbolic bound.
Such analysis is essential to eliminate predicates in
dynamic shape workloads.
We also the int set analysis singlepoint check to avoid recursion
and improve the overall analysis speed.
Added CanProveSinglePoint to serve previous stronger checks.
The new CanProve comes with additinal strength argument
that can only be used in top-level setting with stronger analysis.
Added comment for future implementation efficiency.
Testcases are added to cover the cases.
Hi, this change wants to add some minor updation to region estimator used by buffer compaction:
- Add and clearify among `EstimateRegionStrictBound`, `EstimateRegionLowerBound` and `EstimateRegionUpperBound`
Originally we have `EstimateRegionLowerBound`, actually it implements strict bound estimation IMO. Now add `upper` and `strict` version for where we actually want them.
- When estimating upperbounds (eg. in buffer compaction), try estimate each dimension independently when they are dependent accesses where `EstimateRegionLowerBound` is expected to fail.
Eg, `A[i, i], 3 < i < 16` fails via `EstimateRegionLowerBound` who check indices be independent. But we can still try best to invoke strict bound analysis on each dimension individually.
- If range->extent == 1 for `EvalSet(range, dom)`, invoke `EvalSet(range->min, dom)` instead.
Eg, `EvalSet([k*k, k*k+1), dom_k)` results to [-inf, +inf] due to current algorithm limitation but `EvalSet(k*k, dom_k)` results to a range which makes more sense.
* simplify (x * 96) % 64 to (x * 32) % 64
* adapt merge mulmod opt for OffsetOf computation
* merge DetectIterMap and DetectIterMapPadded
* adjust related interfaces for IterMapLevel
* - check incompatible left paddings
- determine case like x % 16, x in [0, 5) to be non-surjective, since usages may treat the region extent as 16 by mistake.
- skip second round of rewrite when there is no padding
- fix some typo in comments
* rebase upstream
* [Analysis] Exposed Analyzer::CanProveEqual to Python API
Checking for `analyizer.simplify(lhs-rhs) == 0` was a frequent pattern
in Python unit tests, and already had a utility function in the C++
public API. Exposing this utility function to Python allowed this
pattern to be cleaned up.
* Replaced more cases of .simplify with .can_prove_equal
* [Arith] Updated arith::DetectIterMap to keep extent=1 components
Previously, arith::DetectIterMap simplified the output expression by
replacing iteration variables with extent==1 with their value. This
prevented the return value from being used in
arith::InverseAffineIterMap to solve for the variable, as it no longer
existed in the returned expressions.
This commit changes arith::DetectIterMap to keep the iteration
variable even if extent==1, and adds a motivating unit test that
requires this updated behavior.
* Updated to retain default behavior of DetectIterMap
To avoid breaking existing test cases, updated to maintain the same
default behavior, but a flag to maintain trivial iterators in the
result.
* Updated FFI and Python API for DetectIterMap
* [ARITH] Introduce iterator (quasi)affine map detection.
The loop transformations (split, fuse) create bijective
maps from a collection of source iterators to target iterators.
DetectIterMap is a function that detects such bijective mappings
from the lowered index expression.
We choose the term quasi affine to be consistent with the
terminology used by in polyhedral compilation.
DetectIterMap can handle symbolic integers(in split/fuse) to some extent.
The utility can be useful in detecting loop transformation
patterns and data layout change patterns in TIR.
* Update per feedback
* [arith] linear system and equation solver
Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>
* avoid constructing analyzer every time
* generate random test cases and address comments
Co-authored-by: Sergei Grechanik <sergei.grechanik@gmail.com>
* rename linear_system to int_constraints
* add comments and use random seed
* message for reporting failure with seed
* add SEqualReduce to IntConstraints; allow variables & ranges to be None
Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>
Co-authored-by: Sergei Grechanik <sergei.grechanik@gmail.com>
Tianqi Chen