| 182 | } |
| 183 | |
| 184 | int IdentifyLoops(const GraphDef& graph, |
| 185 | std::unordered_map<const NodeDef*, std::vector<int>>* loops) { |
| 186 | int num_components = 0; |
| 187 | std::unordered_map<const NodeDef*, int> components; |
| 188 | StronglyConnectedComponents(graph, &components, &num_components); |
| 189 | if (num_components <= 1) { |
| 190 | if (!components.empty() && components.begin()->second == -1) { |
| 191 | return 0; |
| 192 | } |
| 193 | } |
| 194 | |
| 195 | std::unordered_map<int, std::vector<const NodeDef*>> component_ids; |
| 196 | for (const auto it : components) { |
| 197 | int id = it.second; |
| 198 | if (id < 0) { |
| 199 | continue; |
| 200 | } |
| 201 | component_ids[id].push_back(it.first); |
| 202 | } |
| 203 | |
| 204 | int loop_id = 0; |
| 205 | for (const auto& component : component_ids) { |
| 206 | const std::vector<const NodeDef*>& component_nodes = component.second; |
| 207 | std::vector<std::pair<NodeDef*, string>> next_iter_nodes; |
| 208 | GraphDef subgraph; |
| 209 | std::unordered_map<const NodeDef*, const NodeDef*> subgraph_mapping; |
| 210 | |
| 211 | for (const auto& component_node : component_nodes) { |
| 212 | NodeDef* node = subgraph.add_node(); |
| 213 | *node = *component_node; |
| 214 | subgraph_mapping[node] = component_node; |
| 215 | if (IsNextIteration(*node)) { |
| 216 | CHECK_EQ(1, node->input_size()); |
| 217 | next_iter_nodes.emplace_back(node, node->input(0)); |
| 218 | } |
| 219 | } |
| 220 | if (next_iter_nodes.size() == 1) { |
| 221 | for (const auto& component_node : component_nodes) { |
| 222 | (*loops)[component_node].push_back(loop_id); |
| 223 | } |
| 224 | ++loop_id; |
| 225 | } else { |
| 226 | for (int i = 0; i < next_iter_nodes.size(); ++i) { |
| 227 | for (int j = 0; j < next_iter_nodes.size(); ++j) { |
| 228 | next_iter_nodes[j].first->clear_input(); |
| 229 | if (i == j) { |
| 230 | *next_iter_nodes[j].first->add_input() = next_iter_nodes[j].second; |
| 231 | } |
| 232 | } |
| 233 | int num_components = 0; |
| 234 | std::unordered_map<const NodeDef*, int> components; |
| 235 | StronglyConnectedComponents(subgraph, &components, &num_components); |
| 236 | CHECK_GE(num_components, 1); |
| 237 | for (const auto it : components) { |
| 238 | int id = it.second; |
| 239 | if (id < 0) { |
| 240 | continue; |
| 241 | } |