【CINN】Simplify predicate and remove unreachable branch (#72316)

* simplify predicate after lowering

* remove useless comment

* remove useless comment

* fix undefine bug

* remove unreachable branch

Author: liuruyan

paddle/cinn/hlir/framework/pir/compilation_task.cc

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/cinn/hlir/framework/pir/compilation_task.h"
+#include "paddle/cinn/ir/utils/ir_copy.h"
 
 #include "paddle/cinn/backends/codegen_device_util.h"
 #include "paddle/cinn/common/dim_expr_converter.h"
@@ -241,6 +242,74 @@ void CompilationTask::Lowering() {
 
     context_->broadcast_condition_ = ChangeBroadcastConditionToExpr();
   }
+
+  auto SimplifyPredicate = [](GroupCompilationContext* context) {
+    for (auto& expr : context->predicates_) {
+      optim::SimplifyLogical(&expr);
+    }
+    if (context->broadcast_condition_.defined())
+      optim::SimplifyLogical(&context->broadcast_condition_);
+    for (auto& expr : context->CX86_predicates_) {
+      optim::SimplifyLogical(&expr);
+    }
+  };
+
+  // remove unreachable predicates.
+  auto RemoveUnreachPredicate = [](GroupCompilationContext* context) {
+    // remove unreachable predicate.
+    std::vector<ir::Expr> new_predicates;
+    std::vector<int> new_priorities;
+    std::vector<ir::LoweredFunc> new_lowered_funcs;
+    bool has_true_predicate = false;
+    for (size_t i = 0; i < context->predicates_.size(); ++i) {
+      if (has_true_predicate) continue;
+      if (common::IsZero(context->predicates_[i])) continue;
+      if (common::IsOne(context->predicates_[i])) has_true_predicate = true;
+      new_predicates.push_back(context->predicates_[i]);
+      new_priorities.push_back(context->priorities_[i]);
+      new_lowered_funcs.push_back(context->lowered_funcs_[i]);
+    }
+    // CINN does not support returning an empty module now. if all predicates
+    // are false, we push the first predicate as result.
+    if (new_predicates.empty() && !context->predicates_.empty()) {
+      new_predicates.push_back(context->predicates_[0]);
+      new_priorities.push_back(context->priorities_[0]);
+      new_lowered_funcs.push_back(context->lowered_funcs_[0]);
+    }
+    context->predicates_ = std::move(new_predicates);
+    context->priorities_ = std::move(new_priorities);
+    context->lowered_funcs_ = std::move(new_lowered_funcs);
+
+    // remove unreachable CX86 predicate.
+    std::vector<ir::Expr> new_CX86_predicates;
+    std::vector<ir::LoweredFunc> new_CX86_lowered_funcs;
+    bool has_true_CX86_predicate = false;
+    for (size_t i = 0; i < context->CX86_predicates_.size(); ++i) {
+      if (has_true_CX86_predicate) continue;
+      if (common::IsZero(context->CX86_predicates_[i])) continue;
+      if (common::IsOne(context->CX86_predicates_[i]))
+        has_true_CX86_predicate = true;
+      new_CX86_predicates.push_back(context->CX86_predicates_[i]);
+      new_CX86_lowered_funcs.push_back(context->CX86_lowered_funcs_[i]);
+    }
+    // CINN does not support returning an empty module now. if all predicates
+    // are false, we push the first predicate as result.
+    if (new_CX86_predicates.empty() && !context->CX86_predicates_.empty()) {
+      new_CX86_predicates.push_back(context->CX86_predicates_[0]);
+      new_CX86_lowered_funcs.push_back(context->CX86_lowered_funcs_[0]);
+    }
+    context->CX86_predicates_ = std::move(new_CX86_predicates);
+    context->CX86_lowered_funcs_ = std::move(new_CX86_lowered_funcs);
+  };
+  // Logical Simplifysimplify predicates, such as:
+  // false && ... ==> false
+  // true || ...  ==> true
+  // 1 <= 1       ==> true
+  SimplifyPredicate(context_);
+  // Remove unreachable predicates, unreachable predicates means that predicate
+  // is false, or a true predicate already existed before.
+  RemoveUnreachPredicate(context_);
+
   VLOG(5) << "End to lowering: " << context_->PrintPredicate2Funcs();
 }
 

paddle/cinn/optim/ir_simplify.cc

@@ -256,8 +256,13 @@ struct SimplifyLogicalMutator : public ir::IRMutator<> {
       *expr = Expr(false);
       return;
     }
-    if (common::IsOne(node->a()) && common::IsOne(node->b()))
+    if (common::IsOne(node->a()) && common::IsOne(node->b())) {
       *expr = Expr(true);
+    } else if (common::IsOne(node->a())) {
+      *expr = node->b();
+    } else if (common::IsOne(node->b())) {
+      *expr = node->a();
+    }
     VLOG(7) << "End Visit And op: " << *expr;
   }
 
@@ -276,8 +281,13 @@ struct SimplifyLogicalMutator : public ir::IRMutator<> {
       VLOG(7) << "End visit Or op: " << *expr;
       return;
     }
-    if (common::IsZero(node->a()) && common::IsZero(node->b()))
+    if (common::IsZero(node->a()) && common::IsZero(node->b())) {
       *expr = Expr(false);
+    } else if (common::IsZero(node->a())) {
+      *expr = node->b();
+    } else if (common::IsZero(node->b())) {
+      *expr = node->a();
+    }
     VLOG(7) << "End visit Or op: " << *expr;
   }