1. replace weak_ptr to shared_ptr in child_graph_order_ an parent_graph_

2. erase kAttrChildGraph after finish using

mindspore/ccsrc/backend/session/ascend_control_parser.cc

@@ -63,7 +63,7 @@ static void RecursiveReplaceNode(NotNull<KernelGraphPtr> kg, NotNull<AnfNodePtr>
   }
 
   for (auto &child : kg->child_graph_order()) {
-    RecursiveReplaceNode(NOT_NULL(child), main_parameter, parameter_reuse_set, memo);
+    RecursiveReplaceNode(NOT_NULL(child.lock()), main_parameter, parameter_reuse_set, memo);
   }
 }
 
@@ -177,13 +177,14 @@ void AscendControlParser::AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> g
     return;
   }
   memo->insert(graph.get());
-  const std::vector<std::shared_ptr<KernelGraph>> &child_graph_order = graph->child_graph_order();
+  const std::vector<std::weak_ptr<KernelGraph>> &child_graph_order = graph->child_graph_order();
   if (child_graph_order.empty()) {
     return;
   }
 
   std::vector<AnfNodePtr> depend_inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimPartial->name()))};
-  for (auto &cg : child_graph_order) {
+  for (auto &kg : child_graph_order) {
+    std::shared_ptr<KernelGraph> cg = kg.lock();
     MS_EXCEPTION_IF_NULL(cg);
     auto fg = cg->cast<FuncGraphPtr>();
     MS_EXCEPTION_IF_NULL(fg);
@@ -207,7 +208,7 @@ void AscendControlParser::LinkGraph(NotNull<KernelGraphPtr> kg) {
   memo.clear();
   // assign label resource
   device::ascend::AscendLabelAssign::GetInstance().AssignLabel(kg);
-  AttachChildGraphToReturnNode(kg, NOT_NULL(&memo));
+  // AttachChildGraphToReturnNode(kg, NOT_NULL(&memo));
 }
 
 void AscendControlParser::EraseParameter(NotNull<KernelGraphPtr> root_graph,
@@ -428,7 +429,7 @@ void AscendControlParser::ChildGraphDataAssign(
   }
   kg->SetExecOrderByDefault();
   for (auto &child_graph : kg->child_graph_order()) {
-    ChildGraphDataAssign(NOT_NULL(child_graph), link_list, memo);
+    ChildGraphDataAssign(NOT_NULL(child_graph.lock()), link_list, memo);
   }
 }
 
@@ -772,7 +773,7 @@ std::vector<CNodePtr> AscendControlParser::RecurseGraph(NotNull<KernelGraphPtr>
       MS_LOG(EXCEPTION) << "Index out of range:" << graph->child_graph_order().size();
     }
     auto child_graph = graph->child_graph_order()[child_order_index++];
-    auto child_execution_order = RecurseGraph(NOT_NULL(child_graph), memo);
+    auto child_execution_order = RecurseGraph(NOT_NULL(child_graph.lock()), memo);
     execution_order.insert(execution_order.end(), child_execution_order.begin(), child_execution_order.end());
   };
 
@@ -791,6 +792,10 @@ std::vector<CNodePtr> AscendControlParser::RecurseGraph(NotNull<KernelGraphPtr>
       uint32_t label_index = AnfAlgo::GetNodeAttr<uint32_t>(node, kAttrLabelIndex);
       recurse_child_graph(child_graph_index, label_index, node);
     }
+    // erase kAttrChildGraph after finish using
+    if (AnfAlgo::HasNodeAttr(kAttrChildGraph, node)) {
+      AnfAlgo::EraseNodeAttr(kAttrChildGraph, node);
+    }
   }
   graph->set_execution_order(execution_order);
   graph->PrintGraphExecuteOrder();

mindspore/ccsrc/backend/session/ascend_session.cc

@@ -860,7 +860,7 @@ void AscendSession::CreateMultiBranchOutput(NotNull<KernelGraphPtr> graph, NotNu
   memo->insert(graph.get());
   graph->UpdateChildGraphOrder();
   for (auto &child_graph : graph->child_graph_order()) {
-    CreateMultiBranchOutput(NOT_NULL(child_graph), memo);
+    CreateMultiBranchOutput(NOT_NULL(child_graph.lock()), memo);
   }
   std::map<AnfNodePtr, AnfNodePtr> need_replace_list;
   auto node_list = GetCNodes(TopoSort(graph->get_return()));
@@ -932,7 +932,7 @@ void AscendSession::IrFusionPass(const NotNull<KernelGraphPtr> graph, NotNull<st
   }
 
   for (auto &child_graph : graph->child_graph_order()) {
-    IrFusionPass(NOT_NULL(child_graph), memo);
+    IrFusionPass(NOT_NULL(child_graph.lock()), memo);
   }
 }
 
