Browse Source

!2099 fix summary nodes memory reuse refcount

Merge pull request !2099 from laiyongqiang/r0.3_summary
tags/v0.3.1-alpha
mindspore-ci-bot Gitee 6 years ago
parent
commit
ab39708929
2 changed files with 27 additions and 0 deletions
  1. +26
    -0
      mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc
  2. +1
    -0
      mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.h

+ 26
- 0
mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.cc View File

@@ -267,6 +267,31 @@ void MemReuseUtil::SetReuseRefCount() {
}
}

void MemReuseUtil::SetSummaryNodesRefCount() {
bool summary_exist = graph_->summary_node_exist();
if (!summary_exist) {
return;
}

auto summary_nodes = graph_->summary_nodes();
if (summary_nodes.empty()) {
return;
}

for (auto &node_item : summary_nodes) {
auto node = node_item.second.first;
size_t index = IntToSize(node_item.second.second);
MS_LOG(INFO) << "set summary node's ref count, node: " << node->fullname_with_scope() << " index: " << index;
if (kernel_output_refs_.find(node.get()) != kernel_output_refs_.end()) {
KernelRefCountPtr kernel_ref = kernel_output_refs_[node.get()][index];
kernel_ref->ref_count_ = kMaxRefCount;
kernel_ref->ref_count_dynamic_use_ = kMaxRefCount;
} else {
MS_LOG(WARNING) << "can't find summary node's kernel_def " << node->fullname_with_scope();
}
}
}

void MemReuseUtil::SetGraphOutputRefCount() {
auto nodes = AnfAlgo::GetAllOutput(graph_->output(), {prim::kPrimTupleGetItem});
for (const auto &node : nodes) {
@@ -305,6 +330,7 @@ void MemReuseUtil::SetAllInfo(KernelGraph *graph) {
}
SetKernelDefMap();
SetReuseRefCount();
SetSummaryNodesRefCount();
SetWorkSpaceList();
#ifdef MEM_REUSE_DEBUG
MemReuseChecker::GetInstance().CheckMemReuseIR(total_refs_list_, kernel_def_ptr_list_, graph);


+ 1
- 0
mindspore/ccsrc/pre_activate/mem_reuse/mem_reuse.h View File

@@ -62,6 +62,7 @@ class MemReuseUtil {
void SetOutputMap(const CNodePtr &kernel, KernelDef *kernel_def_ptr);
void SetWkMap(const CNodePtr &kernel, KernelDef *kernel_def_ptr);
void SetReuseRefCount();
void SetSummaryNodesRefCount();
// Set the reference count of graph output specially.
void SetGraphOutputRefCount();
// Reset the dynamic used reference count by ref_count_.


Loading…
Cancel
Save