From 8a035778b6c8990e87fa4d2c3a7eb7072af1c50a Mon Sep 17 00:00:00 2001 From: kswang Date: Thu, 24 Dec 2020 10:31:59 +0800 Subject: [PATCH] add parallel split --- .../ccsrc/backend/session/session_basic.cc | 7 + .../ccsrc/pybind_api/utils/ms_context_py.cc | 1 + mindspore/ccsrc/vm/graph_partition.cc | 444 ++++++++++++++---- mindspore/core/utils/ms_context.cc | 1 + mindspore/core/utils/ms_context.h | 1 + 5 files changed, 367 insertions(+), 87 deletions(-) diff --git a/mindspore/ccsrc/backend/session/session_basic.cc b/mindspore/ccsrc/backend/session/session_basic.cc index dd261a582d..ca6c351de9 100644 --- a/mindspore/ccsrc/backend/session/session_basic.cc +++ b/mindspore/ccsrc/backend/session/session_basic.cc @@ -1421,6 +1421,13 @@ std::vector ExtendNodeUsers(const FuncGraphManagerPtr &front_func_gr continue; } if (IsPrimitiveCNode(user.first, prim::kPrimDepend)) { + auto depend_cnode = user.first->cast(); + if (depend_cnode == nullptr) { + continue; + } + if (front_node != depend_cnode->input(1)) { + continue; + } auto res = ExtendNodeUsers(front_func_graph_manager, user.first); result.insert(result.end(), res.begin(), res.end()); continue; diff --git a/mindspore/ccsrc/pybind_api/utils/ms_context_py.cc b/mindspore/ccsrc/pybind_api/utils/ms_context_py.cc index 140569415e..006f4ab578 100644 --- a/mindspore/ccsrc/pybind_api/utils/ms_context_py.cc +++ b/mindspore/ccsrc/pybind_api/utils/ms_context_py.cc @@ -83,6 +83,7 @@ REGISTER_PYBIND_DEFINE(MsContextPy, ([](const py::module *m) { .value("precompile_only", MsCtxParam::MS_CTX_PRECOMPILE_ONLY) .value("enable_profiling", MsCtxParam::MS_CTX_ENABLE_PROFILING) .value("save_graphs", MsCtxParam::MS_CTX_SAVE_GRAPHS_FLAG) + .value("enable_parallel_split", MsCtxParam::MS_CTX_ENABLE_PARALLEL_SPLIT) .value("max_device_memory", MsCtxParam::MS_CTX_MAX_DEVICE_MEMORY) .value("mode", MsCtxParam::MS_CTX_EXECUTION_MODE) .value("device_target", MsCtxParam::MS_CTX_DEVICE_TARGET) diff --git a/mindspore/ccsrc/vm/graph_partition.cc b/mindspore/ccsrc/vm/graph_partition.cc index 74fa47d9e1..51a975387f 100644 --- a/mindspore/ccsrc/vm/graph_partition.cc +++ b/mindspore/ccsrc/vm/graph_partition.cc @@ -22,9 +22,11 @@ #include #include #include +#include #include "base/core_ops.h" #include "utils/utils.h" #include "utils/ms_context.h" +#include "ps/ps_context.h" #ifdef ENABLE_GE #include "transform/graph_ir/convert.h" #endif @@ -34,6 +36,21 @@ const char kMsVm[] = "vm"; const char kGeVm[] = "ge"; namespace compile { namespace { +std::string GetOtherTarget(const std::vector &nodes) { + auto context_ptr = MsContext::GetInstance(); + MS_EXCEPTION_IF_NULL(context_ptr); + std::string default_target = context_ptr->get_param(MS_CTX_DEVICE_TARGET); + for (auto &node : nodes) { + if (!node->isa()) { + continue; + } + std::string cur_target = GetCNodeTarget(node); + if (cur_target != default_target) { + return cur_target; + } + } + return ""; +} bool ExtractNodes(const FuncGraphPtr &graph, const AnfNodePtr &prior_node, const AnfNodePtr &behind_node, std::vector *prior_nodes, std::vector *depend_nodes) { MS_EXCEPTION_IF_NULL(prior_node); @@ -204,7 +221,7 @@ std::vector SplitSort(const FuncGraphPtr &graph, const std::string & to_visit.swap(next_to_visit); handle_target = next_target; } - auto &node = to_visit.top(); + auto node = to_visit.top(); MS_EXCEPTION_IF_NULL(node); to_visit.pop(); result.emplace_back(node); @@ -248,18 +265,24 @@ std::vector SplitSort(const FuncGraphPtr &graph, const std::string & return result; } -void AddSegmentDependency(const FuncGraphPtr &graph, const std::string &default_target, - const std::map &node_to_segment) { +struct GraphNodesDependencyInfo { + std::stack independent_nodes_; + std::map input_num_; + std::map> output_edges_; +}; + +GraphNodesDependencyInfo GetNodesDependencyInfo(const FuncGraphPtr &graph) { + MS_EXCEPTION_IF_NULL(graph); + GraphNodesDependencyInfo info; std::stack to_visit; std::map nodes_ref; std::map> control_edges; CalcNodeRefCount(graph, &nodes_ref, &control_edges); to_visit.push(graph->get_return()); while (!to_visit.empty()) { - auto &node = to_visit.top(); - MS_EXCEPTION_IF_NULL(node); + auto node = to_visit.top(); to_visit.pop(); - if (!node->isa()) { + if (node == nullptr || !node->isa()) { continue; } auto cnode = node->cast(); @@ -269,22 +292,22 @@ void AddSegmentDependency(const FuncGraphPtr &graph, const std::string &default_ if (ctrl_inputs != control_edges.end()) { node_inputs.insert(node_inputs.end(), ctrl_inputs->second.begin(), ctrl_inputs->second.end()); } - GraphSegmentPtr node_segment{nullptr}; - if (!IsPrimitiveCNode(cnode, prim::kPrimControlDepend)) { - auto node_iter = node_to_segment.find(node); - if (node_iter != node_to_segment.end()) { - node_segment = node_iter->second; - } - } + bool independent = true; for (auto &input : node_inputs) { - if (node_segment != nullptr && !node_segment->is_cut_ && input->isa()) { - GraphSegmentPtr input_segment{nullptr}; - auto input_iter = node_to_segment.find(input); - if (input_iter != node_to_segment.end()) { - input_segment = input_iter->second; + if (input->isa()) { + independent = false; + auto output_edge_iter = info.output_edges_.find(input); + if (output_edge_iter != info.output_edges_.end()) { + auto &edges = output_edge_iter->second; + edges.emplace_back(node); + } else { + info.output_edges_[input] = {node}; } - if (input_segment != nullptr && input_segment != node_segment && !input_segment->is_cut_) { - node_segment->AddPreSegment(input_segment); + auto input_num_iter = info.input_num_.find(node); + if (input_num_iter != info.input_num_.end()) { + input_num_iter->second++; + } else { + info.input_num_[node] = 1; } } auto ref_iter = nodes_ref.find(input); @@ -296,83 +319,193 @@ void AddSegmentDependency(const FuncGraphPtr &graph, const std::string &default_ } to_visit.push(input); } + if (independent) { + info.independent_nodes_.push(node); + } } + return info; +} + +struct VisitNodesInfo { + std::queue default_target_nodes_; + std::queue other_target_nodes_; + std::map seed_cast_next_node_; +}; + +VisitNodesInfo GetVisitNodesInfo(const GraphNodesDependencyInfo &dependency_info, const std::string &default_target, + const std::string &other_target) { + VisitNodesInfo result; + auto independent_nodes = dependency_info.independent_nodes_; + while (!independent_nodes.empty()) { + auto seed_node = independent_nodes.top(); + independent_nodes.pop(); + MS_EXCEPTION_IF_NULL(seed_node); + auto node_target = GetCNodeTarget(seed_node); + if (IsPrimitiveCNode(seed_node, prim::kPrimCast)) { + auto output_edges_iter = dependency_info.output_edges_.find(seed_node); + if (output_edges_iter != dependency_info.output_edges_.end() && output_edges_iter->second.size() == 1) { + auto &cast_next_node = output_edges_iter->second[0]; + auto input_num_iter = dependency_info.input_num_.find(cast_next_node); + if (input_num_iter == dependency_info.input_num_.end()) { + MS_LOG(EXCEPTION) << "Node input num not found!"; + } + if (input_num_iter->second > 1 && node_target == GetCNodeTarget(cast_next_node)) { + result.seed_cast_next_node_[cast_next_node] = seed_node; + continue; + } + } + } + if (node_target == default_target) { + result.default_target_nodes_.push(seed_node); + } else if (node_target == other_target) { + result.other_target_nodes_.push(seed_node); + } else { + MS_LOG(EXCEPTION) << "Only support two difference targets"; + } + } + return result; } -std::vector ParallelSplitSort(const FuncGraphPtr &graph, const std::string &default_target) { +std::string ParallelSortDecideNextHandleTarget(const std::vector &output_edges, + const std::string &node_target, + std::map *node_input_target_map) { + MS_EXCEPTION_IF_NULL(node_input_target_map); + std::string next_target = node_target; + for (auto &dst_node : output_edges) { + auto input_target_iter = node_input_target_map->find(dst_node); + if (input_target_iter != node_input_target_map->end() && input_target_iter->second != node_target) { + next_target = input_target_iter->second; + break; + } + auto dst_node_target = GetCNodeTarget(dst_node); + if (dst_node_target != node_target) { + next_target = dst_node_target; + break; + } + } + for (auto &dst_node : output_edges) { + (*node_input_target_map)[dst_node] = node_target; + } + return next_target; +} + +void ParallelSortVisitNodeEdges(const std::vector &output_edges, GraphNodesDependencyInfo *dependency_info, + VisitNodesInfo *visit_nodes_info, const std::string &default_target) { + MS_EXCEPTION_IF_NULL(dependency_info); + MS_EXCEPTION_IF_NULL(visit_nodes_info); + for (auto &dst_node : output_edges) { + auto dst_node_target = GetCNodeTarget(dst_node); + auto input_num_iter = dependency_info->input_num_.find(dst_node); + if (input_num_iter == dependency_info->input_num_.end()) { + MS_LOG(EXCEPTION) << "Node input num not found!"; + } + input_num_iter->second--; + if (input_num_iter->second == 1 && + visit_nodes_info->seed_cast_next_node_.find(dst_node) != visit_nodes_info->seed_cast_next_node_.end()) { + input_num_iter->second--; + } + if (input_num_iter->second > 0) { + continue; + } + if (dst_node_target == default_target) { + visit_nodes_info->default_target_nodes_.push(dst_node); + } else { + visit_nodes_info->other_target_nodes_.push(dst_node); + } + } +} + +std::vector ParallelSort(const FuncGraphPtr &graph, const std::string &default_target, + const std::string &other_target) { + MS_EXCEPTION_IF_NULL(graph); + auto dependency_info = GetNodesDependencyInfo(graph); + auto visit_nodes_info = GetVisitNodesInfo(dependency_info, default_target, other_target); std::vector result; - std::stack handle_nodes; - std::stack next_handle_nodes; - std::stack wait_handle_nodes; + std::string handle_target; + if (!visit_nodes_info.default_target_nodes_.empty()) { + handle_target = default_target; + } else { + handle_target = other_target; + } + std::map node_input_target_map; + while (!visit_nodes_info.default_target_nodes_.empty() || !visit_nodes_info.other_target_nodes_.empty()) { + AnfNodePtr ready_node; + if ((!visit_nodes_info.default_target_nodes_.empty() && handle_target == default_target) || + visit_nodes_info.other_target_nodes_.empty()) { + ready_node = visit_nodes_info.default_target_nodes_.front(); + visit_nodes_info.default_target_nodes_.pop(); + handle_target = default_target; + } else { + ready_node = visit_nodes_info.other_target_nodes_.front(); + visit_nodes_info.other_target_nodes_.pop(); + handle_target = other_target; + } + MS_EXCEPTION_IF_NULL(ready_node); + auto cast_map_iter = visit_nodes_info.seed_cast_next_node_.find(ready_node); + if (cast_map_iter != visit_nodes_info.seed_cast_next_node_.end()) { + result.emplace_back(cast_map_iter->second); + } + result.emplace_back(ready_node); + auto output_edges_iter = dependency_info.output_edges_.find(ready_node); + if (output_edges_iter == dependency_info.output_edges_.end()) { + continue; + } + auto &output_edges = output_edges_iter->second; + handle_target = ParallelSortDecideNextHandleTarget(output_edges, handle_target, &node_input_target_map); + ParallelSortVisitNodeEdges(output_edges, &dependency_info, &visit_nodes_info, default_target); + } + return result; +} + +void AddSegmentDependency(const FuncGraphPtr &graph, const std::string &default_target, + const std::map &node_to_segment) { + std::stack to_visit; std::map nodes_ref; std::map> control_edges; CalcNodeRefCount(graph, &nodes_ref, &control_edges); - std::string handle_target = default_target; - std::string next_target = ""; - handle_nodes.push(graph->get_return()); - while (!handle_nodes.empty() || !next_handle_nodes.empty() || !wait_handle_nodes.empty()) { - if (handle_nodes.empty()) { - handle_nodes.swap(next_handle_nodes); - handle_target.swap(next_target); - next_handle_nodes.swap(wait_handle_nodes); - } - auto &node = handle_nodes.top(); + to_visit.push(graph->get_return()); + while (!to_visit.empty()) { + auto &node = to_visit.top(); MS_EXCEPTION_IF_NULL(node); - handle_nodes.pop(); - result.emplace_back(node); + to_visit.pop(); if (!node->isa()) { continue; } auto cnode = node->cast(); MS_EXCEPTION_IF_NULL(cnode); auto node_inputs = cnode->inputs(); - std::reverse(node_inputs.begin(), node_inputs.end()); auto ctrl_inputs = control_edges.find(node); if (ctrl_inputs != control_edges.end()) { node_inputs.insert(node_inputs.end(), ctrl_inputs->second.begin(), ctrl_inputs->second.end()); } - std::vector same_target_ready_inputs; - std::vector diff_target_ready_inputs; - for (auto &input : node_inputs) { - auto iter = nodes_ref.find(input); - if (iter != nodes_ref.end()) { - iter->second--; - if (iter->second != 0) { - continue; - } - } - if (!input->isa()) { - handle_nodes.push(input); - continue; + GraphSegmentPtr node_segment{nullptr}; + if (!IsPrimitiveCNode(cnode, prim::kPrimControlDepend)) { + auto node_iter = node_to_segment.find(node); + if (node_iter != node_to_segment.end()) { + node_segment = node_iter->second; } - std::string input_target = GetCNodeTarget(input); - if (input_target == handle_target) { - same_target_ready_inputs.emplace_back(input); - } else { - if (next_target.empty()) { - next_target = input_target; + } + for (auto &input : node_inputs) { + if (node_segment != nullptr && !node_segment->is_cut_ && input->isa()) { + GraphSegmentPtr input_segment{nullptr}; + auto input_iter = node_to_segment.find(input); + if (input_iter != node_to_segment.end()) { + input_segment = input_iter->second; } - if (input_target != next_target) { - MS_LOG(EXCEPTION) << "Only support two different target"; + if (input_segment != nullptr && input_segment != node_segment && !input_segment->is_cut_) { + node_segment->AddPreSegment(input_segment); } - diff_target_ready_inputs.emplace_back(input); } - } - if (diff_target_ready_inputs.size() == 0) { - for (auto &input : same_target_ready_inputs) { - handle_nodes.push(input); - } - } else { - for (auto &input : same_target_ready_inputs) { - wait_handle_nodes.push(input); - } - for (auto &input : diff_target_ready_inputs) { - next_handle_nodes.push(input); + auto ref_iter = nodes_ref.find(input); + if (ref_iter != nodes_ref.end()) { + ref_iter->second--; + if (ref_iter->second != 0) { + continue; + } } + to_visit.push(input); } } - std::reverse(result.begin(), result.end()); - return result; } bool IsSubGraph(const AnfNodePtr &node) { @@ -397,6 +530,150 @@ bool IsSubGraph(const AnfNodePtr &node) { } return false; } + +bool IsShapeDynamic(const abstract::ShapePtr &shape) { + MS_EXCEPTION_IF_NULL(shape); + return std::any_of(shape->shape().begin(), shape->shape().end(), [](int64_t s) { return s < 0; }); +} + +bool IsNodeOutputDynamicShape(const CNodePtr &node) { + MS_EXCEPTION_IF_NULL(node); + auto base_shape = node->Shape(); + if (base_shape == nullptr) { + MS_LOG(INFO) << "Invalid base shape, node: " << node->fullname_with_scope(); + return false; + } + if (base_shape->isa()) { + if (IsShapeDynamic(base_shape->cast())) { + return true; + } + } else if (base_shape->isa()) { + auto tuple_shape = base_shape->cast(); + MS_EXCEPTION_IF_NULL(tuple_shape); + for (size_t i = 0; i < tuple_shape->size(); i++) { + auto b_shape = (*tuple_shape)[i]; + if (!b_shape->isa()) { + continue; + } + if (IsShapeDynamic(b_shape->cast())) { + return true; + } + } + } + return false; +} + +void AddSegment(const std::vector &nodes, std::vector *segments, + std::map *node_to_segment) { + MS_EXCEPTION_IF_NULL(segments); + MS_EXCEPTION_IF_NULL(node_to_segment); + auto segment = std::make_shared(nodes, false); + segments->emplace_back(segment); + for (auto &node : nodes) { + (*node_to_segment)[node] = segment; + } +} + +struct SplitDynamicNodesHelper { + void AddNode(const AnfNodePtr &node, bool is_dynamic_shape) { + if (is_dynamic_shape) { + pre_dynamic_nodes.emplace_back(node); + pre_dynamic_nodes_set.insert(node); + } else { + pre_common_nodes.emplace_back(node); + pre_common_nodes_set.insert(node); + } + pre_nodes.emplace_back(node); + } + + void AddSegments(std::vector *segments, std::map *node_to_segment) { + if (pre_nodes.size() < merge_node_threshold) { + AddSegment(pre_nodes, segments, node_to_segment); + } else { + if (pre_common_nodes.size() > 0) { + AddSegment(pre_common_nodes, segments, node_to_segment); + } + if (pre_dynamic_nodes.size() > 0) { + AddSegment(pre_dynamic_nodes, segments, node_to_segment); + } + } + pre_common_nodes.clear(); + pre_common_nodes_set.clear(); + pre_dynamic_nodes.clear(); + pre_dynamic_nodes_set.clear(); + pre_nodes.clear(); + } + std::vector pre_nodes; + std::vector pre_dynamic_nodes; + std::vector pre_common_nodes; + std::set pre_common_nodes_set; + std::set pre_dynamic_nodes_set; + size_t merge_node_threshold = 6; +}; + +void SplitDynamicNodeSegment(const std::vector &segment_nodes, std::vector *segments, + std::map *node_to_segment, + const std::set &dynamic_nodes_set) { + SplitDynamicNodesHelper helper; + bool is_last_node_dynamic = false; + for (auto &node : segment_nodes) { + auto cnode = node->cast(); + MS_EXCEPTION_IF_NULL(cnode); + if (IsPrimitiveCNode(cnode, prim::kPrimControlDepend)) { + helper.AddNode(node, is_last_node_dynamic); + continue; + } + auto &inputs = cnode->inputs(); + bool has_dynamic_shape = dynamic_nodes_set.find(node) != dynamic_nodes_set.end(); + bool depend_common_node = false; + bool depend_dynamic_node = false; + for (size_t i = 1; i < inputs.size(); ++i) { + if (dynamic_nodes_set.find(inputs[i]) != dynamic_nodes_set.end()) { + has_dynamic_shape = true; + } + if (helper.pre_common_nodes_set.find(inputs[i]) != helper.pre_common_nodes_set.end()) { + depend_common_node = true; + } + if (helper.pre_dynamic_nodes_set.find(inputs[i]) != helper.pre_dynamic_nodes_set.end()) { + depend_dynamic_node = true; + } + } + if (has_dynamic_shape) { + if (depend_common_node) { + helper.AddSegments(segments, node_to_segment); + } + is_last_node_dynamic = true; + } else { + if (depend_dynamic_node) { + helper.AddSegments(segments, node_to_segment); + } + is_last_node_dynamic = false; + } + helper.AddNode(node, is_last_node_dynamic); + } + helper.AddSegments(segments, node_to_segment); +} + +void NodesToSegments(const std::vector &segment_nodes, std::vector *segments, + std::map *node_to_segment) { + if (segment_nodes.size() == 0) { + return; + } + MS_EXCEPTION_IF_NULL(segments); + MS_EXCEPTION_IF_NULL(node_to_segment); + std::set dynamic_nodes_set; + for (auto &node : segment_nodes) { + auto cnode = node->cast(); + if (IsNodeOutputDynamicShape(cnode)) { + (void)dynamic_nodes_set.insert(node); + } + } + if (dynamic_nodes_set.empty()) { + AddSegment(segment_nodes, segments, node_to_segment); + return; + } + SplitDynamicNodeSegment(segment_nodes, segments, node_to_segment, dynamic_nodes_set); +} } // namespace GraphPartition::GraphPartition(const std::vector &cut_list, const std::string &backend_name) @@ -461,31 +738,23 @@ std::vector GraphPartition::Partition(const FuncGraphPtr &graph MS_EXCEPTION_IF_NULL(context_ptr); std::string default_target = context_ptr->get_param(MS_CTX_DEVICE_TARGET); if (contain_multi_target) { - if (graph != nullptr) { - nodes = SplitSort(graph, default_target); + if (context_ptr->get_param(MS_CTX_ENABLE_PARALLEL_SPLIT)) { + auto other_target = GetOtherTarget(nodes); + nodes = ParallelSort(graph, default_target, other_target); } else { - nodes = ParallelSplitSort(graph, default_target); + nodes = SplitSort(graph, default_target); } nodes = OptimizeGetItemOrder(nodes); } std::vector segments; std::vector segment_nodes; std::map node_to_segment; - auto new_segment = [&segments, &segment_nodes, &node_to_segment]() { - if (segment_nodes.size() != 0) { - auto segment = std::make_shared(segment_nodes, false); - segments.emplace_back(segment); - for (auto node : segment_nodes) { - node_to_segment[node] = segment; - } - segment_nodes.clear(); - } - }; std::string last_target; for (auto &node : nodes) { MS_EXCEPTION_IF_NULL(node); if (IsCut(node)) { - new_segment(); + NodesToSegments(segment_nodes, &segments, &node_to_segment); + segment_nodes.clear(); segment_nodes.emplace_back(node); auto segment = std::make_shared(segment_nodes, true); segments.push_back(segment); @@ -494,7 +763,8 @@ std::vector GraphPartition::Partition(const FuncGraphPtr &graph if (contain_multi_target) { std::string cur_target = GetCNodeTarget(node); if (cur_target != last_target && !last_target.empty()) { - new_segment(); + NodesToSegments(segment_nodes, &segments, &node_to_segment); + segment_nodes.clear(); } last_target = cur_target; } diff --git a/mindspore/core/utils/ms_context.cc b/mindspore/core/utils/ms_context.cc index adeac0b9b6..ed12e9111c 100644 --- a/mindspore/core/utils/ms_context.cc +++ b/mindspore/core/utils/ms_context.cc @@ -72,6 +72,7 @@ MsContext::MsContext(const std::string &policy, const std::string &target) { set_param(MS_CTX_PRINT_FILE_PATH, ""); set_param(MS_CTX_ENABLE_GRAPH_KERNEL, false); set_param(MS_CTX_ENABLE_SPARSE, false); + set_param(MS_CTX_ENABLE_PARALLEL_SPLIT, false); backend_policy_ = policy_map_[policy]; } diff --git a/mindspore/core/utils/ms_context.h b/mindspore/core/utils/ms_context.h index 79802e6f6b..98948d10e8 100644 --- a/mindspore/core/utils/ms_context.h +++ b/mindspore/core/utils/ms_context.h @@ -73,6 +73,7 @@ enum MsCtxParam : unsigned { MS_CTX_PRECOMPILE_ONLY, MS_CTX_ENABLE_PROFILING, MS_CTX_SAVE_GRAPHS_FLAG, + MS_CTX_ENABLE_PARALLEL_SPLIT, MS_CTX_TYPE_BOOL_END, // paramater of type int