=================== LayoutTransformPass ======================*/
| 25 | |
| 26 | /* =================== LayoutTransformPass ======================*/ |
| 27 | void LayoutTransformPass::apply(OptState& opt) const { |
| 28 | MIDOUT_B("apply") |
| 29 | opt.set_var_replace_check_flag( |
| 30 | VarReplaceCheckFlag::CHECK_ALL ^ VarReplaceCheckFlag::CHECK_SHAPE); |
| 31 | SubGraphExtractor extractor(m_ctx->opr_list()); |
| 32 | auto partitions = extractor.extract(opt.graph().endpoint_vars()); |
| 33 | |
| 34 | using Solution = SolverBase::Solution; |
| 35 | using OprFormat = SolverBase::OprFormat; |
| 36 | Solution solution; |
| 37 | ThinHashSet<VarNode*> endpoint_vars; |
| 38 | for (auto&& partition : partitions) { |
| 39 | if (solution.empty()) { |
| 40 | solution = m_solver->solve(Problem(partition, *m_ctx)); |
| 41 | } else { |
| 42 | auto s = m_solver->solve(Problem(partition, *m_ctx)); |
| 43 | for (auto&& kv : s) |
| 44 | solution.insert({kv.first, kv.second}); |
| 45 | } |
| 46 | for (auto&& o : partition.output()) { |
| 47 | endpoint_vars.insert(o); |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | auto&& opr_configs = m_ctx->opr_configs(); |
| 52 | auto&& base_fmt = m_ctx->attribute().base_tensor_formats; |
| 53 | auto&& base_cfg_id = m_ctx->attribute().base_config_id; |
| 54 | auto&& reformat_attribute = m_ctx->attribute().reformat_attribute; |
| 55 | ThinHashMap<VarNode*, TensorFormats> var2fmts; |
| 56 | auto rewriter = opt.graph().make_rewriter(); |
| 57 | auto on_opr = [&opr_configs, &base_fmt, &base_cfg_id, &reformat_attribute, |
| 58 | &rewriter, &solution, &var2fmts, |
| 59 | &endpoint_vars](OperatorNodeBase* opr) { |
| 60 | auto it = solution.find(opr); |
| 61 | if (it != solution.end()) { |
| 62 | auto cfg_id = it->second; |
| 63 | auto find = opr_configs.find(opr->dyn_typeinfo()); |
| 64 | Maybe<OprTensorFormatsConfiguration> fmtcfg = None; |
| 65 | Maybe<OprTensorFormatsConfiguration> basecfg = None; |
| 66 | Maybe<OprFormat> opr_fmt = None; |
| 67 | if (find != opr_configs.end()) { |
| 68 | fmtcfg = (*find->second.at(cfg_id))(opr); |
| 69 | auto _ = OprTensorFormatsConfiguration::find_dispatcher_by_type_format( |
| 70 | opr->dyn_typeinfo(), base_cfg_id); |
| 71 | basecfg = (*_)(opr); |
| 72 | opr_fmt = fmtcfg.val().opr_format; |
| 73 | } else { |
| 74 | opr_fmt = |
| 75 | OprTensorFormatsConfiguration::safe_cast_to_opr_format(cfg_id); |
| 76 | } |
| 77 | VarNodeArray new_inp; |
| 78 | size_t nr_inps = opr->input().size(); |
| 79 | TensorFormats out_fmt; |
| 80 | if (fmtcfg.valid()) { |
| 81 | nr_inps = std::min(fmtcfg.val().input_tensor_formats.size(), nr_inps); |
| 82 | out_fmt = fmtcfg.val().output_tensor_formats[0]; |
| 83 | } else { |
| 84 | out_fmt = opr_format_to_tensor_formats(opr_fmt.val()); |
nothing calls this directly
no test coverage detected