Function InsertControlDependencies

tensorflow/core/graph/collective_order.cc:147–180

Insert control dependencies defined by `dependency_edges` in `graph`. If `order_type` is `kEdges`, insert explicit control edges, else if `order_type` is `kAttrs`, encode dependencies as an attribute on collective node.

// `order_type` is `kEdges`, insert explicit control edges, else if `order_type`
// is `kAttrs`, encode dependencies as an attribute on collective node.
Status InsertControlDependencies(
    Graph* graph, GraphCollectiveOrder order_type,
    const absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>>&
        dependency_edges) {
  if (order_type == GraphCollectiveOrder::kEdges) {
    for (const auto& pair : dependency_edges) {
      Node* src_node = pair.first;
      for (Node* dst_node : pair.second) {
        graph->AddControlEdge(src_node, dst_node);
      }
    }
  } else if (order_type == GraphCollectiveOrder::kAttrs) {
    // `wait_for` is the inverse of `dependency_edges`, i.e. `wait_for[node]`
    // contains the list of instance keys for which `node` must wait.
    absl::flat_hash_map<Node*, absl::flat_hash_set<int32>> wait_for;
    for (const auto& pair : dependency_edges) {
      int32 src_instance;
      TF_RETURN_IF_ERROR(
          GetNodeAttr(pair.first->attrs(), "instance_key", &src_instance));
      for (Node* dst_node : pair.second) {
        wait_for[dst_node].insert(src_instance);
      }
    }
    for (const auto& pair : wait_for) {
      std::vector<int32> wait_for_list(pair.second.begin(), pair.second.end());
      pair.first->ClearAttr("wait_for");
      pair.first->AddAttr("wait_for", wait_for_list);
    }
  } else {
    return errors::Internal("Unexpected GraphCollectiveOrder type ",
                            static_cast<int>(order_type));
  }
  return Status::OK();
}

}  // namespace