[CINN] Fix ReduceTree Fusion (#67400)

huangjiyi · web-flow · commit 0c34b1f59a00 · 2024-08-14T18:38:56.000+08:00
* update

* update

* update

* update
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
@@ -181,7 +181,6 @@ std::vector<ir::Var> GetOutputIters(const FusibleOp& op) {
     }
   };
   VLOG(4) << "GetOutputIters";
-  VLOG(4) << "Before AppendBound:" << _GetRootExpr(op);
   return AppendBound(std::visit(Visitor(), op), _GetRootExpr(op));
 }
 
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.cc b/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.cc
@@ -17,6 +17,15 @@
 
 namespace cinn::fusion {
 
+std::vector<ir::Expr> GetFusibleOpsExpr(std::vector<FusibleOp> fusion_ops) {
+  std::vector<ir::Expr> exprs;
+  for (auto& fusion_op : fusion_ops) {
+    auto expr = std::visit(FusibleOp2Expr(), fusion_op).front();
+    exprs.push_back(expr);
+  }
+  return exprs;
+}
+
 // tmp transform for reduce_tree and reduce_tree_trivial.
 std::vector<ir::Tensor> GetOutputTensors(const ir::Expr& op_expr) {
   using cinn::hlir::framework::pir::trivial_fusion_detail::ExprSetFinderUtils::
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.h b/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.h
@@ -60,6 +60,7 @@ struct ApplyTransform {
   ir::Expr expr_;
 };
 
+std::vector<ir::Expr> GetFusibleOpsExpr(std::vector<FusibleOp> fusion_ops);
 std::vector<ir::Expr> TopoSort(const std::vector<ir::Expr>& op_exprs);
 std::vector<FusibleOp> DoPadding(const FusibleOp& fusion_op,
                                  const std::vector<int>& padding_pos);
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/interpreter.cc b/paddle/cinn/operator_fusion/fusion_tracker/interpreter.cc
@@ -33,6 +33,8 @@ void RunCombineInstr(const std::shared_ptr<CombineInstr>& instr,
     const auto& to_insert = interpreter->scope.at(name);
     new_pattern->Extend(to_insert->fusion_ops);
   }
+  VLOG(4) << "After CombineInstr Pattern: \n"
+          << GetFusibleOpsExpr(new_pattern->fusion_ops);
   interpreter->scope[instr->result_] = new_pattern;
 }
 
@@ -69,13 +71,22 @@ void RunTrivialInlineInstr(const std::shared_ptr<TrivialInlineInstr>& instr,
 
 void RunTmpTransformInstr(const std::shared_ptr<TmpTransformInstr>& instr,
                           FusionInterpreter* interpreter) {
-  VLOG(4) << interpreter->scope[instr->upstream_]->fusion_ops.size();
-  PADDLE_ENFORCE_EQ(interpreter->scope[instr->downstream_]->fusion_ops.size(),
-                    1,
-                    ::common::errors::InvalidArgument(
-                        "Downstream op must have only one fusion_op."));
+  PADDLE_ENFORCE_GT(
+      interpreter->scope.count(instr->upstream_),
+      0,
+      ::common::errors::NotFound("Can not find TmpTransformInstr uptream."));
+  PADDLE_ENFORCE_GT(
+      interpreter->scope.count(instr->downstream_),
+      0,
+      ::common::errors::NotFound("Can not find TmpTransformInstr downstream."));
+
+  PADDLE_ENFORCE_EQ(
+      interpreter->scope[instr->downstream_]->fusion_ops.size(),
+      1,
+      ::common::errors::InvalidArgument(
+          "Downstream %s must have only one fusion_op.", instr->downstream_));
   auto upstream_op = std::get<ReduceOp>(
-      interpreter->scope[instr->upstream_]->fusion_ops.front());
+      interpreter->scope[instr->upstream_]->fusion_ops.back());
   auto downstream_op =
       interpreter->scope[instr->downstream_]->fusion_ops.front();
   // inplace set the upstream
diff --git a/paddle/cinn/operator_fusion/pattern.h b/paddle/cinn/operator_fusion/pattern.h
@@ -88,6 +88,7 @@ struct ReduceTreePattern {
                              const FusionTrackerPtr& tracker)
       : childs_(childs), root_(root), tracker_(tracker) {
     id_ = UniqueId();
+    cur_id_ = id_;
   }
   const ReducePattern& GetRootPattern() const { return root_; }
   std::vector<pir::Operation*> ops() const {
@@ -118,44 +119,48 @@ struct ReduceTreePattern {
   std::string id() const { return id_; }
   std::string id_;
 
+  mutable std::string cur_id_;
+  std::string cur_id() const { return cur_id_; }
+  std::string new_tmp_id() const {
+    if (cur_id_ == id_) {
+      cur_id_ = id_ + "_tmp_0";
+    } else {
+      int ith = std::stoi(cur_id_.substr(cur_id_.size() - 1));
+      cur_id_ = id_ + "_tmp_" + std::to_string(ith + 1);
+    }
+    return cur_id_;
+  }
+
   FusionTrackerPtr tracker_;
 
   void update_tracker() const {
-    int counter = 0;
-    std::function<std::string()> gen_name = [&counter]() {
-      return "tmp_" + std::to_string(counter++);
-    };
-    const std::string& root_name = id();
+    const std::string& root_name = GetRootPattern().id();
     std::vector<std::string> names;
-    UpdateTrackerImpl(root_name,
-                      *this,
-                      std::vector<size_t>(),
-                      gen_name,
-                      this->tracker_,
-                      &names);
-    tracker_->append(std::make_shared<CombineInstr>(names, root_name));
+    UpdateTrackerImpl(
+        root_name, *this, std::vector<size_t>(), this->tracker_, &names);
+    tracker_->append(std::make_shared<CombineInstr>(names, cur_id()));
   }
 
   void UpdateTrackerImpl(const std::string root_name,
                          const ReduceTreePattern& root,
                          const std::vector<size_t>& fake_reduce_iter_idx,
-                         const std::function<std::string()>& unique_tmp_name_fn,
                          FusionTrackerPtr tracker,
                          std::vector<std::string>* names) const {
-    // Apply a bunch of tracker to get a output_name of ReduceTreePattern.
+    // Apply a brunch of tracker to get a output_name of ReduceTreePattern.
     // names and trackers collect all the needed fusion nodes.
-    for (const auto& child : childs_) {
-      const std::string& tmp_name = unique_tmp_name_fn();
-      tracker->append(std::make_shared<TmpTransformInstr>(
-          tmp_name, root_name, tmp_name, root_name, fake_reduce_iter_idx));
-      UpdateTrackerImpl(tmp_name,
-                        child,
-                        fake_reduce_iter_idx,
-                        unique_tmp_name_fn,
-                        tracker,
-                        names);
+    for (const auto& child : root.childs()) {
+      auto origin_child_id = child.cur_id();
+      auto new_child_id = child.new_tmp_id();
+      tracker->append(
+          std::make_shared<TmpTransformInstr>(origin_child_id,
+                                              root_name,
+                                              new_child_id,
+                                              root.cur_id(),
+                                              fake_reduce_iter_idx));
+      UpdateTrackerImpl(
+          new_child_id, child, fake_reduce_iter_idx, tracker, names);
     }
-    names->push_back(root_name);
+    names->push_back(root.cur_id());
   }
 
  private:
@@ -190,29 +195,21 @@ struct ReduceTreePlusTrivialPattern {
   FusionTrackerPtr tracker_;
 
   void update_tracker() const {
-    int counter = 0;
-    std::function<std::string()> gen_name = [&counter]() {
-      return "tmp_" + std::to_string(counter++);
-    };
     const std::string& root_name = id();
-    const std::string& tmp_name_for_tree = gen_name();
+    const std::string& origin_tree_id = tree.cur_id();
+    const std::string& new_tree_id = tree.new_tmp_id();
     std::vector<std::string> names;
-    tracker_->append(
-        std::make_shared<TmpTransformInstr>(tree.GetRootPattern().id(),
-                                            sink_trivial.id(),
-                                            tmp_name_for_tree,
-                                            root_name,
-                                            fake_reduce_iter_idx));
-    tree.UpdateTrackerImpl(tmp_name_for_tree,
-                           tree,
-                           fake_reduce_iter_idx,
-                           gen_name,
-                           this->tracker_,
-                           &names);
+    tracker_->append(std::make_shared<TmpTransformInstr>(origin_tree_id,
+                                                         sink_trivial.id(),
+                                                         new_tree_id,
+                                                         root_name,
+                                                         fake_reduce_iter_idx));
+    tree.UpdateTrackerImpl(
+        new_tree_id, tree, fake_reduce_iter_idx, this->tracker_, &names);
     names.push_back(root_name);
     // optimize the loop range of R + T for speed up.
     tracker_->append(std::make_shared<TrivialLoopAlignInstr>(
-        tmp_name_for_tree, root_name, root_name, fake_reduce_iter_idx));
+        new_tree_id, root_name, root_name, fake_reduce_iter_idx));
     // collect all the Expr and represent the root_name.
     tracker_->append(std::make_shared<CombineInstr>(names, root_name));
   }
diff --git a/paddle/cinn/operator_fusion/pattern_fuser.h b/paddle/cinn/operator_fusion/pattern_fuser.h
@@ -132,16 +132,23 @@ static StmtPattern MergePatternImpl(const TrivialPattern& first,
 
 // RR & RT
 
-static int InsertDownstreamIntoTree(const ReduceTreePattern& upstream,
-                                    ReduceTreePattern& downstream) {  // NOLINT
-  if (IsDirectUpstream(upstream.GetRootPattern().GetReduceOp(),
-                       downstream.GetRootPattern().GetReduceOp())) {
+static int InsertUpstreamIntoTree(const ReduceTreePattern& upstream,
+                                  ReduceTreePattern& downstream) {  // NOLINT
+  auto is_direct_upstream = [&](const ReducePattern& upstream,
+                                const ReducePattern& downstream) -> bool {
+    auto upstream_result = upstream.GetReduceOp()->result(0);
+    auto user_ops = FindUserOp(downstream.ops(), upstream_result);
+    return !user_ops.empty();
+  };
+
+  if (is_direct_upstream(upstream.GetRootPattern(),
+                         downstream.GetRootPattern())) {
     downstream.InsertChild(upstream);
     return 1;
   }
   int insert_num = 0;
   for (auto& child : downstream.childs()) {
-    insert_num += InsertDownstreamIntoTree(upstream, child);
+    insert_num += InsertUpstreamIntoTree(upstream, child);
   }
   return insert_num;
 }
@@ -153,7 +160,7 @@ static StmtPattern MergePatternImpl(const ReduceTreePattern& upstream,
       downstream.GetRootPattern(),
       std::make_shared<FusionTracker>(upstream.tracker_,
                                       downstream.tracker_));  // copy first.
-  int insert_num = InsertDownstreamIntoTree(upstream, result);
+  int insert_num = InsertUpstreamIntoTree(upstream, result);
   PADDLE_ENFORCE_EQ(insert_num,
                     1,
                     phi::errors::PreconditionNotMet(
diff --git a/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc b/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
@@ -174,37 +174,17 @@ bool IsProductSmallerOrEqual(const std::vector<DimUsage>& first,
   return shape_analysis.IsEqual(first_product, second_product);
 }
 
-pir::Operation* FindUserOp(const std::vector<pir::Operation*>& candidates,
-                           const pir::Value& value) {
-  std::vector<pir::Operation*> results;
-  for (auto consumer_it = value.use_begin(); consumer_it != value.use_end();
-       ++consumer_it) {
-    pir::Operation* user_op = consumer_it.owner();
-    auto iter = std::find(candidates.begin(), candidates.end(), user_op);
-    if (iter != candidates.end()) {
-      results.emplace_back(*iter);
-    }
-  }
-  PADDLE_ENFORCE_EQ(results.size(),
-                    1,
-                    ::common::errors::InvalidArgument(
-                        "Zero or multiple user operations found in candidates! "
-                        "Expected exactly one, but found %d.",
-                        results.size()));
-  return results.front();
-}
-
 bool RelativeJudgePolicy::ReduceTreeGrownCanMerge(
     const PatternNodePtr& upstream, const PatternNodePtr& downstream) {
   const auto& upstream_tree =
       std::get<ReduceTreePattern>(upstream->stmt_pattern());
   const auto& downstream_tree =
       std::get<ReduceTreePattern>(downstream->stmt_pattern());
 
-  VLOG(4) << "upstream->stmt_pattern():"
-          << OpsDebugStr(GetOpsInPattern(upstream_tree));
-  VLOG(4) << "downstream->stmt_pattern()"
-          << OpsDebugStr(GetOpsInPattern(downstream_tree));
+  VLOG(4) << "upstream: \n" << OpsDebugStr(GetOpsInPattern(upstream_tree));
+  VLOG(4) << "upstream->childs_num: " << upstream_tree.childs().size();
+  VLOG(4) << "downstream: \n" << OpsDebugStr(GetOpsInPattern(downstream_tree));
+  VLOG(4) << "downstream->childs_num: " << downstream_tree.childs().size();
 
   const auto& maybe_downstream_op = GetDownstreamFromCandidate(
       upstream_tree.GetRootPattern(), downstream_tree.FlattenReducePattern());
@@ -220,7 +200,7 @@ bool RelativeJudgePolicy::ReduceTreeGrownCanMerge(
   }
   const pir::Value& reduce_out_value =
       upstream_tree.GetRootPattern().GetReduceOp()->result(0);
-  auto downstream_connect_op =
+  auto downstream_connect_ops =
       FindUserOp(downstream_tree.ops(), reduce_out_value);
   pir::Operation* downstream_reduce_op =
       maybe_downstream_op.value().GetReduceOp();
@@ -229,8 +209,13 @@ bool RelativeJudgePolicy::ReduceTreeGrownCanMerge(
       SplitReduceDims(axes_info_.GetSignature(downstream_reduce_op),
                       downstream_reduce_op);
 
-  const auto& upstream_output_dims = GetValueUsage(
-      reduce_out_value, GetUsageIdx(reduce_out_value, downstream_connect_op));
+  std::vector<DimUsage> upstream_output_dims;
+  for (const auto& op : downstream_connect_ops) {
+    auto dim_usages =
+        GetValueUsage(reduce_out_value, GetUsageIdx(reduce_out_value, op));
+    upstream_output_dims.insert(
+        upstream_output_dims.end(), dim_usages.begin(), dim_usages.end());
+  }
   const auto& [related, _UNUSED] =
       SplitFirstIfRelatedBySecond(downstream_reduce_dims, upstream_output_dims);
   auto res = (related.size() == 0);
diff --git a/paddle/cinn/operator_fusion/utils.h b/paddle/cinn/operator_fusion/utils.h
@@ -453,10 +453,24 @@ static const size_t GetResultIdx(const pir::Value& v, pir::Operation* op) {
       "Can not find the value %s as result of op %s", v.impl(), op->name()));
 }
 
+static std::vector<pir::Operation*> FindUserOp(
+    const std::vector<pir::Operation*>& candidates, const pir::Value& value) {
+  std::vector<pir::Operation*> results;
+  for (auto consumer_it = value.use_begin(); consumer_it != value.use_end();
+       ++consumer_it) {
+    pir::Operation* user_op = consumer_it.owner();
+    auto iter = std::find(candidates.begin(), candidates.end(), user_op);
+    if (iter != candidates.end()) {
+      results.emplace_back(*iter);
+    }
+  }
+  return results;
+}
+
 static bool IsDirectUpstream(const pir::Operation* upstream,
                              const pir::Operation* downstream) {
-  for (const auto& value : downstream->results()) {
-    for (const auto& operand : upstream->operands()) {
+  for (const auto& value : upstream->results()) {
+    for (const auto& operand : downstream->operands()) {
       if (value == operand.source()) {
         return true;
       }

Original file line number	Diff line number	Diff line change
`@@ -181,7 +181,6 @@ std::vector<ir::Var> GetOutputIters(const FusibleOp& op) {`
`181`	`181`	`}`
`182`	`182`	`};`
`183`	`183`	`VLOG(4) << "GetOutputIters";`
`184`		`- VLOG(4) << "Before AppendBound:" << _GetRootExpr(op);`
`185`	`184`	`return AppendBound(std::visit(Visitor(), op), _GetRootExpr(op));`
`186`	`185`	`}`
`187`	`186`