| 32 | } |
| 33 | |
| 34 | void FoldingReduceMeanPass::apply(OptState& opt) const { |
| 35 | MIDOUT_B("FoldingReduceMeanPass::apply"); |
| 36 | FindNext find_tool(opt); |
| 37 | |
| 38 | auto rewriter = opt.graph().make_rewriter(); |
| 39 | |
| 40 | /** |
| 41 | * reshape+---------->reduce(axis, sum)+--------->axis_remove+----------->true_div |
| 42 | * | ^ |
| 43 | * | | |
| 44 | * +--------------> get_var_shape(axis)+------------>type_cvt(fp32)+-------+ |
| 45 | * || |
| 46 | * || |
| 47 | * \/ |
| 48 | * reshape+-------->reduce(axis, mean)+--------->axis_remove |
| 49 | * |
| 50 | * |
| 51 | **/ |
| 52 | auto try_fuse_reduce_mean = [&rewriter, &find_tool](OperatorNodeBase* opr) { |
| 53 | ThinHashSet<OperatorNodeBase*> opr_set; |
| 54 | ThinHashSet<OperatorNodeBase*> reader_set; |
| 55 | MGB_MARK_USED_VAR(rewriter); |
| 56 | // check true_div |
| 57 | auto elemwise = try_cast_as_op<opr::Elemwise>(opr); |
| 58 | CHECK_OR_RETURN(elemwise); |
| 59 | auto mode_ok = elemwise->param().mode == opr::Elemwise::Mode::TRUE_DIV; |
| 60 | CHECK_OR_RETURN(mode_ok); |
| 61 | |
| 62 | auto input0 = elemwise->input(0)->owner_opr(); |
| 63 | auto remove_axis = input0->try_cast_final<opr::AxisAddRemove>(); |
| 64 | auto reduce = input0->try_cast_final<opr::Reduce>(); |
| 65 | if (remove_axis) { |
| 66 | reduce = remove_axis->input(0)->owner_opr()->try_cast_final<opr::Reduce>(); |
| 67 | } |
| 68 | CHECK_OR_RETURN(reduce); |
| 69 | |
| 70 | bool reduce_sum = reduce->param().mode == opr::Reduce::Param::Mode::SUM; |
| 71 | CHECK_OR_RETURN(reduce_sum); |
| 72 | |
| 73 | auto input1 = elemwise->input(1)->owner_opr(); |
| 74 | auto typecvt = input1->try_cast_final<opr::TypeCvt>(); |
| 75 | CHECK_OR_RETURN(typecvt); |
| 76 | auto is_typecvt_f32 = typecvt->param().enumv() == DTypeEnum::Float32; |
| 77 | CHECK_OR_RETURN(is_typecvt_f32); |
| 78 | |
| 79 | auto get_var_shape = |
| 80 | typecvt->input(0)->owner_opr()->try_cast_final<opr::GetVarShape>(); |
| 81 | CHECK_OR_RETURN(get_var_shape); |
| 82 | |
| 83 | bool same_parent = |
| 84 | get_var_shape->input(0)->owner_opr() == reduce->input(0)->owner_opr(); |
| 85 | CHECK_OR_RETURN(same_parent); |
| 86 | |
| 87 | CHECK_OR_RETURN( |
| 88 | find_tool.used_count(get_var_shape->input(0)->owner_opr()) == 2); |
| 89 | |
| 90 | bool same_axis = get_var_shape->param().axis == reduce->param().axis; |
| 91 | CHECK_OR_RETURN(same_axis); |
nothing calls this directly
no test coverage detected