[Arith] Simplify nested if_then_else

Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
apache · Sep 12, 2022 · 175eb97 · 175eb97
1 parent a047e02
commit 175eb97
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
@@ -1667,6 +1667,22 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const CallNode* op) {
       return match.value();
     }
   }
+  if (op->op.same_as(tir::builtin::if_then_else())) {
+    // Simplify nested if_then_else
+    const PrimExpr& cond = op->args[0];
+    const PrimExpr& then_expr = op->args[1];
+    const PrimExpr& else_expr = op->args[2];
+    const CallNode* inner_call = then_expr.as<CallNode>();
+    if (inner_call != nullptr && inner_call->op.same_as(tir::builtin::if_then_else())) {
+      const PrimExpr& inner_cond = inner_call->args[0];
+      const PrimExpr& inner_then_expr = inner_call->args[1];
+      const PrimExpr& inner_else_expr = inner_call->args[2];
+      if (analyzer_->CanProveEqual(else_expr, inner_else_expr)) {
+        return Call(op->dtype, tir::builtin::if_then_else(),
+                    {VisitExpr(cond && inner_cond), inner_then_expr, else_expr});
+      }
+    }
+  }
 
   return ret;
 }

diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py
@@ -992,5 +992,15 @@ def test_sub_bufferload():
     ck.verify(expr, 0.0)
 
 
+def test_if_then_else_simplify():
+    ck = RewriteChecker()
+    x = te.var("x", "int32")
+    z = tvm.tir.if_then_else(x < 5, tvm.tir.if_then_else(x > 1, 1, 0), 0)
+    ck.verify(z, tvm.tir.if_then_else(tvm.tir.And(tvm.tir.LT(x, 5), tvm.tir.LT(1, x)), 1, 0))
+
+    z = tvm.tir.if_then_else(x > 2, tvm.tir.if_then_else(x > 1, 1, 0), 0)
+    ck.verify(z, tvm.tir.if_then_else(tvm.tir.LT(2, x), 1, 0))
+
+
 if __name__ == "__main__":
     pytest.main([__file__])