Find all CollectiveReduce nodes and the existing data dependencies between them.
| 24 | // Find all CollectiveReduce nodes and the existing data dependencies between |
| 25 | // them. |
| 26 | Status DiscoverDataDependencies( |
| 27 | const Graph* graph, std::vector<Node*>* collective_nodes, |
| 28 | std::vector<int32>* instance_keys, |
| 29 | absl::flat_hash_map<Node*, absl::flat_hash_set<int32>>* data_dependencies) { |
| 30 | Status s; |
| 31 | // Algorithm: do Reverse DFS starting at sink. `node_leave` is called when |
| 32 | // all parents of `node` have been visited. At that point, |
| 33 | // `data_dependencies[node]` is a list containing `instance_key` of every |
| 34 | // `CollectiveReduce` on which `node` has a data dependency. |
| 35 | // For this node's children, add all these instance keys. Also, if this node |
| 36 | // is collective, add as a dependency for the children. |
| 37 | auto node_leave = [collective_nodes, instance_keys, data_dependencies, |
| 38 | &s](Node* node) { |
| 39 | int32 instance_key; |
| 40 | bool enter_node = |
| 41 | node->IsCollective() && node->type_string() == "CollectiveReduce"; |
| 42 | if (enter_node) { |
| 43 | Status get_attr_status = |
| 44 | GetNodeAttr(node->attrs(), "instance_key", &instance_key); |
| 45 | s.Update(get_attr_status); |
| 46 | collective_nodes->push_back(node); |
| 47 | instance_keys->push_back(instance_key); |
| 48 | VLOG(2) << "collective node " << node->DebugString(); |
| 49 | } |
| 50 | // Avoid reference invalidation of `node_deps`. |
| 51 | data_dependencies->reserve(data_dependencies->size() + 1 + |
| 52 | node->out_edges().size()); |
| 53 | const auto& node_deps = (*data_dependencies)[node]; |
| 54 | for (const Edge* out_edge : node->out_edges()) { |
| 55 | auto& child_deps = (*data_dependencies)[out_edge->dst()]; |
| 56 | child_deps.insert(node_deps.begin(), node_deps.end()); |
| 57 | if (enter_node && s.ok()) { |
| 58 | child_deps.insert(instance_key); |
| 59 | } |
| 60 | } |
| 61 | }; |
| 62 | ReverseDFS(*graph, nullptr, node_leave); |
| 63 | return s; |
| 64 | } |
| 65 | |
| 66 | // Given a list of `collective_nodes` and `data_dependencies` between the |
| 67 | // collective nodes, create control dependencies between concurrent collectives |
no test coverage detected