Only check constant cases

apache · Sep 14, 2022 · 23dca66 · 23dca66
1 parent 62a6803
commit 23dca66
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
@@ -1667,8 +1667,11 @@ 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
+    // if (cond) { if (inner_cond) { inner_then_expr } else { inner_else_expr } } else { else_expr }
+    // => if (cond && inner_cond) { inner_then_expr } else { else_expr }
     const PrimExpr& cond = op->args[0];
     const PrimExpr& then_expr = op->args[1];
     const PrimExpr& else_expr = op->args[2];
@@ -1677,9 +1680,10 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const CallNode* op) {
       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(),
-                    {cond && inner_cond, inner_then_expr, else_expr});
+      // Only check constant cases to avoid recursion
+      if (is_const_number(inner_else_expr) && is_const_number(else_expr) &&
+          analyzer_->CanProve(inner_else_expr == else_expr)) {
+        return if_then_else(cond && inner_cond, inner_then_expr, else_expr);
       }
     }
   }