p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!5461 add atomic clean for every communication op's inputs 

Merge pull request !5461 from laiyongqiang/atomic_clean_every_master
This commit is contained in:
mindspore-ci-bot 2020-09-01 14:26:05 +08:00 committed by Gitee
parent 660aa8e60d 0c1e391556
commit 402bb4c053
2 changed files with 46 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse.h
View File

@ -49,7 +49,7 @@ class MemReuseUtil {
}
MS_LOG(INFO) << "Total Dynamic Memory Size: " << total_dy_size_;
MS_LOG(INFO) << "Total WorkSpace Memory Size: " << total_workspace_size_;
MS_LOG(INFO) << "Total Reused WorkSpafce Memory Size: " << total_reuseworkspace_size_;
MS_LOG(INFO) << "Total Reused WorkSpace Memory Size: " << total_reuseworkspace_size_;
}
void SetAllInfo(const KernelGraph *graph);

45
mindspore/ccsrc/runtime/device/ascend/kernel_build_ascend.cc
View File

@ -205,6 +205,10 @@ static bool IsAtomicNode(const CNodePtr &kernel_node) {
}
// process output
std::vector<size_t> output_indexs = {};
if (AnfAlgo::HasNodeAttr(kAttrAtomicOutputIndexs, kernel_node)) {
output_indexs = AnfAlgo::GetNodeAttr<std::vector<size_t>>(kernel_node, kAttrAtomicOutputIndexs);
}
for (size_t i = 0; i < output_num; ++i) {
auto param_output = parameters_indexs.at(input_num + workspace_num + i);
if (param_output == 1) {
@ -212,7 +216,10 @@ static bool IsAtomicNode(const CNodePtr &kernel_node) {
MS_LOG(INFO) << "Atomic clear output index: " << i;
}
}
if (!output_indexs.empty()) {
std::set<size_t> s(output_indexs.begin(), output_indexs.end());
output_indexs.assign(s.begin(), s.end());
AnfAlgo::SetNodeAttr(kAttrAtomicOutputIndexs, MakeValue(output_indexs), kernel_node);
}
// process workspace
@ -244,11 +251,49 @@ bool KernelBuild(const mindspore::session::KernelGraph *kernel_graph_ptr) {
return ret;
}
std::map<AnfNodePtr, std::vector<size_t>> GetCommunicationOpInputInfo(
const mindspore::session::KernelGraph *kernel_graph) {
std::map<AnfNodePtr, std::vector<size_t>> comm_input_info_map;
for (auto &kernel : kernel_graph->execution_order()) {
auto input_num = AnfAlgo::GetInputTensorNum(kernel);
if (mindspore::session::AnfRuntimeAlgorithm::IsCommunicationOp(kernel)) {
for (size_t i = 0; i < input_num; i++) {
auto input_node = kernel->input(i + 1);
auto kernel_input = AnfAlgo::VisitKernelWithReturnType(input_node, 0, true);
MS_LOG(INFO) << " Add atomic clean for single communication op input, comm:" << kernel->fullname_with_scope()
<< " input_node: " << kernel_input.first->fullname_with_scope()
<< " index: " << kernel_input.second;
auto iter = comm_input_info_map.find(kernel_input.first);
if (iter != comm_input_info_map.end()) {
iter->second.push_back(kernel_input.second);
} else {
std::vector<size_t> indexes = {kernel_input.second};
comm_input_info_map[kernel_input.first] = indexes;
}
}
}
}
// remove duplicate index
for (auto &info : comm_input_info_map) {
std::set<size_t> s(info.second.begin(), info.second.end());
info.second.assign(s.begin(), s.end());
}
return comm_input_info_map;
}
void KernelBuildPreprocess(mindspore::session::KernelGraph *kernel_graph) {
MS_EXCEPTION_IF_NULL(kernel_graph);
std::vector<CNodePtr> new_nodes;
std::map<AnfNodePtr, std::vector<size_t>> comm_input_info_map = GetCommunicationOpInputInfo(kernel_graph);
for (const auto &anf_node : kernel_graph->execution_order()) {
std::string apply_function_name = AnfAlgo::GetCNodeName(anf_node);
if (comm_input_info_map.find(anf_node) != comm_input_info_map.end()) {
auto indexes = comm_input_info_map[anf_node];
AnfAlgo::SetNodeAttr(kAttrAtomicOutputIndexs, MakeValue(indexes), anf_node);
}
if (apply_function_name == prim::kPrimMaxPoolGrad->name() &&
AnfAlgo::GetKernelType(anf_node) == KernelType::AKG_KERNEL) {
auto clear_zero_prim = std::make_shared<Primitive>(kClearZeroOpName);