| 1400 | } |
| 1401 | |
| 1402 | Status Partition(const PartitionOptions& opts, Graph* g, |
| 1403 | std::unordered_map<string, GraphDef>* partitions) { |
| 1404 | Status status; |
| 1405 | partitions->clear(); |
| 1406 | |
| 1407 | GraphInfo g_info; |
| 1408 | if (!opts.control_flow_added) { |
| 1409 | // Add the "code" for distributed execution of control flow. Code is |
| 1410 | // added only for the frames that are placed on multiple devices. The |
| 1411 | // new graph is an equivalent transformation of the original graph and |
| 1412 | // has the property that it can be subsequently partitioned arbitrarily |
| 1413 | // (down to the level of individual device) for distributed execution. |
| 1414 | status = AddControlFlow(opts, g, &g_info); |
| 1415 | if (!status.ok()) return status; |
| 1416 | } |
| 1417 | |
| 1418 | // At this point, all the graph mutations have been done. Build memory |
| 1419 | // and device type info for every node and edge in the graph. |
| 1420 | status = BuildMemoryDeviceInfo(*g, &g_info); |
| 1421 | if (!status.ok()) return status; |
| 1422 | |
| 1423 | string dstp; |
| 1424 | std::vector<const Edge*> inputs; |
| 1425 | DupRecvTable dup_recv(3); |
| 1426 | // For a node dst, 'ref_recvs' remembers the recvs introduced by a ref |
| 1427 | // edge to dst. 'ref_control_inputs' remembers the inputs by a non-ref |
| 1428 | // edge to dst. We will add a control edge for every pair in |
| 1429 | // (ref_recvs x ref_control_inputs). |
| 1430 | std::vector<NodeDef*> ref_recvs; |
| 1431 | std::vector<string> ref_control_inputs; |
| 1432 | |
| 1433 | int32 num_data = 0; |
| 1434 | int32 num_control = 0; |
| 1435 | for (const Node* dst : g->op_nodes()) { |
| 1436 | dstp = opts.node_to_loc(dst); |
| 1437 | GraphDef* dst_graph = &(*partitions)[dstp]; |
| 1438 | NodeDef* dst_def = dst_graph->add_node(); |
| 1439 | *dst_def = dst->def(); |
| 1440 | MergeDebugInfo(NodeDebugInfo(dst->def()), dst_def); |
| 1441 | dst_def->set_device(dst->assigned_device_name()); |
| 1442 | dst_def->clear_input(); // Inputs are filled below |
| 1443 | if (opts.need_to_record_start_times) { |
| 1444 | int64 start_time; |
| 1445 | status = GetNodeAttr(*dst_def, "_start_time", &start_time); |
| 1446 | if (errors::IsNotFound(status)) { |
| 1447 | start_time = opts.start_times[dst->id()].value(); |
| 1448 | AddNodeAttr("_start_time", start_time, dst_def); |
| 1449 | } else if (!status.ok()) { |
| 1450 | return status; |
| 1451 | } |
| 1452 | } |
| 1453 | |
| 1454 | // Arrange the incoming edges to dst so that input[i] holds the |
| 1455 | // input flowing into slot numbered i. Trailing entries in input[] |
| 1456 | // hold control edges. |
| 1457 | inputs.clear(); |
| 1458 | inputs.resize(dst->num_inputs(), nullptr); |
| 1459 | ref_recvs.clear(); |