Given a list of `collective_nodes` and `data_dependencies` between the collective nodes, create control dependencies between concurrent collectives and store in `dependency_edges`. If there exists an edge a -> b then `dependency_edges[a]` contains `b`
| 68 | // and store in `dependency_edges`. |
| 69 | // If there exists an edge a -> b then `dependency_edges[a]` contains `b` |
| 70 | Status CreateControlDependencies( |
| 71 | const std::vector<Node*>& collective_nodes, |
| 72 | const std::vector<int32>& instance_keys, |
| 73 | absl::flat_hash_map<Node*, absl::flat_hash_set<int32>>* data_dependencies, |
| 74 | absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>>* dependency_edges) { |
| 75 | // If there exists some path a -> ... -> b then `all_paths[a]` contains `b` |
| 76 | absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>> all_paths; |
| 77 | for (int i = 0; i < collective_nodes.size() - 1; i++) { |
| 78 | if (!collective_nodes[i]->IsCollective() || |
| 79 | collective_nodes[i]->type_string() != "CollectiveReduce") { |
| 80 | return errors::Internal("Unexpected node ", |
| 81 | collective_nodes[i]->DebugString()); |
| 82 | } |
| 83 | const auto& deps_i = (*data_dependencies)[collective_nodes[i]]; |
| 84 | for (int j = i + 1; j < collective_nodes.size(); j++) { |
| 85 | if (collective_nodes[i]->requested_device() != |
| 86 | collective_nodes[j]->requested_device()) { |
| 87 | continue; |
| 88 | } |
| 89 | if (instance_keys[i] == instance_keys[j]) { |
| 90 | return errors::Internal("Unexpected same instance_key ", |
| 91 | instance_keys[i], |
| 92 | " on 2 nodes with the same device ", |
| 93 | collective_nodes[i]->requested_device()); |
| 94 | } |
| 95 | const auto& deps_j = (*data_dependencies)[collective_nodes[j]]; |
| 96 | if (deps_i.find(instance_keys[j]) == deps_i.end() && |
| 97 | deps_j.find(instance_keys[i]) == deps_j.end()) { |
| 98 | int src_idx = instance_keys[i] > instance_keys[j] ? i : j; |
| 99 | int dst_idx = instance_keys[i] > instance_keys[j] ? j : i; |
| 100 | Node* src_node = collective_nodes[src_idx]; |
| 101 | Node* dst_node = collective_nodes[dst_idx]; |
| 102 | VLOG(1) << "Adding control dependency from node " << src_node->name() |
| 103 | << " instance " << instance_keys[src_idx] << " to node " |
| 104 | << dst_node->name() << " instance " << instance_keys[dst_idx]; |
| 105 | (*dependency_edges)[src_node].insert(dst_node); |
| 106 | auto& src_paths = all_paths[src_node]; |
| 107 | src_paths.insert(dst_node); |
| 108 | for (Node* downstream_node : all_paths[dst_node]) { |
| 109 | src_paths.insert(downstream_node); |
| 110 | } |
| 111 | } |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | // Prune dependency edges so that if there are edges a -> b, b -> c, and a -> |
| 116 | // c, then remove a -> c. This dependency would be handled naturally during |
| 117 | // op scheduling. |
| 118 | for (int i = 0; i < collective_nodes.size(); ++i) { |
| 119 | Node* node = collective_nodes[i]; |
| 120 | auto& neighbor_set = (*dependency_edges)[node]; |
| 121 | std::vector<Node*> neighbor_list(neighbor_set.begin(), neighbor_set.end()); |
| 122 | // For all n1, n2 in `neighbor_list` if there is a path from n1 -> n2 then |
| 123 | // eliminate n2 from `neighbor_set` and `neighbor_list`. We remove from |
| 124 | // `neighbor_list` by replacing with a `nullptr`, hence the `nullptr` checks |
| 125 | // below. |
| 126 | for (int j = 0; j < neighbor_list.size(); ++j) { |
| 127 | Node* n1 = neighbor_list[j]; |
no test coverage detected