From 20083679a0d5d4f349c6496e64296e844d170d61 Mon Sep 17 00:00:00 2001
From: limingqi107 <limingqi@huawei.com>
Date: Tue, 16 Jun 2020 15:43:22 +0800
Subject: [PATCH] gpu memreuse supports summary node

---
 mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc    | 2 ++
 mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc | 4 ++++
 mindspore/ccsrc/session/gpu_session.cc              | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc b/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc
index 0e90b90f2d..da96ce8a54 100644
--- a/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc
+++ b/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc
@@ -190,6 +190,8 @@ void GPUKernelRuntime::InitKernelRefCount(const session::KernelGraph *graph) {
   mem_reuse_util_ptr->SetReuseRefCount();
   // Can't free the device address of graph output, so set the reference count of graph output specially.
   mem_reuse_util_ptr->SetGraphOutputRefCount();
+  // Can't free the device address of summary nodes, so set the reference count of summary nodes specially.
+  mem_reuse_util_ptr->SetSummaryNodesRefCount();
   auto graph_id = graph->graph_id();
   mem_reuse_util_map_[graph_id] = mem_reuse_util_ptr;
 }
diff --git a/mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc b/mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc
index a08c228220..1ee2254c01 100644
--- a/mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc
+++ b/mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc
@@ -323,6 +323,10 @@ void MemReuseUtil::SetSummaryNodesRefCount() {
       MS_LOG(WARNING) << "can't find summary node's kernel_def " << node->fullname_with_scope();
     }
   }
+#ifdef MEM_REUSE_DEBUG
+  auto graph = *graph_;
+  MemReuseChecker::GetInstance().CheckMemReuseIR(total_refs_list_, kernel_def_ptr_list_, &graph);
+#endif
 }
 
 void MemReuseUtil::SetGraphOutputRefCount() {
diff --git a/mindspore/ccsrc/session/gpu_session.cc b/mindspore/ccsrc/session/gpu_session.cc
index 999c5ba163..9597717caa 100644
--- a/mindspore/ccsrc/session/gpu_session.cc
+++ b/mindspore/ccsrc/session/gpu_session.cc
@@ -162,6 +162,8 @@ GraphId GPUSession::CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList
   auto execution_order = graph->execution_order();
   Reorder(&execution_order);
   graph->set_execution_order(execution_order);
+  // Get summary nodes.
+  GetSummaryNodes(graph.get());
   // Remove NoOp from execution graph
   opt::RemoveNopNode(graph.get());
   // Alloc memory, including static memory and dynamic memory