This is an implementation of Tarjan's Strongly Connected Components DFS algorithm. Most of the hard work is done in the function StrongConnect, which is an iterative reimplementation of the recursive version described here: https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm The edges for the purpose of this algorithm are directed from input to op (the reverse of the
| 124 | // to op (the reverse of the declarations of the NodeDef, which |
| 125 | // contain in-edges) |
| 126 | void StronglyConnectedComponents( |
| 127 | const GraphDef& graph, std::unordered_map<const NodeDef*, int>* components, |
| 128 | int* num_components) { |
| 129 | std::stack<SCCNodeData*> stack; |
| 130 | std::unordered_map<string, SCCNodeData*> name_to_data; |
| 131 | std::vector<SCCNodeData> node_data_container; |
| 132 | node_data_container.reserve(graph.node_size()); |
| 133 | std::unordered_map<const NodeDef*, SCCNodeData*> node_to_data; |
| 134 | |
| 135 | for (const NodeDef& node : graph.node()) { |
| 136 | SCCNodeData node_data; |
| 137 | node_data.node = &node; |
| 138 | node_data_container.push_back(node_data); |
| 139 | name_to_data[node.name()] = &(*node_data_container.rbegin()); |
| 140 | node_to_data[&node] = &(*node_data_container.rbegin()); |
| 141 | } |
| 142 | |
| 143 | // Create a list of top-level parents (add them to object queue) |
| 144 | // Also create a mapping from nodes to their children. |
| 145 | // Inputs might not be present if called on a subgraph. |
| 146 | for (const NodeDef& node : graph.node()) { |
| 147 | for (const string& input : node.input()) { |
| 148 | auto it = name_to_data.find(NodeName(input)); |
| 149 | if (it != name_to_data.end()) { |
| 150 | it->second->children.push_back(node_to_data[&node]); |
| 151 | } |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | components->clear(); |
| 156 | *num_components = 0; |
| 157 | int index = 0; |
| 158 | for (auto& v : node_data_container) { |
| 159 | if (v.index == -1) { |
| 160 | // Node has not yet been visited. Start a DFS at v. |
| 161 | StrongConnect(&v, &stack, &index, components, num_components); |
| 162 | } |
| 163 | } |
| 164 | |
| 165 | std::vector<int> counts_per_component(*num_components, 0); |
| 166 | for (auto& component : *components) { |
| 167 | DCHECK(component.second >= 0); |
| 168 | DCHECK(component.second < *num_components); |
| 169 | counts_per_component[component.second]++; |
| 170 | } |
| 171 | bool has_single_element_component = false; |
| 172 | for (auto& component : *components) { |
| 173 | if (counts_per_component[component.second] == 1) { |
| 174 | component.second = -1; |
| 175 | (*num_components)--; |
| 176 | has_single_element_component = true; |
| 177 | } |
| 178 | } |
| 179 | if (has_single_element_component) { |
| 180 | (*num_components) += 1; |
| 181 | } |
| 182 | } |
| 183 |