From dd3191026f8b6dafffed2ba5f0a95ed2faf6d7d5 Mon Sep 17 00:00:00 2001
From: hwjiaorui <jiaorui3@huawei.com>
Date: Wed, 4 Aug 2021 15:21:32 +0800
Subject: [PATCH] fix stream find target

---
 .../device/ascend/ascend_stream_assign.cc     | 36 ++++++++++++-------
 .../device/ascend/ascend_stream_assign.h      |  3 +-
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc
index e1a773864c..a2850cdc33 100644
--- a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc
+++ b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.cc
@@ -1992,6 +1992,28 @@ CNodePtr AscendStreamAssign::CreateRecvApplyKernel(const NotNull<KernelGraphPtr>
   return recv_node_ptr;
 }
 
+bool AscendStreamAssign::IsNopNodeTarget(const AnfNodePtr &nop_node, const CNodePtr &target_node,
+                                         const CNodePtr &cur_node, bool exclude_hcom) {
+  MS_EXCEPTION_IF_NULL(nop_node);
+  auto cnode = nop_node->cast<CNodePtr>();
+  auto new_inputs = cnode->inputs();
+  for (size_t i = 1; i < new_inputs.size(); i++) {
+    if (opt::IsNopNode(new_inputs[i])) {
+      if (IsNopNodeTarget(new_inputs[i], target_node, cur_node, exclude_hcom)) {
+        return true;
+      }
+    } else {
+      auto new_real_input = AnfAlgo::VisitKernel(new_inputs[i], 0);
+      if (target_node == new_real_input.first) {
+        if (!(exclude_hcom && IsHcom(cur_node))) {
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+
 vector<CNodePtr>::iterator AscendStreamAssign::FindTargetOp(vector<CNodePtr>::iterator begin,
                                                             vector<CNodePtr>::iterator end, const CNodePtr &node,
                                                             bool exclude_hcom) {
@@ -2000,18 +2022,8 @@ vector<CNodePtr>::iterator AscendStreamAssign::FindTargetOp(vector<CNodePtr>::it
     for (size_t i = 1; i < inputs.size(); i++) {
       auto input = inputs[i];
       if (opt::IsNopNode(input)) {
-        CNodePtr cnode = input->cast<CNodePtr>();
-        auto new_inputs = cnode->inputs();
-        for (size_t j = 1; j < new_inputs.size(); j++) {
-          auto new_real_input = AnfAlgo::VisitKernel(new_inputs[j], 0);
-          // find target node except hcom op. insert event for hcom in:InsertEventHcomDependCommonBak function
-          // only insert one time
-          if (node == new_real_input.first) {
-            if (!(exclude_hcom && IsHcom(*begin))) {
-              MS_LOG(DEBUG) << "Nop node find target op[" << (*begin)->DebugString() << "]";
-              return begin;
-            }
-          }
+        if (IsNopNodeTarget(input, node, *begin, exclude_hcom)) {
+          return begin;
         }
       } else {
         auto real_input = AnfAlgo::VisitKernel(input, 0);
diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.h b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.h
index 8f7773e77b..bfe55a440d 100644
--- a/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.h
+++ b/mindspore/ccsrc/runtime/device/ascend/ascend_stream_assign.h
@@ -175,7 +175,8 @@ class AscendStreamAssign {
   uint32_t GetIndexByKey(const NotNull<KernelGraphPtr> &graph_ptr, const CNodeKey &key);
   uint32_t GetIndependentStreamSwitchStreamId(const NotNull<KernelGraphPtr> &graph_ptr);
   void GetIndependentMaxTarget(const NotNull<KernelGraphPtr> &graph_ptr);
-
+  bool IsNopNodeTarget(const AnfNodePtr &nop_node, const CNodePtr &target_node, const CNodePtr &cur_node,
+                       bool exclude_hcom);
   bool IsTaskSink();
   bool IsHcom(const CNodePtr &cur_cnode_ptr);
   bool IsIndependentNode(const CNodePtr &node_ptr);