From a410300f2fe89f9e2d55c8e5082403637dd1b8c3 Mon Sep 17 00:00:00 2001
From: wangxiaotian22 <root@DESKTOP-3JISBIL.localdomain>
Date: Mon, 22 Feb 2021 14:46:47 +0800
Subject: [PATCH] add control edge between iter active node and fpbp assign_add
 node

---
 ge/graph/passes/flow_ctrl_pass.cc | 23 ++++++++++++++++++-----
 ge/graph/passes/flow_ctrl_pass.h  |  3 +++
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/ge/graph/passes/flow_ctrl_pass.cc b/ge/graph/passes/flow_ctrl_pass.cc
index 5a294aa2..9d441ed5 100755
--- a/ge/graph/passes/flow_ctrl_pass.cc
+++ b/ge/graph/passes/flow_ctrl_pass.cc
@@ -80,6 +80,16 @@ Status FlowCtrlPass::Run(ComputeGraphPtr compute_graph) {
       graph_change = true;
     }
   }
+
+  // add edge operation below depends on memcpy node in itertor loop set single stream,or may cause block
+  for (auto &active_node : active_nodes_in_iter_loop_) {
+    auto ret = GraphUtils::AddEdge(active_node->GetOutControlAnchor(),
+                                   assign_add_node_in_fpbp_loop_->GetInControlAnchor());
+    if (ret != GRAPH_SUCCESS) {
+      GELOGW("add control edge between iter_loop_node:%s and fpbp_loop_node:%s fail, may cause block",
+              active_node->GetName().c_str(), assign_add_node_in_fpbp_loop_->GetName().c_str());
+    }
+  }
   GELOGI("FlowCtrl pass end, graph is %s.", graph_change ? "changed" : "not changed");
   return graph_change ? SUCCESS : NOT_CHANGED;
 }
@@ -279,16 +289,16 @@ Status FlowCtrlPass::CreateIterCtrlTrueBranch(ComputeGraphPtr &compute_graph, co
    *         loopIncrement
    */
   // Insert AssignAdd node
-  NodePtr assign_add_node =
+  assign_add_node_in_fpbp_loop_ =
       InsertAssignOp(compute_graph, ASSIGNADD, NODE_NAME_FLOWCTRL_LOOP_ASSIGNADD, loop_cond_node, loop_inc_node);
-  if (assign_add_node == nullptr || switch_node == nullptr) {
+  if (assign_add_node_in_fpbp_loop_ == nullptr || switch_node == nullptr) {
     GELOGE(PARAM_INVALID, "assign add node or switch node is null");
     return FAILED;
   }
 
   string active_name = switch_node->GetName() + "_StreamActive";
   // add attr for stream assign model to break branch.
-  GE_CHK_STATUS_RET(SetStreamLabel(assign_add_node, active_name), "set stream label failed");
+  GE_CHK_STATUS_RET(SetStreamLabel(assign_add_node_in_fpbp_loop_, active_name), "set stream label failed");
 
   // used for stream assign to find true branch
   GE_CHK_STATUS_RET(SetActiveLabelList(switch_node, { active_name }), "set active label list failed");
@@ -304,13 +314,15 @@ Status FlowCtrlPass::CreateIterCtrlTrueBranch(ComputeGraphPtr &compute_graph, co
                     DOMI_LOGE("set ATTR_NAME_IS_LOOP_ACTIVE failed"); return FAILED);
 
   // add ctrl edges
-  graphStatus add_ret = GraphUtils::AddEdge(switch_node->GetOutControlAnchor(), assign_add_node->GetInControlAnchor());
+  graphStatus add_ret = GraphUtils::AddEdge(switch_node->GetOutControlAnchor(),
+                                            assign_add_node_in_fpbp_loop_->GetInControlAnchor());
   if (add_ret != GRAPH_SUCCESS) {
     GELOGE(FAILED, "Add switch_node to assign_add_node ctrl edge failed, add_ret=%u.", add_ret);
     return FAILED;
   }
 
-  add_ret = GraphUtils::AddEdge(assign_add_node->GetOutControlAnchor(), active_node->GetInControlAnchor());
+  add_ret = GraphUtils::AddEdge(assign_add_node_in_fpbp_loop_->GetOutControlAnchor(),
+                                active_node->GetInControlAnchor());
   if (add_ret != GRAPH_SUCCESS) {
     GELOGE(FAILED, "Add assign_add_node to active_node ctrl edge failed, add_ret=%u.", add_ret);
     return FAILED;
@@ -533,6 +545,7 @@ Status FlowCtrlPass::AddSpecialNodeIteratorCtrl(ComputeGraphPtr &compute_graph,
   GE_CHK_STATUS_RET(SetActiveLabelList(switch_node, { active_name }), "set active label list failed");
   // used for stream assign to find active stream
   GE_CHK_STATUS_RET(SetActiveLabelList(active_node, { loop_pre_node->GetName() }), "set active label list failed");
+  active_nodes_in_iter_loop_.push_back(active_node);
   return SUCCESS;
 }
 }  // namespace ge
diff --git a/ge/graph/passes/flow_ctrl_pass.h b/ge/graph/passes/flow_ctrl_pass.h
index 35270946..74f3cce0 100755
--- a/ge/graph/passes/flow_ctrl_pass.h
+++ b/ge/graph/passes/flow_ctrl_pass.h
@@ -142,6 +142,9 @@ class FlowCtrlPass : public GraphPass {
   ///         false: only one dataSet exist
   ///
   bool CheckMultiDataSet(ComputeGraphPtr &compute_graph);
+
+  NodePtr assign_add_node_in_fpbp_loop_ = nullptr;
+  std::vector<NodePtr> active_nodes_in_iter_loop_;
 };
 }  // namespace ge