| 48 | } |
| 49 | |
| 50 | void DealWithNode(Graph* g, Node* n, |
| 51 | const std::vector<bool>& is_var_relate, |
| 52 | std::queue<Node*>& queue, |
| 53 | std::unordered_set<Node*>& has_visit_node, |
| 54 | std::unordered_set<Node*>& boundary_node_set) { |
| 55 | if (!n->IsConstant()) { |
| 56 | for (auto edge : n->out_edges()) { |
| 57 | Node* dst = edge->dst(); |
| 58 | if (is_var_relate[dst->id()]) { |
| 59 | boundary_node_set.emplace(n); |
| 60 | } else { |
| 61 | queue.emplace(dst); |
| 62 | } |
| 63 | } |
| 64 | return; |
| 65 | } |
| 66 | |
| 67 | // classify the output edges of const op |
| 68 | bool is_connect_to_marked_graph = false; |
| 69 | std::unordered_set<const Edge*> data_edge_to_unmarked_graph; |
| 70 | for (auto edge : n->out_edges()) { |
| 71 | // skip control edge |
| 72 | if (edge->IsControlEdge()) { |
| 73 | continue; |
| 74 | } |
| 75 | |
| 76 | Node* dst = edge->dst(); |
| 77 | if (is_var_relate[dst->id()]) { |
| 78 | is_connect_to_marked_graph = true; |
| 79 | } else { |
| 80 | data_edge_to_unmarked_graph.emplace(edge); |
| 81 | queue.emplace(dst); |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | // const op connects to marked graph and unmarked graph at the same time, |
| 86 | // duplicate a new const op node to avoid memcpy bewteen cpu and gpu. |
| 87 | if (is_connect_to_marked_graph && !data_edge_to_unmarked_graph.empty()) { |
| 88 | Node* new_node = g->CopyNode(n); |
| 89 | std::string new_name(n->name() + "_duplicate"); |
| 90 | new_node->set_name(new_name); |
| 91 | has_visit_node.emplace(new_node); |
| 92 | |
| 93 | for (auto edge : data_edge_to_unmarked_graph) { |
| 94 | Node* dst_node = edge->dst(); |
| 95 | g->AddEdge(new_node, edge->src_output(), dst_node, |
| 96 | edge->dst_input()); |
| 97 | g->RemoveEdge(edge); |
| 98 | } |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | } // namespace |
| 103 |
no test coverage detected