!14533 enable reuse for non-fusion AtomicAddrClean node
From: @zhaosida_hw Reviewed-by: @kisnwang,@zhoufeng54 Signed-off-by: @zhoufeng54
This commit is contained in:
commit
0a870440e0
|
@ -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";
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue