From f0db47aba82e91510f251fda9f0fbcfa68a5e475 Mon Sep 17 00:00:00 2001
From: wenchunjiang <wenchunjiang@huawei.com>
Date: Thu, 17 Sep 2020 17:32:36 +0800
Subject: [PATCH] remove ValueNode<KernelGraphPtr> after replace call to
 lagelgoto in order to fix mutual reference in while

---
 .../ccsrc/backend/session/ascend_control_parser.cc  |  1 +
 mindspore/ccsrc/backend/session/kernel_graph.cc     | 13 +++++++++++++
 mindspore/ccsrc/backend/session/kernel_graph.h      |  1 +
 3 files changed, 15 insertions(+)

diff --git a/mindspore/ccsrc/backend/session/ascend_control_parser.cc b/mindspore/ccsrc/backend/session/ascend_control_parser.cc
index 0bd79b54b5..f4d91f5858 100644
--- a/mindspore/ccsrc/backend/session/ascend_control_parser.cc
+++ b/mindspore/ccsrc/backend/session/ascend_control_parser.cc
@@ -565,6 +565,7 @@ void AscendControlParser::RecurseCall(NotNull<KernelGraphPtr> kg, NotNull<CNodeP
   cur_node->set_inputs(new_inputs);
   cur_node->set_abstract(nullptr);
   AnfAlgo::SetNodeAttr(kAttrChildGraph, MakeValue<std::vector<KernelGraphPtr>>({call_kg}), cur_node.get());
+  kg->RemoveNodeFromGraph(origin_inputs[kCNodeCallArg]);
   origin_inputs.assign(origin_inputs.begin() + kCNodeCallArg + 1, origin_inputs.end());
   AttachOriginalInputsToGraph(kg, origin_inputs);
   MS_LOG(INFO) << "Succeed processing call func " << cur_node->DebugString();
diff --git a/mindspore/ccsrc/backend/session/kernel_graph.cc b/mindspore/ccsrc/backend/session/kernel_graph.cc
index c2bcc58e39..c0c0e0a6ff 100644
--- a/mindspore/ccsrc/backend/session/kernel_graph.cc
+++ b/mindspore/ccsrc/backend/session/kernel_graph.cc
@@ -1188,6 +1188,19 @@ void KernelGraph::UpdateChildGraphOrder() {
   child_graph_order_ = child_graph_order;
 }
 
+void KernelGraph::RemoveNodeFromGraph(const AnfNodePtr &node) {
+  if (backend_front_anf_map_.find(node) != backend_front_anf_map_.end()) {
+    auto front_node = backend_front_anf_map_[node];
+    (void)backend_front_anf_map_.erase(node);
+    (void)front_backend_anf_map_.erase(front_node);
+  }
+  if (node->isa<ValueNode>()) {
+    if (graph_value_nodes_.find(node->cast<ValueNodePtr>()) != graph_value_nodes_.end()) {
+      (void)graph_value_nodes_.erase(node->cast<ValueNodePtr>());
+    }
+  }
+}
+
 std::string KernelGraph::ToString() const { return std::string("kernel_graph_").append(std::to_string(graph_id_)); }
 
 KernelGraph::~KernelGraph() {
diff --git a/mindspore/ccsrc/backend/session/kernel_graph.h b/mindspore/ccsrc/backend/session/kernel_graph.h
index 152ae498cd..ff01501bbc 100644
--- a/mindspore/ccsrc/backend/session/kernel_graph.h
+++ b/mindspore/ccsrc/backend/session/kernel_graph.h
@@ -174,6 +174,7 @@ class KernelGraph : public FuncGraph {
       return nullptr;
     }
   }
+  void RemoveNodeFromGraph(const AnfNodePtr &node);
 
  private:
   // remove value node form graph