| 271 | } |
| 272 | |
| 273 | void NetworkImplDft::replace_src_discrete_input_opr_pass() { |
| 274 | mgb::ThinHashMap<mgb::SymbolVar, mgb::SymbolVar> out_var_map; |
| 275 | |
| 276 | auto dest_with_extra_deps = |
| 277 | get_dest_vars_with_extra_deps(m_load_result.output_var_list); |
| 278 | gopt::SubGraph graph{dest_with_extra_deps}; |
| 279 | auto rewriter = graph.make_rewriter(); |
| 280 | |
| 281 | auto on_opr = [&](cg::OperatorNodeBase* opr) { |
| 282 | bool replace_output = false; |
| 283 | for (auto inp : opr->input()) { |
| 284 | if ((inp->owner_opr()->same_type<mgb::opr::Host2DeviceCopy>() || |
| 285 | inp->owner_opr()->same_type<mgb::opr::VolatileSharedDeviceTensor>()) && |
| 286 | inp->name() == m_user_config->discrete_input_name) { |
| 287 | bool is_h2d = true; |
| 288 | if (inp->owner_opr()->same_type<mgb::opr::Host2DeviceCopy>()) { |
| 289 | is_h2d = true; |
| 290 | } else { |
| 291 | is_h2d = false; |
| 292 | } |
| 293 | |
| 294 | SymbolVarArray srcs; |
| 295 | if (is_h2d) { |
| 296 | auto h2d = inp->owner_opr(); |
| 297 | for (auto&& i : |
| 298 | get_discrete_tensors(m_user_config->discrete_input_name)) { |
| 299 | auto val = TensorHelper::implement(i) |
| 300 | ->cast_final_safe<TensorImplDft>() |
| 301 | .m_host_tensor; |
| 302 | LITE_ASSERT(val); |
| 303 | srcs.push_back(mgb::opr::Host2DeviceCopy::make( |
| 304 | *m_load_result.graph, val, h2d->config())); |
| 305 | } |
| 306 | } else { |
| 307 | auto volatiled = inp->owner_opr(); |
| 308 | for (auto&& i : |
| 309 | get_discrete_tensors(m_user_config->discrete_input_name)) { |
| 310 | auto val = TensorHelper::implement(i) |
| 311 | ->cast_final_safe<TensorImplDft>() |
| 312 | .m_dev_tensor; |
| 313 | LITE_ASSERT(val); |
| 314 | srcs.push_back(mgb::opr::VolatileSharedDeviceTensor::make( |
| 315 | *m_load_result.graph, val, volatiled->config())); |
| 316 | } |
| 317 | } |
| 318 | |
| 319 | if (opr->same_type<mgb::opr::WarpPerspective>()) { |
| 320 | auto& warp = opr->cast_final<mgb::opr::WarpPerspective>(); |
| 321 | SymbolVar new_out; |
| 322 | if (opr->input().size() == 3) { |
| 323 | new_out = mgb::opr::WarpPerspective::make( |
| 324 | srcs, warp.input(1), warp.input(2), warp.param(), |
| 325 | warp.config()); |
| 326 | } else { |
| 327 | LITE_ASSERT(opr->input().size() == 4); |
| 328 | new_out = mgb::opr::WarpPerspective::make( |
| 329 | srcs, warp.input(1), warp.input(2), warp.input(3), |
| 330 | warp.param(), warp.config()); |
nothing calls this directly
no test coverage detected