Function GetGroupNodes

tensorflow/core/common_runtime/colocation_graph.cc:766–794

Given a node requiring placer inspection and its IOColocationGroups, computes `group_nodes`. group_nodes[i] contains the nodes that are members of colocation group i. These nodes are inputs or outputs of `node`. group_nodes[i][j] is a pair containing a node and whether this node has a resource input from `node`. Note: The same node can be added multiple times to the same group. The same node can b

// The same node can be added multiple times to the same group.
// The same node can be added to multiple groups.
Status GetGroupNodes(const IOColocationGroups& groups, const Node& node,
                     std::vector<std::vector<NodeAndBool>>* group_nodes) {
  group_nodes->reserve(groups.group_devices.size());
  for (int arg_idx = 0; arg_idx < groups.input_groups.size(); ++arg_idx) {
    const Node* src;
    TF_RETURN_IF_ERROR(node.input_node(arg_idx, &src));
    int group_id = groups.input_groups[arg_idx];
    (*group_nodes)[group_id].emplace_back(src, false);
  }

  for (const Edge* edge : node.out_edges()) {
    if (edge->IsControlEdge()) {
      continue;
    }

    int group_id = groups.output_groups[edge->src_output()];
    (*group_nodes)[group_id].emplace_back(
        edge->dst(), edge->dst()->input_type(edge->dst_input()) == DT_RESOURCE);
  }

  if (VLOG_IS_ON(2)) {
    VLOG(2) << "Colocated inputs/outputs of node: " << node.DebugString();
    for (const std::vector<NodeAndBool>& nodes : *group_nodes) {
      VLOG(2) << "\t[" << absl::StrJoin(NodeAndBoolToString(nodes), "\t\n")
              << "]";
    }
  }
  return Status::OK();
}

}  // namespace