!14533 enable reuse for non-fusion AtomicAddrClean node

From: @zhaosida_hw
Reviewed-by: @kisnwang,@zhoufeng54
Signed-off-by: @zhoufeng54
This commit is contained in:
mindspore-ci-bot 2021-04-01 19:20:13 +08:00 committed by Gitee
commit 0a870440e0
1 changed files with 13 additions and 8 deletions

View File

@ -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";
}
}
}
}