| 1716 | } // namespace |
| 1717 | |
| 1718 | Status PartitionWithTensorFuse(const PartitionOptions& opts, Graph* g, |
| 1719 | std::unordered_map<string, GraphDef>* partitions) { |
| 1720 | |
| 1721 | Status status; |
| 1722 | partitions->clear(); |
| 1723 | |
| 1724 | std::vector<ControlFlowInfo> cf_info; |
| 1725 | GraphInfo g_info; |
| 1726 | if (!opts.control_flow_added) { |
| 1727 | // Add the "code" for distributed execution of control flow. Code is |
| 1728 | // added only for the frames that are placed on multiple devices. The |
| 1729 | // new graph is an equivalent transformation of the original graph and |
| 1730 | // has the property that it can be subsequently partitioned arbitrarily |
| 1731 | // (down to the level of individual device) for distributed execution. |
| 1732 | status = AddControlFlow(opts, g, &g_info); |
| 1733 | cf_info = g_info.cf_info; |
| 1734 | if (!status.ok()) return status; |
| 1735 | } else { |
| 1736 | status = BuildControlFlowInfo(g, &cf_info); |
| 1737 | if (!status.ok()) return status; |
| 1738 | } |
| 1739 | |
| 1740 | // At this point, all the graph mutations have been done. Build memory |
| 1741 | // and device type info for every node and edge in the graph. |
| 1742 | status = BuildMemoryDeviceInfo(*g, &g_info); |
| 1743 | if (!status.ok()) return status; |
| 1744 | |
| 1745 | std::unordered_set<Node*> done; |
| 1746 | std::unordered_map<std::string, std::unordered_set<Node*> > subgraph_nodes; |
| 1747 | for (Node* n : g->nodes()) { |
| 1748 | if (!n->IsOp()) continue; |
| 1749 | std::string loc = opts.node_to_loc(n); |
| 1750 | subgraph_nodes[loc].insert(n); |
| 1751 | } |
| 1752 | |
| 1753 | int32 num_data = 0; |
| 1754 | int32 num_control = 0; |
| 1755 | // Used to store control edges. |
| 1756 | DupRecvTable dup_recv(3); |
| 1757 | GlobalDupFuseRecvTable dup_global_fuse_recv; |
| 1758 | |
| 1759 | // topological partition |
| 1760 | while (true) { |
| 1761 | size_t left_nodes_count = 0; |
| 1762 | for (auto nodes : subgraph_nodes) { |
| 1763 | left_nodes_count += nodes.second.size(); |
| 1764 | } |
| 1765 | if (left_nodes_count == 0) break; |
| 1766 | |
| 1767 | std::unordered_map<std::string, std::unordered_set<Node*> > ready_subgraph_nodes; |
| 1768 | for (auto& node_set : subgraph_nodes) { |
| 1769 | GetReadyNodes(done, node_set.first, &(node_set.second), &ready_subgraph_nodes); |
| 1770 | } |
| 1771 | |
| 1772 | for (auto nodes : ready_subgraph_nodes) { |
| 1773 | string dstp = nodes.first; |
| 1774 | std::vector<const Edge*> inputs; |
| 1775 | // Used to store non-control inputs |