Copies a subgraph from `graph` to `output` by performing a reverse DFS starting at nodes in vector `stack`. `node_map` is a vector indexed by source node ID to dest nodes. Does not traverse into nodes in `node_map`, so by adding nodes to `node_map` before the traversal clients can cut the graph. If a frame is provided (frame != nullptr), then this functions will return an error if the traversal le
| 58 | // nodes only have one output. By adding the Switch node to `squash_src_outputs` |
| 59 | // we rewrite the src_output of the corresponding edge to be 0. |
| 60 | Status CopySubgraph(const Graph& graph, const WhileLoopFrame* frame, |
| 61 | std::vector<Node*> stack, |
| 62 | const std::vector<bool>& squash_src_outputs, |
| 63 | std::vector<Node*>* node_map, Graph* output) { |
| 64 | VLOG(3) << "Stack: " << NodesToString(stack); |
| 65 | std::vector<bool> visited(graph.num_node_ids(), false); |
| 66 | while (!stack.empty()) { |
| 67 | Node* n = stack.back(); |
| 68 | stack.pop_back(); |
| 69 | |
| 70 | VLOG(5) << "Copying node " << n->name(); |
| 71 | |
| 72 | if (visited[n->id()]) continue; |
| 73 | visited[n->id()] = true; |
| 74 | |
| 75 | // Sort "n->in_edges()" to make sure nodes are copied in a deterministic |
| 76 | // order. |
| 77 | std::vector<const Edge*> sorted_edges(n->in_edges().begin(), |
| 78 | n->in_edges().end()); |
| 79 | std::sort(sorted_edges.begin(), sorted_edges.end(), |
| 80 | [](const Edge* a, const Edge* b) { |
| 81 | int a_src_output = a->src_output(), |
| 82 | b_src_output = b->src_output(); |
| 83 | StringPiece a_name(a->src()->name()), b_name(b->src()->name()); |
| 84 | return std::tie(a_src_output, a_name) < |
| 85 | std::tie(b_src_output, b_name); |
| 86 | }); |
| 87 | for (const Edge* e : sorted_edges) { |
| 88 | Node* src = e->src(); |
| 89 | if (frame != nullptr && frame->nodes.find(src) == frame->nodes.end()) { |
| 90 | // We traversed out of the loop frame, without encountering a cut node. |
| 91 | return errors::Internal("Graph traversal of loop frame ", frame->name, |
| 92 | " escaped frame at ", src->name(), |
| 93 | " without encountering an argument node."); |
| 94 | } |
| 95 | if ((*node_map)[src->id()] == nullptr) { |
| 96 | (*node_map)[src->id()] = output->CopyNode(src); |
| 97 | stack.push_back(src); |
| 98 | } |
| 99 | Node* src_copy = (*node_map)[e->src()->id()]; |
| 100 | int src_output = squash_src_outputs[e->src()->id()] && !e->IsControlEdge() |
| 101 | ? 0 |
| 102 | : e->src_output(); |
| 103 | Node* dst_copy = (*node_map)[e->dst()->id()]; |
| 104 | output->AddEdge(src_copy, src_output, dst_copy, e->dst_input()); |
| 105 | } |
| 106 | } |
| 107 | return Status::OK(); |
| 108 | } |
| 109 | |
| 110 | StatusOr<Node*> BuildArgNode(Graph* graph, DataType type, int index) { |
| 111 | const char* const kArgOp = "_Arg"; |