refactor ifthenelse optimize (#56274)

Courtesy-Xs · web-flow · commit 4526f61ef922 · 2023-08-15T14:24:03.000+08:00
diff --git a/paddle/cinn/optim/ir_simplify.cc b/paddle/cinn/optim/ir_simplify.cc
@@ -243,8 +243,31 @@ struct SimplifyIfThenElseMutator : public ir::IRMutator<> {
     auto* node = expr->As<ir::IfThenElse>();
     node->condition = common::AutoSimplify(node->condition);
 
-    if (node->true_case.defined()) Visit(&node->true_case, &node->true_case);
-    if (node->false_case.defined()) Visit(&node->false_case, &node->false_case);
+    auto* condition_int = node->condition.As<ir::IntImm>();
+    auto* condition_uint = node->condition.As<ir::UIntImm>();
+    int64_t value;
+    if (condition_int || condition_uint) {
+      if (condition_int) {
+        value = condition_int->value;
+      } else {
+        value = condition_uint->value;
+      }
+      if (value) {
+        *expr = op->true_case;
+      } else {
+        if (op->false_case.defined()) {
+          *expr = op->false_case;
+        } else {
+          // null condition
+          *expr = ir::Block::Make({});
+        }
+      }
+    }
+    if (expr->As<ir::IfThenElse>()) {
+      if (node->true_case.defined()) Visit(&node->true_case, &node->true_case);
+      if (node->false_case.defined())
+        Visit(&node->false_case, &node->false_case);
+    }
   }
 };
 
diff --git a/paddle/cinn/optim/optimize.cc b/paddle/cinn/optim/optimize.cc
@@ -79,8 +79,9 @@ Expr Optimize(Expr e,
   Simplify(&copied);
   VLOG(10) << "After Optimize Simplify:" << copied;
 
-  IfSimplify(&copied);
-  VLOG(10) << "After Optimize IfSimplify:" << copied;
+  // TODO(LiuYang): I attends to remove this part code, I integate it into
+  // ifthenelse part IfSimplify(&copied); VLOG(10) << "After Optimize
+  // IfSimplify:" << copied;
 
   if (runtime_debug_info) {
     LOG(WARNING) << "Turn on runtime debug information output";
