================== SubGraphExtractor =================*/
| 199 | |
| 200 | /* ================== SubGraphExtractor =================*/ |
| 201 | std::vector<GraphPartition> SubGraphExtractor::extract( |
| 202 | const SymbolVarArray& endpoint_vars) const { |
| 203 | ThinHashMap<OperatorNodeBase*, std::pair<OperatorNodeBase*, int>> parent; |
| 204 | thin_function<OperatorNodeBase*(OperatorNodeBase*)> union_find; |
| 205 | union_find = [&parent, &union_find](OperatorNodeBase* o) { |
| 206 | if (parent[o].first == o) |
| 207 | return o; |
| 208 | else { |
| 209 | auto p = union_find(parent[o].first); |
| 210 | parent[o].first = p; |
| 211 | return p; |
| 212 | } |
| 213 | }; |
| 214 | auto union_merge = [&parent, &union_find]( |
| 215 | OperatorNodeBase* x, OperatorNodeBase* y) { |
| 216 | auto root_x = union_find(x), root_y = union_find(y); |
| 217 | if (root_x != root_y) { |
| 218 | OperatorNodeBase *large, *small; |
| 219 | if (parent[root_x].second < parent[root_y].second) { |
| 220 | small = root_x, large = root_y; |
| 221 | } else { |
| 222 | small = root_y, large = root_x; |
| 223 | } |
| 224 | parent[small].first = large; |
| 225 | if (parent[large].second == parent[small].second) { |
| 226 | parent[large].second += 1; |
| 227 | } |
| 228 | } |
| 229 | }; |
| 230 | |
| 231 | std::vector<OperatorNodeBase*> topo; |
| 232 | auto cb = [this, &parent, &union_merge, &topo](OperatorNodeBase* opr) { |
| 233 | topo.push_back(opr); |
| 234 | if (m_opr_list.count(opr->dyn_typeinfo()) == 0) |
| 235 | return; |
| 236 | auto find = parent.find(opr); |
| 237 | if (find == parent.end()) { |
| 238 | parent.insert(std::make_pair(opr, std::make_pair(opr, 0))); |
| 239 | } |
| 240 | for (auto&& i : opr->input()) { |
| 241 | auto&& o = i->owner_opr(); |
| 242 | if (m_opr_list.count(o->dyn_typeinfo()) == 0) |
| 243 | continue; |
| 244 | union_merge(opr, o); |
| 245 | } |
| 246 | }; |
| 247 | cg::DepOprIter iter{cb}; |
| 248 | for (const auto& v : endpoint_vars) |
| 249 | iter.add(v.node()->owner_opr()); |
| 250 | |
| 251 | std::vector<GraphPartition> partitions; |
| 252 | partitions.reserve(topo.size()); |
| 253 | ThinHashMap<OperatorNodeBase*, GraphPartition*> roots; |
| 254 | /// backward pass |
| 255 | for (const auto& opr : reverse_adaptor(topo)) { |
| 256 | if (m_opr_list.count(opr->dyn_typeinfo()) > 0) { |
| 257 | auto root = union_find(opr); |
| 258 | auto find = roots.find(root); |