This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c067b74367 [FIX][TIRx] Avoid stale constraints from mutable memory
predicates (#20393)
c067b74367 is described below
commit c067b743676508ddebcc58c77c219406b17fc6a7
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Sep 20 09:42:51 2026 -0400
[FIX][TIRx] Avoid stale constraints from mutable memory predicates (#20393)
This pr fixes incorrect simplification of conditions that read mutable
memory. An enclosing branch can currently cause `StmtSimplify` to remove
a nested loop bound even when stores inside the branch or loop change
the value being tested. Assertions can similarly cause later expressions
to be folded incorrectly.
Only record pure branch and assertion predicates as analyzer facts, and
exclude memory-dependent branch conditions from iteration predicates.
---
src/tirx/ir/ir_mutator_with_analyzer.cc | 19 +++++++--
.../tirx-transform/test_tir_transform_simplify.py | 46 ++++++++++++++++++++++
2 files changed, 61 insertions(+), 4 deletions(-)
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
index d4ca4797ec..6006f331b9 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -159,18 +159,26 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
IfThenElseNode* op,
}
}
+ // A branch only establishes a memory-dependent predicate at entry. Loads
+ // may change after stores, opaque calls, or another iteration of a nested
+ // loop, so they cannot be facts for the entire branch scope.
+ bool condition_is_pure = SideEffect(real_condition) <=
CallEffectKind::kPure;
Stmt then_case;
ffi::Optional<Stmt> else_case;
constraint_scope_.WithNewScope([&]() {
- EnterConstraintFacts(&constraint_scope_.Current(), analyzer_,
real_condition);
- WithRecordIterPredicate(real_condition, [&] {
+ if (condition_is_pure) {
+ EnterConstraintFacts(&constraint_scope_.Current(), analyzer_,
real_condition);
+ }
+ WithRecordIterPredicate(condition_is_pure ? real_condition :
IntImm::Bool(true), [&] {
then_case = this->Mutate(op->then_case,
inplace_mode).ValueOrUnchanged(op->then_case);
});
});
if (op->else_case) {
PrimExpr neg_condition =
analyzer_->rewrite_simplify(prim::Not(real_condition));
constraint_scope_.WithNewScope([&]() {
- constraint_scope_.Current().Emplace(analyzer_, neg_condition);
+ if (condition_is_pure) {
+ constraint_scope_.Current().Emplace(analyzer_, neg_condition);
+ }
else_case = this->Mutate(op->else_case.value(), inplace_mode)
.ValueOrUnchanged(op->else_case.value());
});
@@ -220,7 +228,10 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
AssertStmtNode* op,
auto condition_result = this->Mutate(op->condition, inplace_mode);
bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition);
PrimExpr condition =
std::move(condition_result).ValueOrUnchanged(op->condition);
- constraint_scope_.Current().Emplace(analyzer_, condition);
+ // Like branch predicates, assertions about mutable memory are snapshots.
+ if (SideEffect(condition) <= CallEffectKind::kPure) {
+ constraint_scope_.Current().Emplace(analyzer_, condition);
+ }
if (condition_unchanged) {
return ffi::Unchanged();
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py
b/tests/python/tirx-transform/test_tir_transform_simplify.py
index fc74467177..9f32aa1ed2 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import pytest
+
import tvm
import tvm.testing
from tvm.script import ir as I
@@ -1313,5 +1315,49 @@ def test_nested_if_elimination():
tvm.ir.assert_structural_equal(after, expected)
[email protected]("else_branch", [False, True])
[email protected]("write_before_loop", [False, True])
+def test_mutable_branch_predicate_preserves_while_bound(else_branch,
write_before_loop):
+ # Build both branch directions from the same body. A store before the loop
+ # and a store on its back edge both invalidate the entry predicate.
+ from tvm import tirx
+
+ x = tirx.decl_buffer((1,), "int32", name="x")
+ count = tirx.decl_buffer((1,), "int32", name="count")
+ loop = tirx.While(
+ T.And(x[0] < 8, count[0] == 0),
+ tirx.BufferStore(x, x[0] + 1, [0]),
+ )
+ body = tirx.SeqStmt([tirx.BufferStore(x, x[0] + 1, [0]), loop]) if
write_before_loop else loop
+ branch = (
+ tirx.IfThenElse(T.int32(8) <= x[0], tirx.Evaluate(0), body)
+ if else_branch
+ else tirx.IfThenElse(x[0] < 8, body, None)
+ )
+ func = tirx.PrimFunc([x, count], branch)
+ after = _apply_simplify(func)
+ tvm.ir.assert_structural_equal(after, func)
+
+
+def test_mutable_branch_predicate_preserves_later_load():
+ @T.prim_func(private=True, s_tir=True)
+ def before(x: T.Buffer((1,), "int32"), out: T.Buffer((1,), "int32")):
+ if x[0] < 8:
+ x[0] = x[0] + 1
+ out[0] = T.Select(x[0] < 8, 1, 0)
+
+ tvm.ir.assert_structural_equal(_apply_simplify(before), before)
+
+
+def test_mutable_assert_does_not_constrain_later_load():
+ @T.prim_func(private=True, s_tir=True)
+ def before(x: T.Buffer((1,), "int32"), out: T.Buffer((1,), "int32")):
+ assert x[0] < 8, "initial bound"
+ x[0] = x[0] + 1
+ out[0] = T.Select(x[0] < 8, 1, 0)
+
+ tvm.ir.assert_structural_equal(_apply_simplify(before), before)
+
+
if __name__ == "__main__":
tvm.testing.main()
