You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by sa...@apache.org on 2023/12/17 16:08:17 UTC
(tvm) branch main updated: [ARITH] Simplify nested if_then_else when constant is appearing in then_expr (#16227)
This is an automated email from the ASF dual-hosted git repository.
sanirudh 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 799e81036d [ARITH] Simplify nested if_then_else when constant is appearing in then_expr (#16227)
799e81036d is described below
commit 799e81036d43df552efdaecf75c680532ea85144
Author: rutkoor <12...@users.noreply.github.com>
AuthorDate: Sun Dec 17 21:38:11 2023 +0530
[ARITH] Simplify nested if_then_else when constant is appearing in then_expr (#16227)
Simplify nested if_then_else when constant is appearing in then_expr
---
src/arith/ir_mutator_with_analyzer.cc | 5 +++--
tests/python/tir-transform/test_tir_transform_simplify.py | 12 ++++++++++++
2 files changed, 15 insertions(+), 2 deletions(-)
diff --git a/src/arith/ir_mutator_with_analyzer.cc b/src/arith/ir_mutator_with_analyzer.cc
index 2ee427beb8..d26ac36676 100644
--- a/src/arith/ir_mutator_with_analyzer.cc
+++ b/src/arith/ir_mutator_with_analyzer.cc
@@ -173,8 +173,9 @@ PrimExpr IRMutatorWithAnalyzer::VisitExpr_(const CallNode* op) {
WithRecordIterPredicate(cond, [&] { true_value = this->VisitExpr(op->args[1]); });
}
{
- With<ConstraintContext> constraint(analyzer_, analyzer_->rewrite_simplify(Not(cond)));
- false_value = this->VisitExpr(op->args[2]);
+ PrimExpr not_cond = Not(cond);
+ With<ConstraintContext> constraint(analyzer_, not_cond);
+ WithRecordIterPredicate(not_cond, [&] { false_value = this->VisitExpr(op->args[2]); });
}
if (is_zero(cond)) {
return false_value;
diff --git a/tests/python/tir-transform/test_tir_transform_simplify.py b/tests/python/tir-transform/test_tir_transform_simplify.py
index c779d92f9c..6bad817c49 100644
--- a/tests/python/tir-transform/test_tir_transform_simplify.py
+++ b/tests/python/tir-transform/test_tir_transform_simplify.py
@@ -1757,5 +1757,17 @@ class TestBufferShapeConstraintWithOffset(BaseBeforeAfter):
A[T.int64(1)] = T.float32(0)
+class TestNestedIfElimination(BaseBeforeAfter):
+ def before(a: T.Buffer((2, 8), "int32"), b: T.Buffer((2, 8), "int32")):
+ for i0, j0 in T.grid(2, 8):
+ b[i0, j0] = T.if_then_else(
+ i0 == 1 and 6 <= j0, 0, T.max(0, T.if_then_else(i0 == 1 and 6 <= j0, 0, a[i0, j0]))
+ )
+
+ def expected(a: T.Buffer((2, 8), "int32"), b: T.Buffer((2, 8), "int32")):
+ for i0, j0 in T.grid(2, 8):
+ b[i0, j0] = T.if_then_else(i0 == 1 and 6 <= j0, 0, T.max(0, a[i0, j0]))
+
+
if __name__ == "__main__":
tvm.testing.main()