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

Method replace_src_discrete_input_opr_pass

lite/src/mge/network_impl.cpp:273–365  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

271}
272
273void 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());

Callers

nothing calls this directly

Calls 15

make_rewriterMethod · 0.80
apply_inplaceMethod · 0.80
resizeMethod · 0.80
renameMethod · 0.80
makeFunction · 0.50
inputMethod · 0.45
owner_oprMethod · 0.45
nameMethod · 0.45
push_backMethod · 0.45
sizeMethod · 0.45
paramMethod · 0.45
replace_varMethod · 0.45

Tested by

no test coverage detected