| 2263 | }; |
| 2264 | |
| 2265 | void ShuffleShuffleRemovePass::Impl::detect_shuffle_operations() { |
| 2266 | auto rewriter = m_opt_state.graph().make_rewriter(); |
| 2267 | auto uniq_reader_check = UniqReaderCheck{m_opt_state.graph()}; |
| 2268 | auto try_reshape_shuffle = [&rewriter, &uniq_reader_check](OperatorNodeBase* opr) { |
| 2269 | // check shuffle |
| 2270 | auto shuffle = try_cast_as_op<opr::Dimshuffle>(opr); |
| 2271 | if (shuffle == nullptr) |
| 2272 | return false; |
| 2273 | auto&& param = shuffle->param(); |
| 2274 | if (param.pattern_len != 5) |
| 2275 | return false; |
| 2276 | bool is_nchw2nchw4 = param.pattern[0] == 0 && param.pattern[1] == 1 && |
| 2277 | param.pattern[2] == 3 && param.pattern[3] == 4 && |
| 2278 | param.pattern[4] == 2 && opr->output(0)->shape()[4] == 4; |
| 2279 | if (!is_nchw2nchw4) |
| 2280 | return false; |
| 2281 | if (!uniq_reader_check(shuffle->input(0))) |
| 2282 | return false; |
| 2283 | |
| 2284 | // check reshape |
| 2285 | auto reshape = try_cast_as_op<opr::Reshape>(opr->input(0)->owner_opr()); |
| 2286 | if (reshape == nullptr) |
| 2287 | return false; |
| 2288 | auto inp_var = rewriter.get_var(reshape->input(0)); |
| 2289 | auto abstract_shuffle = AbstractShuffleOpr::make( |
| 2290 | inp_var, ReformatKey{TensorFormats::NCHW, TensorFormats::NCHWc4}); |
| 2291 | rewriter.replace_var( |
| 2292 | opr->output(0), abstract_shuffle.node(), |
| 2293 | mgb_cstr_log("replace reformat(nchw -> nchw4) to " |
| 2294 | "AbstractShuffleOpr(nchw -> nchw4).")); |
| 2295 | return true; |
| 2296 | }; |
| 2297 | |
| 2298 | auto try_reshape_shuffle_reshape = [&rewriter, |
| 2299 | &uniq_reader_check](OperatorNodeBase* opr) { |
| 2300 | // check reshape |
| 2301 | auto reshape1 = try_cast_as_op<opr::Reshape>(opr); |
| 2302 | if (reshape1 == nullptr) |
| 2303 | return false; |
| 2304 | if (!uniq_reader_check(reshape1->input(0))) |
| 2305 | return false; |
| 2306 | |
| 2307 | // check shuffle |
| 2308 | auto shuffle = try_cast_as_op<opr::Dimshuffle>(opr->input(0)->owner_opr()); |
| 2309 | if (shuffle == nullptr) |
| 2310 | return false; |
| 2311 | auto&& param = shuffle->param(); |
| 2312 | if (param.pattern_len != 6) |
| 2313 | return false; |
| 2314 | bool is_nchw42nchw32 = param.pattern[0] == 0 && param.pattern[1] == 1 && |
| 2315 | param.pattern[2] == 3 && param.pattern[3] == 4 && |
| 2316 | param.pattern[4] == 2 && param.pattern[5] == 5 && |
| 2317 | shuffle->input(0)->shape()[5] == 4 && |
| 2318 | shuffle->input(0)->shape()[2] == 8; |
| 2319 | bool is_nchw322nchw4 = param.pattern[0] == 0 && param.pattern[1] == 1 && |
| 2320 | param.pattern[2] == 4 && param.pattern[3] == 2 && |
| 2321 | param.pattern[4] == 3 && param.pattern[5] == 5 && |
| 2322 | shuffle->input(0)->shape()[4] == 8 && |
nothing calls this directly
no test coverage detected