MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / update_out_specs

Method update_out_specs

src/core/impl/graph/cg_impl_partial.cpp:391–515  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

389}
390
391void ComputingGraphImpl::MultiPartCompiler::update_out_specs() {
392 // 1. Determine the overall computing sequence
393 // 2. Determine inter-dependencies between parts
394 // 3. Copy oprs of each part into a new graph, and use VarSinkOpr to pass
395 // values between graphs
396
397 auto&& opr_seq = *concat_and_prepare();
398 init_opr_trait_and_var_reader_type(opr_seq);
399
400 auto owner_graph = m_out_specs.at(0).at(0).first.node()->owner_graph();
401
402 // var replacement map for current part; cleared for at each part iter
403 ThinHashMap<VarNode*, VarNode*> cur_part_var_repl;
404
405 VarNodeArray tmp_new_inputs;
406 size_t cur_part;
407 serialization::OprShallowCopyContext copy_ctx;
408 // copy each opr and record in cur_part_var_repl
409 auto on_opr = [&](OperatorNodeBase* opr) {
410 VarNodeArray& new_inputs = tmp_new_inputs;
411 new_inputs.clear();
412 for (auto inp : opr->input()) {
413 new_inputs.emplace_back(cur_part_var_repl.at(inp));
414 }
415 OperatorNodeBase* new_opr = serialization::copy_opr_shallow(
416 *opr, new_inputs, opr->config(), copy_ctx);
417 auto&& new_dep_map = const_cast<OperatorNodeBase::NodeProp::DepMap&>(
418 new_opr->node_prop().dep_map());
419 // copy dep entries added by TopoSorter (and maybe others)
420 for (auto&& dep_entry : opr->node_prop().dep_map()) {
421 if (dep_entry.second == DepType::DEV_COMP_ORDER &&
422 !has_different_comp_node(opr, dep_entry.first->comp_node())) {
423 continue;
424 }
425 auto new_var = cur_part_var_repl.at(dep_entry.first);
426 auto iter = new_dep_map.find(new_var);
427 if (iter == new_dep_map.end()) {
428 iter = new_dep_map.insert({new_var, DepType{}}).first;
429 }
430 if (iter->second != dep_entry.second) {
431 mgb_assert((iter->second & dep_entry.second) == iter->second);
432 iter->second = dep_entry.second;
433 }
434 }
435
436 new_opr->node_prop().attribute().priority = m_opr_trait.at(opr).priority;
437
438 auto&& out0 = opr->output();
439 auto&& out1 = new_opr->output();
440 mgb_assert(out0.size() == out1.size());
441 for (size_t i = 0; i < out0.size(); ++i) {
442 cur_part_var_repl[out0[i]] = out1[i];
443 }
444 };
445 DepOprIter opr_iter{on_opr};
446
447 // map from outer graph var to the var used as VarSinkOpr input (i.e. when
448 // it is first produced in a part and recorded in the VarSinkOpr)

Callers

nothing calls this directly

Calls 15

emplace_backMethod · 0.80
backMethod · 0.80
safe_sink_oprMethod · 0.80
makeFunction · 0.70
owner_graphMethod · 0.45
nodeMethod · 0.45
atMethod · 0.45
clearMethod · 0.45
inputMethod · 0.45
comp_nodeMethod · 0.45
findMethod · 0.45

Tested by

no test coverage detected