!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) {
|
void Somas::InitSomasInputTensors(const session::KernelGraph *graph) {
|
||||||
MS_EXCEPTION_IF_NULL(graph);
|
MS_EXCEPTION_IF_NULL(graph);
|
||||||
bool is_all_nop_node = opt::IsAllNopNode(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();
|
auto kernel_cnodes = graph->execution_order();
|
||||||
for (const auto &kernel : kernel_cnodes) {
|
for (const auto &kernel : kernel_cnodes) {
|
||||||
if (AnfAlgo::GetCNodeName(kernel) != kAtomicAddrCleanOpName) {
|
if (AnfAlgo::GetCNodeName(kernel) != kAtomicAddrCleanOpName) {
|
||||||
InitCommonNodeInputs(is_all_nop_node, kernel);
|
InitCommonNodeInputs(is_all_nop_node, kernel);
|
||||||
} else {
|
} 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()];
|
auto node = nodes_map_[kernel.get()];
|
||||||
MS_EXCEPTION_IF_NULL(node);
|
MS_EXCEPTION_IF_NULL(node);
|
||||||
auto stream = node->GetStream();
|
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];
|
auto input_somas_tensor = pre_somas_node->output_tensors_[index];
|
||||||
MS_EXCEPTION_IF_NULL(input_somas_tensor);
|
MS_EXCEPTION_IF_NULL(input_somas_tensor);
|
||||||
node->input_tensors_.push_back(input_somas_tensor);
|
node->input_tensors_.push_back(input_somas_tensor);
|
||||||
input_somas_tensor->lifelong_value_ = kLifeLongGraphAll;
|
if (enable_fusion_clear) {
|
||||||
MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_
|
input_somas_tensor->lifelong_value_ = kLifeLongGraphAll;
|
||||||
<< " 's output" << index << " to lifelong";
|
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
|
// 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];
|
auto input_somas_tensor = pre_somas_node->workspace_tensors_[index];
|
||||||
MS_EXCEPTION_IF_NULL(input_somas_tensor);
|
MS_EXCEPTION_IF_NULL(input_somas_tensor);
|
||||||
node->input_tensors_.push_back(input_somas_tensor);
|
node->input_tensors_.push_back(input_somas_tensor);
|
||||||
input_somas_tensor->lifelong_value_ = kLifeLongGraphAll;
|
if (enable_fusion_clear) {
|
||||||
MS_LOG(INFO) << "Set " << node->scope_full_name_ << "'s Input node " << pre_somas_node->scope_full_name_
|
input_somas_tensor->lifelong_value_ = kLifeLongGraphAll;
|
||||||
<< " 's workspace" << index << " to lifelong";
|
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