This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
## 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
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
## Why
has_attr pattern matching failed for integer attributes because
StructuralEqual requires exact dtype match for IntImm,
causing mismatches between Python-specified pattern values and actual
operator attributes.
## How
- Migrate NN conv/pooling/grad attrs from Array<IntImm> to
Array<int64_t>
---------
Signed-off-by: Guan-Ming Chiu <guanmingchiu@gmail.com>
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
* cleanup relay c++
* [REFACTOR] Phase out relay c++ components
This PR phases out the relay C++ components and
simplifies the overall codegen runtime logic.
---------
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
This PR starts the step 0 to phase out relay from the current
development main branch. This PR focuses on the python
components of relay, autotvm, auto_scheduler. To make the change
manageable, we will also do followup steps on te.Schedule and
c++ components in followup PRs.
To continue support community members who depends on
legacy flows, the [v0.19.0](https://github.com/apache/tvm/tree/v0.19.0)
branch will continue contain these components.
As noted in [discussion on phasing out legacy components](https://discuss.tvm.apache.org/t/phasing-out-legacy-components/17703/30),
this would help us to do two purposes:
- By removing outdated or redundant elements, we can significantly
reduce complexity and improve maintainability.
- Unify our focus: Concentrating our efforts on the new unity flow
will allow for more efficient development and innovation.
It is also a good opportunity for us to revisit and reduce CI time.
The past relay legacy flow contains a lot of end to end tests that
requires hardware resources to run and causing long CI time.
Moving onwards, we can focus more on unit-tests that focuses
on structural equality and runs within seconds, while be mindful
about tests that requires hardware resources (by restricting them
to specific folders and CI nightly in some cases).
---
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
Prior to this commit, the Relax well-formed checker validated
arguments provided to Relax functions, but did not validate arguments
provided to `R.call_tir`. As a result, incorrect arguments from Relax
to TIR would not be checked until runtime, if at all.
This commit updates the well-formed checker to verify that
`R.call_tir` has received the correct arguments, and has the correct
output shape specified in the `out_sinfo` parameter.
Initial implementation performed the validation as part of
`FNormalize`, to maximize coverage of this check. This increased
end-to-end compilation time by ~10%, and so the check was requested to
be restricted to the well-formed checker. Expensive operator-specific
validation is now performed in the new `FValidate` attribute.
Prior to this commit, the `tvm.relax.dpl.rewrite_bindings` utility
would segfault if its input contained a `DataflowBlock` whose first
binding was a `MatchCast`.
The root cause is use of an unintialized `const VarNode* cur_user_;`
when collecting the variable usage. This variable is only initialized
for `VarBinding` nodes, and may be used uninitialized if a `MatchCast`
node is encountered before the first `VarBinding`. This uninitialized
value is later dereferenced during while pattern-matching, causing a
segfault.
This commit provides a default value of `nullptr` for
`MatcherUseDefAnalysis::cur_user_`, preventing the segfault.
* [Relax] Express dynamic arguments of strided_slice as arguments
Prior to this commit, `relax.op.strided_slice` stored the `axes`,
`begin`, `end`, and `strides` in the `CallNode::attrs`. However, the
attributes are only intended to store static values. The indices used
used for `relax.op.strided_slice` must frequently be in terms of
symbolic shape variables, which should not be stored in the
attributes. While some utilities have special handling for
`relax.op.strided_slice` (e.g. `tvm::relax::Bind`), many do
not (e.g. `tvm::relax::WellFormed` and
`tvm::relax::FreeSymbolicVars`). As a result, the symbolic
expressions in `relax.op.strided_slice` will fail to be updated in
generic utilities, and will fail to trigger safeguards when this
occurs.
This commit changes the representation of `relax.op.strided_slice` to
store all arguments in the `relax::CallNode::args`, rather than the
`relax::CallNode::attrs`. As mentioned in a comment from
https://github.com/apache/tvm/pull/13987, which initially implemented
`relax.op.strided_slice`, this was an intended refactor once
`relax::PrimValue` was fully supported.
* Undo unnecessary changes in const_int_bound
* Remove unnecessary changes to rewrite_simplify
* lint fixes
* Fix unit tests
* Improve error message
* Fix additional unit tests
* Mark MSC tests with xfail
* remove commented-out code
* Resolve failing unit test
* Remove unused imports
* [Relax][Bugfix] Provide the full Expr to pattern-match rewriter
This resolves a bug that was introduced in
https://github.com/apache/tvm/pull/16732. If a rewriter function
returned a no-op, and the pattern-match continued, then the `matches`
provided to the rewriter function in subsequent calls would contain
a variable to which the matched expression was bound, not the matched
expression itself. (e.g. For a match of `C = R.add(A,B)`, passing `C`
to the rewriter instead of `R.add(A,B)`.)
This bug was caused by incorrect re-wrapping of `OrPattern` in
`ExprPatternRewriter`. Prior to
https://github.com/apache/tvm/pull/16732, all pattern-match results
were populated by `ExtractMatchExpr`, and contained the result after
applying `TryGetValOfVar`. When re-wrapping the result of an
`OrPattern`, https://github.com/apache/tvm/pull/16732 populated the
additional matches with the result before applying `TryGetValOfVar`.
This commit fixes the bug by applying `TryGetValOfVar`.
* Update with PR link of bugfix
[Relax] Allow composition of DFPattern replacements
The `rewrite_call` function accepts a `DFPattern`, and a function to
rewrite expressions matching that pattern. Often, the rewriting
function will perform additional validation that cannot be expressed
within the `DFPattern` itself. If this additional validation fails,
the rewriter function will return the matched expression unmodified.
Prior to this commit, an `OrPattern` that matches on the first branch,
but whose rewriter function does not apply a modification, would
prevent the second branch from being checked. This commit updates the
`ExprPatternRewriter` to check both branches of a `OrPattern`, if the
rewriter function of the first branch does not modify the result.
* Check well-formedness in the parser
* Correct packed funcs in NN frontend
* Support the check_well_formed optional argument to I.ir_module
* Also check well-formedness in TIR
* Enable normalization for individual Relax functions and PrimFuncs
* Use the error raised by the TIR well-formed checker for the message
* Fix tvmscript test failures
* Whitespace
* Fix errors in verify_well_formed test
* Include a more helpful error message
* Fix TIR test failures
* Address well-formed failures in test_tir_specialize
* Correct well-formedness error in test_tir_analysis_oob
* Correct further well-formedness failures
* Remove __tvm_meta__ from test case to avoid parsing error
* Avoid circular import in entryy.py
* Formatting fixes
* lint fix
* Add pylint exceptions
* Fix whitespace
* Fix more failed test cases
* Catch inappropriate use of decl_function instead of segfaulting
* Fix test_lower.py
* Mark purity in test_relax_2d_buffer_allocation.py
* Mark purity in test_dma_builtin.py
* Remove __tvm_meta___ from test_tir_usmp_analysis_extract_bufferinfo.py
* Suppress well-formed check in test_tir_transform_convert_blocks_to_opaque.py
* Remove __tvm_meta__ in test_tir_usmp_algo.py
* Remove __tvm_meta__ from more USMP tests
* Fix incorrect var in test_tir_transform_storage_flatten.py
* Remove all remaining instances of __tvm_meta__
* Fix purity error in test_dataflow_pattern.py
* Fix purity error in test_ast_printer
* Fix test_arith_domain_touched example
* Okay to set check_well_formed to True in test_tir_analysis_identify_mcmcpy
* Define variable in test_tir_analysis_oob
* Typo fix
* Add explanatory comment to test case
* Define the undefined vars in test_tir_transform_common_subexpr_elim
* Exception no longer necessary in test_tir_transform_inject_rolling_buffer
* Remove unnecessary check exemption in test_tir_transform_convert_ssa
* Avoid checking exemption in test_inject_ptx_ldg32
* Note special case in test_distributed_transform_propagate_sharding
* Exempt well-formed error in dlight/test_benchmark
* Exempt well-formedness errors in test_ethosu/, mostly uninitialized vars
* Whitespace
* Include non-CUDA GPUs in IsScheduledOnGPU
* Fix thread binding bug by changing thread binding var dtype
* Include overrides in test_runtime_builtin_paged_attention_kv_cache.py
* add exemptions in test_ethosu/test_replace_conv2d
* Add more ethosu exemptions
* More exemptions for ethosu tests
* Remove unused reference
* Indicate purity in test_transform_rewrite_cuda_graph
* Indicate purity in test_transform_normalize
* Reorder MergeSharedMemoryAllocations in GPU codegen
* Add target parameter for FP8StorageLegalize and FP8ComputeLegalize
* Don't re-import Target in tvm/tir/transform/transform.py
This commit implements `StructInfoPattern`, which can be applied to
any existing `DFPattern`, and requires the expression to have a
specific struct info. Any symbolic variables that occur in the struct
info are treated as free parameters, to be defined by the match.
Before this PR, dynamic shapes require upper bound of
variables to be provided in order to use storage planning.
We can relax this requirement, for shapes with unknown bound,
we can look up other tensors with the same symbolic
shapes. This can be helpful for deep learning models
where the layers with the same configurations are usually
repeated since there are many objects with the same shapes.
This PR changed the `StorageToken` to use `PrimExpr`
bytes which can be integer or symbolic. For symbolic
shapes, we put the tokens into a special buckets for looking up.
`std::regex` in TVM codebase may cause a symbol conflict with PyTorch,
we temporarily disable it before we find a better solution, meanwhile
the current usage of `std::regex` is not necessary.
Prior to this commit, `relax.transform.LegalizeOps` needed to be
called prior to `relax.build`. This commit adds `LegalizeOps` to the
lowering flow, to simplify the calling steps for an end-user. If the
`IRModule` contains no legalizable functions, a second legalization
pass has no effect.
Some test cases relied on this behavior as an implicit assertion that
operator fusion patterns applied. That is, by omitting `LegalizeOps`,
a successful compilation `relax.build` would only occur if all
legalizable operators have already been removed, and so an incorrect
fusion pattern would result in a failure to build the module. While
these tests would be better expressed by comparing against an expected
fused pattern, updating the tests is outside the scope of this PR. To
allow these tests to keep their implicit assertions, a
`"relax.transform.apply_legalize_ops"` config can be used to disable
the `LegalizeOps` pass.
* [Unity][Transform] Improved canonicalization of non-dataflow Var
Prior to this commit, `relax.transform.CanonicalizeBindings` removed
trivial bindings `var_y = var_x` where a `var_y: relax.DataflowVar`
and `var_x: relax.Var`, but did not remove trivial bindings when
`var_y: relax.Var` and `var_x: relax.DataflowVar`. This was to avoid
invalid use of a `relax.DataflowVar` outside of a dataflow block.
This commit updates `CanonicalizeBindings` to handle this type of
binding as well. To ensure that no `relax.DataflowVar` instances are
used outside of a dataflow block, this is done by replacing `var_y:
relax.DataflowVar` at its point of definition, instead of replacing
`var_x: relax.Var` at its point of use.
This commit also canonicalizes `relax.Var` definitions to
`relax.DataflowVar`, if the binding occurs within a dataflow block,
and the variable is never used outside of a dataflow block.
* Simplify unwrapping of known bindings
* Updated to use Map<Id,Var>, to avoid while(true) loops
* [Unity][Transform] Canonicalize and use CSE between pattern matches
The `PatternRewriter` is intended to iterate until no matching
patterns remain. Prior to this commit, this only involved repeating
the pattern match rewrite rules. However, intermediate results
produced by pattern replacement could cause the iterative pattern
matching to terminate early.
* If two rewrite rules each introduce the same intermediate, there
will exist two copies of that intermediate, which can prevent
`only_used_by` patterns from matching. Applying
`EliminateCommonSubexpr` allows the pattern matching to continue.
* Applying a rewrite rule may result in dangling intermediates that
are no longer used. These dangling intermediates may prevent the
next application of a rewrite rule that uses the `only_used_by`
constraint. Applying `RemoveAllUnused` allows the pattern matching
to continue.
* A rewrite rule that returns a `relax::Var` or `relax::TupleGetItem`
as the replacement introduces trivial var-to-var rebinding, which
are not tracked by `PatternRewriter`. Applying
`CanonicalizeBindings` allows the pattern matching to continue.
While this could be fixed externally by repeatedly applying
`rewrite_call`, this would require re-inspecting the entire function,
and not just the dataflow block in which the replacement occurred.
* Fix tests for removing redundant reshapes
* Fixed failing unit tests, along with edge case in CSE
* [Unity] Implemented SameShapeConstraint for dataflow pattern matches
Prior to this commit, a shape could be explicitly specified using the
`ShapePattern`, but could not be specified relative to the shape of
another expression. As a result, patterns with restricted shapes
became very difficult to
This commit implements `SameShapeConstraint`, which can be applied
between any patterns that participate in the match. For example,
matching against `R.add(lhs,rhs)` where `lhs` and `rhs` have the same
non-broadcasted shape. Because these constraints operate between
patterns that do not necessarily share a consumer/producer
relationship, they could not previously be expressed using the
existing `PairCons` functionality, and were instead implemented in
terms of a new `DFConstraint` base class.
* Removed empty line
* Update API based on PR discussion
Provide a `AsPrimExpr` instead of `IsConstraintSatisfied` function for
constraints. A constraint can return a necessary-and-sufficient
condition, or a necessary-but-not-sufficient condition, and the
calling scope can decide how to interpret those results.
* Lint fix
* lint fixes
* [Unity] Avoid trivial `var2 = var1` bindings in pattern matcher
Prior to this commit, pattern matches that returned a
variable (e.g. rewriting `R.add(x, R.const(0.0))` into `x`) would
result in rewriting `y = R.add(x, R.const(0.0))` to the trivial
binding `y = x`. This commit updates the pattern matcher to instead
remove the `y` variable entirely in this case.
* Updated unit test to preserve dataflow, avoid infinite replacement
Prior to this commit, pattern matches were evaluated a single time, in
the order of the block bindings. If a pattern rewrite produced a node
which would also match the pattern (e.g. recursive reordering), the
pattern would still only be evaluated once.
This commit updates `PatternMatcher` to iterate until convergence when
provided with a pattern, analogous to its current behavior when
called through `rewrite_bindings`.
* [Unity] Update relax.Function.ret_struct_info when mutated
Prior to this commit, the `relax::ExprMutator` forwards the
original `ret_struct_info` when visiting a `relax::Function`. While
this allows the resulting function to have a more specific return type
if the mutation no longer allows the shapes to be propagated across
the entire body (e.g. preserving `R.Tensor(shape=[16,16])` even if
shape propagation resulted in `R.Tensor(ndim=2)`), this also preserves
information that is no longer correct due to the mutation.
For example, a mutator that implements `VisitVarDef_` to replace the
shape of a function parameter would expect that updated shape to
propagate through to the function return type. By retaining the
original return type, the mutator produces an incorrect return type.
This commit updates `relax::ExprMutator` to only forward the original
`ret_struct_info` if the mutated body's struct info is compatible with
it.
* Test updates with ret struct info
---------
Co-authored-by: Farshid <fparizi@octoml.ai>
Prior to this commit, the commutative pattern matching was enabled
based on the operation in the pattern. As a result,
commutative matches would only be checked if the match checked for a
single operator, but not if the operator was itself a pattern that
resolved to a commutative operator.
```python
pattern_add = ExprPattern(Op.get("relax.add"))
pattern_mul = ExprPattern(Op.get("relax.multiply"))
uses_commutative_matching = pattern_add(lhs, rhs)
no_commutative_matching = OrPattern(pattern_add, pattern_mul)(lhs, rhs)
```
This commit updates the pattern matcher to check against the matched
operator, rather than the pattern, to determine whether to check for
commutative matches.
This PR implements the privacy annotation proposal. Namely, the @R.function decorator now has an optional private attribute. If a function is marked as private, then it will not have a global symbol attached to it and thus will not be externally accessible. By default, functions are not private, so the parser does insert a global symbol for them.
* remove start_hint from MatchGraph
* improve graph matching algorithm
* remove side effect from matching algo
* pylint
* add comments
* add more const now that we can
* cpplint
* fix compile warning
* Update src/relax/ir/dataflow_matcher.cc
Co-authored-by: Jiawei Liu <jaway.liu@gmail.com>
* Update src/relax/ir/dataflow_matcher.cc
Co-authored-by: Jiawei Liu <jaway.liu@gmail.com>
* use insert for merging MatchState
* fix
* parent check is not specific to wildcard
* use map merge
* cpplint
* Pass and check current_match in TryMatch
---------
Co-authored-by: Jiawei Liu <jaway.liu@gmail.com>
* stub
* wip
* works
* restore binding
* attention test work
* use RemoveAllUnused
* simplified callback api
* pass original call node to callback
* clean test
* add doc
* add test for the case where the original call is returned
* callback -> rewriter and other doc improvement
The current cross-function calls in TVMScript will cause PyLint warnings,
since the GlobalVar will be marked as undefined vars, e.g.:
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
gv0 = R.call_tir(tir_func, x, R.Tensor((128,), dtype="float32")) # <= `tir_func` is not defined in Python syntax.
return gv0
```
This PR changes the behavior into `TestModule.tir_func` instead of direct `tir_func`
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
cls = TestModule # Use `cls` to refer the current Module
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
return gv0
```
NOTE: It's a breaking change, the old style is deprecated.
Additionally, this PR contains the following minor fixes:
- mark `R.function` as staticmethod as what we do for `T.prim_func`
- make `I`, `R`, `T`, `cls` be the builtin keywords for the printer
- define names for functions, modules to prevent naming conflict
- checking the var names is valid via regex expression
- fix typos
The dataflow pattern language for Relax (originally from https://github.com/tlc-pack/relax/pull/163).
The implementation splits patterns into two parts:
- Match an Expression: match an expression syntactically (MatchExprPattern, i.e., DFPatternMatcher);
- Match a Graph: match a graph (cross multiple VarBinding) topologically (MatchGraphPattern);
Tianqi Chen