| 2430 | } |
| 2431 | |
| 2432 | void ShuffleShuffleRemovePass::Impl::do_replace() { |
| 2433 | auto rewriter = m_opt_state.graph().make_rewriter(); |
| 2434 | auto uniq_reader_check = UniqReaderCheck{m_opt_state.graph()}; |
| 2435 | ThinHashSet<OperatorNodeBase*> writers; |
| 2436 | ThinHashSet<OperatorNodeBase*> root; |
| 2437 | ThinHashSet<VarNode*> trt_opr_inps; |
| 2438 | SmallVector<OperatorNodeBase*> topo_order; |
| 2439 | |
| 2440 | auto cb = [&topo_order, &trt_opr_inps](OperatorNodeBase* opr) { |
| 2441 | topo_order.push_back(opr); |
| 2442 | MGB_MARK_USED_VAR(trt_opr_inps); |
| 2443 | #if MGB_ENABLE_TENSOR_RT |
| 2444 | if (opr->same_type<opr::TensorRTOpr>()) { |
| 2445 | for (auto&& inp : opr->input()) |
| 2446 | trt_opr_inps.insert(inp); |
| 2447 | } |
| 2448 | #endif |
| 2449 | }; |
| 2450 | m_opt_state.graph().iter(cb); |
| 2451 | |
| 2452 | for (auto&& opr : reverse_adaptor(topo_order)) { |
| 2453 | if (opr->same_type<opr::TypeCvt>() || opr->same_type<AbstractShuffleOpr>()) { |
| 2454 | writers.insert(opr->input(0)->owner_opr()); |
| 2455 | if (writers.count(opr) > 0) { |
| 2456 | if (!uniq_reader_check(opr->output(0))) { |
| 2457 | root.insert(opr); |
| 2458 | } |
| 2459 | } else { |
| 2460 | root.insert(opr); |
| 2461 | } |
| 2462 | } |
| 2463 | } |
| 2464 | |
| 2465 | auto on_opr = [&rewriter, &uniq_reader_check, &trt_opr_inps, |
| 2466 | &root](OperatorNodeBase* opr) { |
| 2467 | MGB_MARK_USED_VAR(trt_opr_inps); |
| 2468 | bool cond_opr = |
| 2469 | opr->same_type<opr::TypeCvt>() || opr->same_type<AbstractShuffleOpr>(); |
| 2470 | if (cond_opr) { |
| 2471 | bool cond_endpoint = root.count(opr) > 0; |
| 2472 | if (!cond_endpoint) { |
| 2473 | return; |
| 2474 | } |
| 2475 | auto cur = opr; |
| 2476 | auto var = opr->output(0), inp_var = opr->input(0); |
| 2477 | bool force_folding_typecvt = false; |
| 2478 | bool first_shuffle = false; |
| 2479 | // initialize inp_format and out_format |
| 2480 | TensorFormats out_format = TensorFormats::NCHW, inp_format = out_format; |
| 2481 | megdnn::DType inp_dtype = cur->input(0)->dtype(), |
| 2482 | out_dtype = cur->output(0)->dtype(); |
| 2483 | SmallVector<megdnn::DType> out_dtype_vec; |
| 2484 | while (cond_opr) { |
| 2485 | if (cur->same_type<AbstractShuffleOpr>()) { |
| 2486 | auto shuffle = try_cast_as_op<AbstractShuffleOpr>(cur); |
| 2487 | inp_format = shuffle->key().input_format; |
| 2488 | if (!first_shuffle) { |
| 2489 | out_format = shuffle->key().output_format; |
nothing calls this directly
no test coverage detected