| 182 | } // namespace |
| 183 | |
| 184 | Status OrderCollectives(Graph* graph, GraphCollectiveOrder order_type) { |
| 185 | // `instance_keys[i]` corresponds to `collective_nodes[i]` |
| 186 | std::vector<Node*> collective_nodes; |
| 187 | std::vector<int32> instance_keys; |
| 188 | // node -> set of collectives on which node depends. |
| 189 | absl::flat_hash_map<Node*, absl::flat_hash_set<int32>> data_dependencies; |
| 190 | TF_RETURN_IF_ERROR(DiscoverDataDependencies( |
| 191 | graph, &collective_nodes, &instance_keys, &data_dependencies)); |
| 192 | |
| 193 | if (collective_nodes.empty()) return Status::OK(); |
| 194 | |
| 195 | absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>> dependency_edges; |
| 196 | // For all pairs of collective nodes n1 and n2 on the same device, if n1 does |
| 197 | // not depend on n2 and n2 does not depend on n1, then they are potentially |
| 198 | // concurrent. Create an arbitrary, deterministic ordering between them. |
| 199 | TF_RETURN_IF_ERROR(CreateControlDependencies( |
| 200 | collective_nodes, instance_keys, &data_dependencies, &dependency_edges)); |
| 201 | |
| 202 | return InsertControlDependencies(graph, order_type, dependency_edges); |
| 203 | } |
| 204 | |
| 205 | } // namespace tensorflow |