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

Method apply

src/gopt/impl/folding_conv_dimshuffle.cpp:35–416  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

33}
34
35void FoldingConvBiasDimshufflePass::apply(OptState& opt) const {
36 MIDOUT_B("FoldingConvBiasDimshufflePass::apply");
37 using DepType = cg::OperatorNodeProp::DepType;
38 ThinHashMap<OperatorNodeBase*, SmallVector<std::pair<OperatorNodeBase*, DepType>>>
39 readers;
40 static const ThinHashSet<Typeinfo*> opr_type_list = {
41 opr::TypeCvt::typeinfo(), opr::Dimshuffle::typeinfo(),
42 opr::Reshape::typeinfo(), opr::ConvBias::typeinfo()};
43 opt.graph().iter([&readers](OperatorNodeBase* opr) {
44 for (auto&& i : opr->node_prop().dep_map()) {
45 if (opr_type_list.count(i.first->owner_opr()->dyn_typeinfo())) {
46 readers[i.first->owner_opr()].emplace_back(opr, i.second);
47 }
48 }
49 });
50
51 auto rewriter = opt.graph().make_rewriter();
52 auto try_conv_dimshuffle_reshape_typecvt = [&rewriter,
53 &readers](OperatorNodeBase* opr) {
54 ThinHashSet<OperatorNodeBase*> opr_set;
55 ThinHashSet<OperatorNodeBase*> reader_set;
56 // check typecvt
57 auto typecvt = try_cast_as_op<opr::TypeCvt>(opr);
58 if (typecvt == nullptr)
59 return false;
60 auto inp_dtype = typecvt->input(0)->dtype(),
61 out_dtype = typecvt->output(0)->dtype();
62 bool is_s82f32 = inp_dtype.enumv() == DTypeEnum::QuantizedS8 &&
63 out_dtype.enumv() == DTypeEnum::Float32;
64 if (!is_s82f32)
65 return false;
66 opr_set.insert(opr);
67
68 // check reshape
69 auto reshape = try_cast_as_op<opr::Reshape>(typecvt->input(0)->owner_opr());
70 if (reshape == nullptr)
71 return false;
72 opr_set.insert(reshape);
73 for (auto&& i : readers[reshape]) {
74 if (i.second & DepType::DEV_VALUE) {
75 reader_set.insert(i.first);
76 }
77 }
78
79 // check shuffle
80 auto shuffle = try_cast_as_op<opr::Dimshuffle>(reshape->input(0)->owner_opr());
81 if (shuffle == nullptr)
82 return false;
83 auto&& param = shuffle->param();
84 if (param.pattern_len != 5)
85 return false;
86 bool is_nchw42nchw = param.pattern[0] == 0 && param.pattern[1] == 1 &&
87 param.pattern[2] == 4 && param.pattern[3] == 2 &&
88 param.pattern[4] == 3 &&
89 shuffle->input(0)->shape()[4] == 4;
90 if (!is_nchw42nchw)
91 return false;
92 opr_set.insert(shuffle);

Callers

nothing calls this directly

Calls 15

emplace_backMethod · 0.80
make_rewriterMethod · 0.80
apply_inplaceMethod · 0.80
makeFunction · 0.50
iterMethod · 0.45
graphMethod · 0.45
countMethod · 0.45
owner_oprMethod · 0.45
dtypeMethod · 0.45
inputMethod · 0.45
outputMethod · 0.45
enumvMethod · 0.45

Tested by

no test coverage detected