[CINN]Apply broadcast device lowering for too many broadcast tree bugs (#66207)

Hongqing-work · web-flow · commit 985b7eb33f27 · 2024-07-22T16:31:18.000+08:00
diff --git a/paddle/cinn/common/broadcast_tree.cc b/paddle/cinn/common/broadcast_tree.cc
@@ -18,11 +18,8 @@
 #include <unordered_map>
 
 #include "paddle/common/enforce.h"
-#include "paddle/common/flags.h"
 #include "paddle/pir/include/dialect/shape/utils/dim_expr_util.h"
 
-COMMON_DECLARE_int64(pir_broadcast_tree_limit);
-
 namespace cinn::common {
 
 namespace {
@@ -95,9 +92,6 @@ template <typename DoEachT>
 bool SearchBroadcastImpl(const symbol::Broadcast<symbol::DimExpr>& variadic,
                          const DoEachT& DoEach) {
   const auto& operands = *(variadic.operands);
-  if (operands.size() > 3) {
-    PADDLE_THROW(phi::errors::Fatal("Too many broadcast leaves to compile!"));
-  }
   for (const auto& operand : operands) {
     CHECK(!operand.isa<int64_t>());
     if (SearchBroadcast(operand, DoEach)) return true;
@@ -310,13 +304,13 @@ std::optional<symbol::Broadcastable<symbol::DimExpr>> GetFirstCstrBroadcastable(
 
 BroadcastTree ConstructBroadcastTree(const BroadcastLeaf& leaves,
                                      int* num_of_leaves) {
+  if (*num_of_leaves > FLAGS_pir_broadcast_tree_limit) {
+    return leaves;
+  }
   std::optional<symbol::Broadcastable<symbol::DimExpr>>
       broadcastable_condition = GetFirstCstrBroadcastable(leaves);
   if (!broadcastable_condition.has_value()) {
     (*num_of_leaves)++;
-    if (*num_of_leaves > FLAGS_pir_broadcast_tree_limit) {
-      PADDLE_THROW(phi::errors::Fatal("Too many broadcast leaves to compile!"));
-    }
     return leaves;
   }
   return ConstructBroadcastBranch(
diff --git a/paddle/cinn/common/broadcast_tree.h b/paddle/cinn/common/broadcast_tree.h
@@ -15,8 +15,11 @@
 #pragma once
 
 #include "paddle/cinn/adt/tree.h"
+#include "paddle/common/flags.h"
 #include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
 
+COMMON_DECLARE_int64(pir_broadcast_tree_limit);
+
 namespace cinn::common {
 
 template <typename T>
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc
@@ -15,6 +15,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/utils.h"
 #include "paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
@@ -26,6 +27,8 @@ using OpLoweringGroupPtr = std::shared_ptr<OpLoweringGroup>;
 using cinn::dialect::ir::details::CompileGroupAsOpAttribute;
 using cinn::dialect::ir::details::GetBlockOutsideInput;
 
+PD_DECLARE_bool(cinn_bc_branch_optimize);
+
 namespace {
 std::vector<pir::Value> GetOpOuputValues(const pir::Operation* op) {
   std::vector<pir::Value> outputs;
@@ -438,15 +441,19 @@ void SimplyConditionBlock(
 
 namespace cinn::dialect::ir::details {
 
-std::shared_ptr<BroadcastTree> ConstructBroadcastTree(
+std::optional<std::shared_ptr<BroadcastTree>> ConstructBroadcastTree(
     const cinn::common::BroadcastLeaf& leaves) {
   VLOG(6) << "before constructed. broadcast-leaf: \n"
           << ToTxtString(cinn::common::BroadcastTree(leaves));
   int num_of_leaves = 0;
   auto broadcast_tree = std::make_shared<cinn::common::BroadcastTree>(
       cinn::common::ConstructBroadcastTree(cinn::common::BroadcastLeaf(leaves),
                                            &num_of_leaves));
-  VLOG(4) << "num of broadcast tree leaves:" << num_of_leaves;
+  if (num_of_leaves > FLAGS_pir_broadcast_tree_limit) {
+    LOG(WARNING) << "the number of leaf nodes in broadcast tree exceeds "
+                    "limit.";
+    return std::nullopt;
+  }
   VLOG(4) << "broadcast-tree: \n" << ToTxtString(*broadcast_tree);
   return broadcast_tree;
 }
@@ -478,13 +485,24 @@ GroupDimExprInfo GetGroupDimExprInfo(const OpLoweringGroupPtr& group) {
   return group_dim_expr_info;
 }
 
-bool NeedBroadcastWithCF(const OpLoweringGroupPtr& group) {
-  GroupDimExprInfo group_dim_expr_info = GetGroupDimExprInfo(group);
-  const auto& leaves = group_dim_expr_info.all_value_dim_exprs;
-  return NeedBroadcastWithCF(leaves);
+std::optional<std::shared_ptr<BroadcastTree>> GetBroadcastTreeForOptimize(
+    const OpLoweringGroupPtr& group) {
+  if (!FLAGS_cinn_bc_branch_optimize) return std::nullopt;
+
+  const common::BroadcastLeaf leaves = [&]() {
+    // NOTE(dev): Need UpdateShapeOrDataExprs firstly and the logic
+    // will be migated into BucketLower later.
+    UpdateGroupShapeOrDataExprs(const_cast<OpLoweringGroupPtr&>(group));
+    GroupDimExprInfo group_dim_expr_info = GetGroupDimExprInfo(group);
+    return group_dim_expr_info.all_value_dim_exprs;
+  }();
+
+  if (!ContainBroadcastShape(leaves)) return std::nullopt;
+
+  return ConstructBroadcastTree(leaves);
 }
 
-bool NeedBroadcastWithCF(const cinn::common::BroadcastLeaf& leaves) {
+bool ContainBroadcastShape(const cinn::common::BroadcastLeaf& leaves) {
   std::optional<symbol::Broadcastable<symbol::DimExpr>>
       broadcastable_condition = cinn::common::GetFirstCstrBroadcastable(leaves);
   return broadcastable_condition.has_value();
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.h b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.h
@@ -27,11 +27,12 @@ struct GroupDimExprInfo {
   std::unordered_map<pir::Value, size_t> value_to_dim_expr_idx;
 };
 
-std::shared_ptr<BroadcastTree> ConstructBroadcastTree(
+std::optional<std::shared_ptr<BroadcastTree>> ConstructBroadcastTree(
     const common::BroadcastLeaf& leaves);
 
-bool NeedBroadcastWithCF(const OpLoweringGroupPtr& group);
-bool NeedBroadcastWithCF(const common::BroadcastLeaf& leaves);
+std::optional<std::shared_ptr<BroadcastTree>> GetBroadcastTreeForOptimize(
+    const OpLoweringGroupPtr& group);
+bool ContainBroadcastShape(const common::BroadcastLeaf& leaves);
 GroupDimExprInfo GetGroupDimExprInfo(const OpLoweringGroupPtr& group);
 
 pir::Operation* CompileBroadcastTreeToConditionBlock(
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/lower_cinn_fusion_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/lower_cinn_fusion_op_pass.cc
@@ -28,24 +28,17 @@
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/pass/pass_registry.h"
 
-PD_DECLARE_bool(cinn_bc_branch_optimize);
-
 namespace cinn::dialect::ir::details {
 
 pir::Operation* ProcessDyShapeGroup(const OpLoweringGroupPtr& group,
                                     pir::PatternRewriter& rewriter) {  // NOLINT
-  // NOTE(dev): Need UpdateShapeOrDataExprs firstly and the logic
-  // will be migated into BucketLower later.
-  UpdateGroupShapeOrDataExprs(const_cast<OpLoweringGroupPtr&>(group));
-  auto group_inputs = GetBlockOutsideInput(group->ops());
-  GroupDimExprInfo group_dim_expr_info = GetGroupDimExprInfo(group);
-  const auto& leaves = group_dim_expr_info.all_value_dim_exprs;
-  // has multiple branch
-  if (FLAGS_cinn_bc_branch_optimize && NeedBroadcastWithCF(leaves)) {
-    const auto& value_to_dim_expr_idx =
-        group_dim_expr_info.value_to_dim_expr_idx;
+  const auto& group_inputs = GetBlockOutsideInput(group->ops());
+  const auto& optional_broadcast_tree = GetBroadcastTreeForOptimize(group);
+  if (optional_broadcast_tree.has_value()) {
     const std::shared_ptr<BroadcastTree> broadcast_tree =
-        ConstructBroadcastTree(leaves);
+        optional_broadcast_tree.value();
+    const auto& value_to_dim_expr_idx =
+        GetGroupDimExprInfo(group).value_to_dim_expr_idx;
     std::vector<pir::Type> output_types;
     auto group_output_values = group->GetGroupOutputValues();
     for (size_t i = 0; i < group_output_values.size(); ++i) {
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/pre_analysis.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/pre_analysis.cc
@@ -51,7 +51,9 @@ void FusionOpAnalysis::PreCompileGroup() {
 
   std::vector<OpLoweringGroupPtr> groups;
   for (auto& group_info : *group_infos_) {
-    if (is_dy_shape_ && NeedBroadcastWithCF(group_info.second)) continue;
+    if (is_dy_shape_ &&
+        GetBroadcastTreeForOptimize(group_info.second).has_value())
+      continue;
     groups.push_back(group_info.second);
   }
   // Build and trigger compilaion cache.
