| 1996 | } |
| 1997 | |
| 1998 | Status GraphPartitionerBase::SplitGraphInternal( |
| 1999 | std::vector<SubGraph> *sub_graphs, |
| 2000 | SubGraph *main_graph, |
| 2001 | bool needResetSwitchOp) |
| 2002 | { |
| 2003 | Status status; |
| 2004 | |
| 2005 | GraphInfo g_info; |
| 2006 | if (!opts_.control_flow_added) { |
| 2007 | // Add the "code" for distributed execution of control flow. Code is |
| 2008 | // added only for the frames that are placed on multiple devices. The |
| 2009 | // new graph is an equivalent transformation of the original graph and |
| 2010 | // has the property that it can be subsequently partitioned arbitrarily |
| 2011 | // (down to the level of individual device) for distributed execution. |
| 2012 | status = AddControlFlow(opts_, graph_, &g_info); |
| 2013 | if (!status.ok()) return status; |
| 2014 | } |
| 2015 | |
| 2016 | // At this point, all the graph mutations have been done. Build memory |
| 2017 | // and device type info for every node and edge in the graph. |
| 2018 | RETURN_IF_NOT_OK(BuildMemoryDeviceInfo(*graph_, &g_info)); |
| 2019 | if (needResetSwitchOp) { |
| 2020 | RETURN_IF_NOT_OK(ResetSwitchOpDevice()); |
| 2021 | } |
| 2022 | |
| 2023 | std::vector<bool> node_visited(graph_->num_node_ids(), false); |
| 2024 | for (const Node* cur : graph_->nodes()) { |
| 2025 | if (!cur->IsOp()) { |
| 2026 | continue; // Skip Sink/Source nodes. |
| 2027 | } |
| 2028 | |
| 2029 | if (node_visited[cur->id()]) { |
| 2030 | continue; |
| 2031 | } |
| 2032 | |
| 2033 | std::string cur_loc = opts_.node_to_loc(cur); |
| 2034 | SubGraph cur_sub_graph(cur_loc); |
| 2035 | cur_sub_graph.SetLoc(cur_loc); |
| 2036 | node_visited[cur->id()] = true; |
| 2037 | |
| 2038 | if ((cur->IsVariable() || cur->IsKvVarHandle ()) && |
| 2039 | !is_main_loc_func_(cur_loc)) { |
| 2040 | for (const auto& out_edge : cur->out_edges()) { |
| 2041 | if (!out_edge->dst()->IsOp()) continue; |
| 2042 | if (opts_.node_to_loc(out_edge->dst()) != cur_loc) { |
| 2043 | cur_sub_graph.AddNode(cur); |
| 2044 | cur_sub_graph.SetOnlyVariable(); |
| 2045 | sub_graphs->push_back(cur_sub_graph); |
| 2046 | break; |
| 2047 | } |
| 2048 | } |
| 2049 | continue; |
| 2050 | } |
| 2051 | |
| 2052 | if (cur->IsIdentity() && !is_main_loc_func_(cur_loc)) { |
| 2053 | const Node* var = GetVarOfIdentity(opts_, cur); |
| 2054 | if (var != NULL) { |
| 2055 | bool is_boundary = false; |
nothing calls this directly
no test coverage detected