From dfe90257c299cf1d8e36bd26888e124c68d946b4 Mon Sep 17 00:00:00 2001 From: zhaosida Date: Thu, 1 Apr 2021 11:00:14 +0800 Subject: [PATCH] enable reuse for non-fusion AtomicAddrClean node --- .../ccsrc/backend/optimizer/somas/somas.cc | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/mindspore/ccsrc/backend/optimizer/somas/somas.cc b/mindspore/ccsrc/backend/optimizer/somas/somas.cc index 64c9c62ed4..9e84fa5aae 100644 --- a/mindspore/ccsrc/backend/optimizer/somas/somas.cc +++ b/mindspore/ccsrc/backend/optimizer/somas/somas.cc @@ -483,12 +483,13 @@ void Somas::InitSomasOutputAndWorkspaceTensors(const session::KernelGraph *graph void Somas::InitSomasInputTensors(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(graph); bool is_all_nop_node = opt::IsAllNopNode(graph); + static const auto enable_fusion_clear = (common::GetEnv("ENV_FUSION_CLEAR") == "1"); auto kernel_cnodes = graph->execution_order(); for (const auto &kernel : kernel_cnodes) { if (AnfAlgo::GetCNodeName(kernel) != kAtomicAddrCleanOpName) { InitCommonNodeInputs(is_all_nop_node, kernel); } else { - InitAtomicCleanInputs(is_all_nop_node, kernel); + InitAtomicCleanInputs(enable_fusion_clear, kernel); } } } @@ -561,7 +562,7 @@ void Somas::InitCommonNodeInputs(bool is_all_nop_node, const CNodePtr &kernel) { } } -void Somas::InitAtomicCleanInputs(bool is_all_nop_node, const CNodePtr &kernel) { +void Somas::InitAtomicCleanInputs(bool enable_fusion_clear, const CNodePtr &kernel) { auto node = nodes_map_[kernel.get()]; MS_EXCEPTION_IF_NULL(node); auto stream = node->GetStream(); @@ -588,9 +589,11 @@ void Somas::InitAtomicCleanInputs(bool is_all_nop_node, const CNodePtr &kernel) auto input_somas_tensor = pre_somas_node->output_tensors_[index]; MS_EXCEPTION_IF_NULL(input_somas_tensor); node->input_tensors_.push_back(input_somas_tensor); - input_somas_tensor->lifelong_value_ = kLifeLongGraphAll; - MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_ - << " 's output" << index << " to lifelong"; + if (enable_fusion_clear) { + input_somas_tensor->lifelong_value_ = kLifeLongGraphAll; + MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_ + << " 's output" << index << " to lifelong"; + } } } // set clean workspace tensors @@ -604,9 +607,11 @@ void Somas::InitAtomicCleanInputs(bool is_all_nop_node, const CNodePtr &kernel) auto input_somas_tensor = pre_somas_node->workspace_tensors_[index]; MS_EXCEPTION_IF_NULL(input_somas_tensor); node->input_tensors_.push_back(input_somas_tensor); - input_somas_tensor->lifelong_value_ = kLifeLongGraphAll; - MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_ - << " 's workspace" << index << " to lifelong"; + if (enable_fusion_clear) { + input_somas_tensor->lifelong_value_ = kLifeLongGraphAll; + MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_ + << " 's workspace" << index << " to lifelong"; + } } } }