| 55 | } |
| 56 | |
| 57 | void MarkEmbeddingGraph(Graph* g, int num_streams) { |
| 58 | bool train_graph = false; |
| 59 | |
| 60 | // trained graph |
| 61 | if (!g->IsTrainingGraph()) { |
| 62 | return; |
| 63 | } |
| 64 | |
| 65 | //for (Node* n : g->nodes()) { |
| 66 | // if (n->type_string() == "IsVariableInitialized" && |
| 67 | // n->name() != "global_step/IsVariableInitialized") { |
| 68 | // return; |
| 69 | // } |
| 70 | //} |
| 71 | |
| 72 | std::unordered_map<std::string, Node*> name_to_node; |
| 73 | // User marked subgraph |
| 74 | for (Node* n : g->nodes()) { |
| 75 | name_to_node[n->name()] = n; |
| 76 | |
| 77 | if (n->assigned_device_name().find("device:GPU:") == std::string::npos || |
| 78 | n->def().attr().find("_stream_id") == n->def().attr().end()) { |
| 79 | continue; |
| 80 | } |
| 81 | |
| 82 | int stream_id = n->def().attr().at("_stream_id").i(); |
| 83 | std::string required_device = |
| 84 | GetDeviceNamePrefix(n->assigned_device_name()) + |
| 85 | std::to_string(stream_id); |
| 86 | if (n->assigned_device_name() != required_device) { |
| 87 | n->set_assigned_device_name(required_device); |
| 88 | } |
| 89 | } |
| 90 | |
| 91 | |
| 92 | // Colocate nodes |
| 93 | std::unordered_map<Node*, std::vector<Node*>> node_colocate_childs; |
| 94 | std::unordered_set<Node*> colocate_nodes; |
| 95 | for (Node* n : g->nodes()) { |
| 96 | std::vector<string> constraints; |
| 97 | bool has_constraints = GetColocationConstraints(n, &constraints); |
| 98 | if (has_constraints) { |
| 99 | for (const auto& constraint : constraints) { |
| 100 | node_colocate_childs[name_to_node[constraint.substr(5)]].emplace_back(n); |
| 101 | } |
| 102 | colocate_nodes.insert(n); |
| 103 | } |
| 104 | } |
| 105 | std::queue<Node*> q; |
| 106 | for (auto pair : node_colocate_childs) { |
| 107 | if (colocate_nodes.find(pair.first) == colocate_nodes.end()) { |
| 108 | q.push(pair.first); |
| 109 | } |
| 110 | } |
| 111 | while (!q.empty()) { |
| 112 | Node* curr = q.front(); |
| 113 | q.pop(); |
| 114 | for (auto& n : node_colocate_childs[curr]) { |
no test coverage detected