From 23cc01f21d23b5a5fc885de9d4bb8e2c9685078f Mon Sep 17 00:00:00 2001
From: zhoufeng <zhoufeng54@huawei.com>
Date: Fri, 17 Jul 2020 16:09:02 +0800
Subject: [PATCH] dumpir can dump subgraphs of ascend kernel graph

Signed-off-by: zhoufeng <zhoufeng54@huawei.com>
---
 .../backend/session/ascend_control_parser.cc  | 29 +++++++++++++++++--
 .../backend/session/ascend_control_parser.h   |  3 +-
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/mindspore/ccsrc/backend/session/ascend_control_parser.cc b/mindspore/ccsrc/backend/session/ascend_control_parser.cc
index 4c6c7ab9cf4..4306a3d7855 100644
--- a/mindspore/ccsrc/backend/session/ascend_control_parser.cc
+++ b/mindspore/ccsrc/backend/session/ascend_control_parser.cc
@@ -107,7 +107,7 @@ static void ReuseParameter(NotNull<KernelGraphPtr> root_kg,
 
 static CNodePtr GetNextRealKernel(const std::vector<CNodePtr> &list, size_t start) {
   for (size_t i = start; i < list.size() - 1; ++i) {
-    if (!IsPrimitiveCNode(list[i], prim::kPrimPartial) && AnfAlgo::IsRealKernel(list[i])) {
+    if (AnfAlgo::IsRealKernel(list[i])) {
       return list[i];
     }
   }
@@ -171,18 +171,43 @@ static void EraseNodeFromExecOrder(const AnfNodePtr &node, const NotNull<std::ve
   exec_order->erase(exec_iter);
 }
 
+void AscendControlParser::AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> graph,
+                                                       const NotNull<std::set<KernelGraphPtr> *> memo) {
+  if (memo->find(graph) != memo->end()) {
+    return;
+  }
+  memo->insert(graph.get());
+  const std::vector<std::shared_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) {
+    MS_EXCEPTION_IF_NULL(cg);
+    auto fg = cg->cast<FuncGraphPtr>();
+    MS_EXCEPTION_IF_NULL(fg);
+    depend_inputs.emplace_back(NewValueNode(fg));
+    AttachChildGraphToReturnNode(NOT_NULL(cg), memo);
+  }
+  auto child_graphs = graph->NewCNode(depend_inputs);
+  InsertDependToGraph(graph, NOT_NULL(child_graphs));
+}
+
 void AscendControlParser::LinkGraph(NotNull<KernelGraphPtr> kg) {
   std::set<KernelGraphPtr> memo;
   std::vector<std::pair<AnfNodePtr, AnfNodePtr>> link_list;
   // Insert Assign
   ChildGraphDataAssign(kg, NOT_NULL(&link_list), NOT_NULL(&memo));
+  memo.clear();
   // Reuse Parameter
   ReuseParameter(kg, link_list);
   // replace call by label goto / label switch
-  memo.clear();
   (void)ProcessKernelGraph(kg, nullptr, nullptr, NOT_NULL(&memo));
+  memo.clear();
   // assign label resource
   device::ascend::AscendLabelAssign::GetInstance().AssignLabel(kg);
+  AttachChildGraphToReturnNode(kg, NOT_NULL(&memo));
 }
 
 void AscendControlParser::EraseParameter(NotNull<KernelGraphPtr> root_graph,
diff --git a/mindspore/ccsrc/backend/session/ascend_control_parser.h b/mindspore/ccsrc/backend/session/ascend_control_parser.h
index ac247351390..4e62629d985 100644
--- a/mindspore/ccsrc/backend/session/ascend_control_parser.h
+++ b/mindspore/ccsrc/backend/session/ascend_control_parser.h
@@ -66,7 +66,8 @@ class AscendControlParser {
   static AnfNodePtr InsertAssignToGraph(NotNull<KernelGraphPtr> kg, NotNull<AnfNodePtr> from, NotNull<AnfNodePtr> to);
   static std::vector<std::pair<KernelGraphPtr, std::vector<AnfNodePtr>>> ParseCallNode(NotNull<CNodePtr> call_node);
   static std::tuple<KernelGraphPtr, std::vector<AnfNodePtr>> ParsePartial(NotNull<AnfNodePtr> node);
-
+  static void AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> graph,
+                                           const NotNull<std::set<KernelGraphPtr> *> memo);
   // root graph order
   static bool CheckLabelIndex(uint32_t order_index, uint32_t label_index, const CNodePtr &cnode,
                               NotNull<KernelGraphPtr> graph);