| 82 | } |
| 83 | |
| 84 | void SubGraph::CompleteVariables(const PartitionOptions& opts) { |
| 85 | set<int> var_id_set; |
| 86 | int size = nodes_.size(); |
| 87 | for (int i = 0; i < size; i++) { |
| 88 | auto node = nodes_[i]; |
| 89 | if (!node->IsOp()) continue; |
| 90 | std::string loc = opts.node_to_loc(node); |
| 91 | for (const Edge* edge : node->in_edges()) { |
| 92 | Node* src = edge->src(); |
| 93 | if (!src->IsOp()) continue; |
| 94 | if ((src->IsVariable() || src->IsKvVarHandle ()) |
| 95 | && opts.node_to_loc(src) == loc) { |
| 96 | ADD_VAR_IF_NOT_EXIST(src); |
| 97 | } |
| 98 | |
| 99 | if (src->IsIdentity() && opts.node_to_loc(src) == loc) { |
| 100 | const Node *var = GetVarOfIdentity(opts, src); |
| 101 | if (var == nullptr) { |
| 102 | continue; |
| 103 | } |
| 104 | |
| 105 | ADD_VAR_IF_NOT_EXIST(src); |
| 106 | ADD_VAR_IF_NOT_EXIST(var); |
| 107 | } |
| 108 | } |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | GraphPartitionerBase::GraphPartitionerBase( |
| 113 | const PartitionOptions& popts, Graph* g, |
no test coverage detected