| 40 | } |
| 41 | |
| 42 | const CompNode::UnorderedMap<OprNodeArray>* MemoryOptimizerHelper::split_into_cn2oprseq( |
| 43 | const OprNodeArray& oprseq, const SubGraphConfig& config) { |
| 44 | auto BAD_OPR_FLAG = config.bad_opr_flag; |
| 45 | auto BAD_VAR_FLAG = config.bad_var_flag; |
| 46 | |
| 47 | m_cn2oprseq.clear(); |
| 48 | m_var_memsize.clear(); |
| 49 | for (auto i : oprseq) { |
| 50 | if (i->node_prop().contain(BAD_OPR_FLAG)) { |
| 51 | continue; |
| 52 | } |
| 53 | |
| 54 | auto cn = i->output(0)->comp_node(); |
| 55 | auto cn_loc = cn.locator(); |
| 56 | |
| 57 | bool have_static_shape_out = false, multi_out_cn = false, |
| 58 | different_device_inp = false; |
| 59 | |
| 60 | // check whether there are inputs from different device (if so, this opr |
| 61 | // should never be duplciated) |
| 62 | for (auto j : i->input()) { |
| 63 | auto loc = j->comp_node().locator(); |
| 64 | if (loc.type != cn_loc.type || loc.device != cn_loc.device) { |
| 65 | different_device_inp = true; |
| 66 | break; |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | if (different_device_inp) { |
| 71 | continue; |
| 72 | } |
| 73 | |
| 74 | // check same comp node and known shape for outputs |
| 75 | for (auto j : i->output()) { |
| 76 | if (j->comp_node() != cn) { |
| 77 | multi_out_cn = true; |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | if (multi_out_cn) { |
| 82 | continue; |
| 83 | } |
| 84 | |
| 85 | auto&& infer_mgr = m_owner_graph->static_infer_manager(); |
| 86 | for (auto j : i->output()) { |
| 87 | if (!j->contain_flag(BAD_VAR_FLAG)) { |
| 88 | // omit infer type check |
| 89 | // inferred shape will be used as-is |
| 90 | if (auto shape = infer_mgr.infer_shape_fallible(j)) { |
| 91 | have_static_shape_out = true; |
| 92 | m_var_memsize[j] = j->dtype().size(shape->total_nr_elems()); |
| 93 | } |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | if (have_static_shape_out) { |
| 98 | m_cn2oprseq[cn].push_back(i); |
| 99 | } |
no test coverage detected