@@ -1016,7 +1016,7 @@ void AscendSession::HardwareOptimize(NotNull<KernelGraphPtr> graph,
 
   HardwareOptimize(graph.get());
   for (auto &child_graph : graph->child_graph_order()) {
-    HardwareOptimize(NOT_NULL(child_graph), memo);
+    HardwareOptimize(NOT_NULL(child_graph.lock()), memo);
   }
   MS_LOG(INFO) << "Finish doing HardwareOptimize in graph: " << graph->graph_id();
 }
@@ -1035,7 +1035,7 @@ void AscendSession::AssignStaticMemory(NotNull<KernelGraphPtr> graph,
   runtime_instance->AssignStaticMemoryInput(graph.get().get());
   runtime_instance->AssignStaticMemoryValueNode(graph.get().get());
   for (auto &child_graph : graph->child_graph_order()) {
-    AssignStaticMemory(NOT_NULL(child_graph), memo);
+    AssignStaticMemory(NOT_NULL(child_graph.lock()), memo);
   }
   MS_LOG(INFO) << "Finish assigning static memory for parameter in graph: " << graph->graph_id();
 }
@@ -1048,9 +1048,11 @@ void AscendSession::UpdateRefOutputMap(NotNull<KernelGraphPtr> graph,
   memo->insert(graph.get());
 
   for (auto &child_graph : graph->child_graph_order()) {
-    UpdateRefOutputMap(NOT_NULL(child_graph), memo);
+    std::shared_ptr<KernelGraph> child_graph_ptr = child_graph.lock();
+    MS_EXCEPTION_IF_NULL(child_graph_ptr);
+    UpdateRefOutputMap(NOT_NULL(child_graph_ptr), memo);
     // copy ref map to final graph
-    auto child_ref_map = child_graph->GetRefMap();
+    auto child_ref_map = child_graph_ptr->GetRefMap();
     for (auto &item : child_ref_map) {
       if (graph->IsInRefOutputMap(item.first)) {
         MS_LOG(WARNING) << "The ref pair <" << item.first.first->DebugString() << ", " << item.first.second

mindspore/ccsrc/backend/session/kernel_graph.cc

@@ -922,8 +922,9 @@ std::vector<std::shared_ptr<KernelGraph>> KernelGraph::GetLeafGraphOrder() {
     leaf_graph_order.push_back(shared_from_this()->cast<KernelGraphPtr>());
   } else {
     for (const auto &child_graph : child_graph_order_) {
-      MS_EXCEPTION_IF_NULL(child_graph);
-      auto child_leaf_graph_order = child_graph->GetLeafGraphOrder();
+      std::shared_ptr<KernelGraph> child_graph_ptr = child_graph.lock();
+      MS_EXCEPTION_IF_NULL(child_graph_ptr);
+      auto child_leaf_graph_order = child_graph_ptr->GetLeafGraphOrder();
       std::copy(child_leaf_graph_order.begin(), child_leaf_graph_order.end(), std::back_inserter(leaf_graph_order));
     }
   }
@@ -1103,13 +1104,13 @@ void KernelGraph::UpdateChildGraphOrder() {
   SetExecOrderByDefault();
   auto call_nodes = FindNodeByPrimitive(
     {std::make_shared<Primitive>(prim::kPrimCall->name()), std::make_shared<Primitive>(prim::kPrimSwitch->name())});
-  std::vector<KernelGraphPtr> child_graph_order;
+  std::vector<std::weak_ptr<KernelGraph>> child_graph_order;
   for (auto &call_node : call_nodes) {
     MS_EXCEPTION_IF_NULL(call_node);
     auto call_child_graphs = AnfAlgo::GetCallSwitchKernelGraph(call_node->cast<CNodePtr>());
     for (const auto &child_graph : call_child_graphs) {
       MS_EXCEPTION_IF_NULL(child_graph);
-      if (child_graph != parent_graph_) {
+      if (child_graph != parent_graph_.lock()) {
         auto shared_this = std::dynamic_pointer_cast<KernelGraph>(shared_from_this());
         MS_EXCEPTION_IF_NULL(shared_this);
         child_graph->set_parent_graph(shared_this);
@@ -1118,7 +1119,9 @@ void KernelGraph::UpdateChildGraphOrder() {
     }
   }
   for (size_t i = 0; i < child_graph_order.size(); ++i) {
-    MS_LOG(INFO) << "Child graph[" << i << "][id:" << child_graph_order[i]->graph_id() << "]";
+    std::shared_ptr<KernelGraph> child_graph = child_graph_order[i].lock();
+    MS_EXCEPTION_IF_NULL(child_graph);
+    MS_LOG(INFO) << "Child graph[" << i << "][id:" << child_graph->graph_id() << "]";
   }
   child_graph_order_ = child_graph_order;
 }

mindspore/ccsrc/backend/session/kernel_graph.h

@@ -114,8 +114,8 @@ class KernelGraph : public FuncGraph {
   // calculate the leaf graph order of root graph
   std::vector<std::shared_ptr<KernelGraph>> GetLeafGraphOrder();
   // the child graph of current graph
-  const std::vector<std::shared_ptr<KernelGraph>> &child_graph_order() const { return child_graph_order_; }
-  void set_child_graph_order(const std::vector<std::shared_ptr<KernelGraph>> &order) { child_graph_order_ = order; }
+  const std::vector<std::weak_ptr<KernelGraph>> &child_graph_order() const { return child_graph_order_; }
+  void set_child_graph_order(const std::vector<std::weak_ptr<KernelGraph>> &order) { child_graph_order_ = order; }
   // checkout whether current graph is leaf graph
   bool IsLeafGraph() const;
 
@@ -126,9 +126,9 @@ class KernelGraph : public FuncGraph {
   // get input_tensors pointer of control parameter
   std::shared_ptr<std::vector<tensor::TensorPtr>> input_ctrl_tensors() const { return input_ctrl_tensors_; }
   // get parent kernel graph
-  std::shared_ptr<KernelGraph> parent_graph() const { return parent_graph_; }
+  std::weak_ptr<KernelGraph> parent_graph() const { return parent_graph_; }
   // set parent kernel graph
-  void set_parent_graph(const std::shared_ptr<KernelGraph> &parent_graph) { parent_graph_ = parent_graph; }
+  void set_parent_graph(const std::weak_ptr<KernelGraph> &parent_graph) { parent_graph_ = parent_graph; }
   // find anf node in graph
   std::vector<CNodePtr> FindNodeByPrimitive(const PrimitivePtr &primitive) const;
   std::vector<CNodePtr> FindNodeByPrimitive(const std::vector<PrimitivePtr> &primitive_list) const;
@@ -227,13 +227,13 @@ class KernelGraph : public FuncGraph {
   std::vector<bool> valid_inputs_;
 
   // child graph execute order in root graph
-  std::vector<std::shared_ptr<KernelGraph>> child_graph_order_;
+  std::vector<std::weak_ptr<KernelGraph>> child_graph_order_;
 
   // input_tensors of control parameter
   std::shared_ptr<std::vector<tensor::TensorPtr>> input_ctrl_tensors_;
 
   // parameter graph
-  std::shared_ptr<KernelGraph> parent_graph_;
+  std::weak_ptr<KernelGraph> parent_graph_;
 
   CNodePtr start_label_;
   CNodePtr end_goto_;

mindspore/ccsrc/runtime/device/ascend/ascend_label_assign.cc

@@ -89,7 +89,7 @@ static void AssignLabelForLabelSet(NotNull<std::shared_ptr<session::KernelGraph>
   }
 
   for (auto &cg : graph->child_graph_order()) {
-    AssignLabelForLabelSet(NOT_NULL(cg), label_id, memo);
+    AssignLabelForLabelSet(NOT_NULL(cg.lock()), label_id, memo);
   }
 }
 
@@ -120,7 +120,7 @@ static void AssignLabelForGotoSwitch(NotNull<std::shared_ptr<session::KernelGrap
     }
   }
   for (auto &cg : graph->child_graph_order()) {
-    AssignLabelForGotoSwitch(NOT_NULL(cg), memo);
+    AssignLabelForGotoSwitch(NOT_NULL(cg.lock()), memo);
   }
 }