| 970 | } |
| 971 | |
| 972 | Status GraphPartitionerBase::CompleteSubGraph(SubGraph &sub_graph, int graph_idx) { |
| 973 | const std::string &loc = sub_graph.GetLoc(); |
| 974 | GraphDef &graph_def = sub_graph.GetGraphDef(); |
| 975 | |
| 976 | NodeDef* new_fuse_recv_node = nullptr; |
| 977 | std::unordered_map<std::string, int> src_to_slot; |
| 978 | bool use_fuse_recv = UseFuseRecv(sub_graph, graph_idx, |
| 979 | &new_fuse_recv_node, |
| 980 | src_to_slot); |
| 981 | |
| 982 | bool use_send_recv = false; |
| 983 | std::map<InputSrcKey, NodeDef*> input_src_nodes_map; |
| 984 | for (const Node* node: sub_graph.GetNodes()) { |
| 985 | NodeDef *node_def = graph_def.add_node(); |
| 986 | Status s = ConstructNodeDef(node, node_def); |
| 987 | RETURN_IF_NOT_OK(s); |
| 988 | if ((node->IsVariable() || node->IsKvVarHandle ()) |
| 989 | && !sub_graph.IsOnlyVariable()) { |
| 990 | continue; |
| 991 | } |
| 992 | |
| 993 | if (node->IsIdentity() && |
| 994 | GetVarOfIdentity(opts_, node) != NULL && |
| 995 | !sub_graph.IsOnlyVariable()) { |
| 996 | continue; |
| 997 | } |
| 998 | |
| 999 | std::vector<std::pair<std::string, const Edge*>> boundary_input_edges; |
| 1000 | if (use_fuse_recv) { |
| 1001 | s = ProcessSubNodeInputs(node, loc, node_def, |
| 1002 | &graph_def, &input_src_nodes_map, |
| 1003 | &boundary_input_edges, |
| 1004 | use_send_recv, |
| 1005 | new_fuse_recv_node, |
| 1006 | src_to_slot); |
| 1007 | } else { |
| 1008 | s = ProcessSubNodeInputs(node, loc, node_def, |
| 1009 | &graph_def, &input_src_nodes_map, |
| 1010 | &boundary_input_edges, |
| 1011 | use_send_recv); |
| 1012 | } |
| 1013 | RETURN_IF_NOT_OK(s); |
| 1014 | sub_graph.AddInputs(boundary_input_edges); |
| 1015 | |
| 1016 | std::vector<std::pair<std::string, const Edge*>> boundary_output_edges; |
| 1017 | s = ProcessSubNodeOutputs(node, loc, node_def, &graph_def, |
| 1018 | &boundary_output_edges, |
| 1019 | use_send_recv); |
| 1020 | RETURN_IF_NOT_OK(s); |
| 1021 | sub_graph.SetSendRecvFlag(use_send_recv); |
| 1022 | sub_graph.AddOutputs(boundary_output_edges); |
| 1023 | } |
| 1024 | |
| 1025 | return Status::OK(); |
| 1026 | } |
| 1027 | |
| 1028 | Status GraphPartitionerBase::CompleteSubGraphV2( |
| 1029 | SubGraph &sub_graph, int graph_idx, |
nothing calls this directly
no test coverage detected