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

Method deduce_exec_layout

dnn/src/common/relayout_format.cpp:438–721  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

436}
437
438void RelayoutFormat::deduce_exec_layout(
439 const TensorLayout& src, const TensorLayout& dst, TensorLayout& exec_workspace,
440 TensorLayout& exec_src, TensorLayout& exec_dst) {
441 check_layout_fwd(src, dst);
442 using Param = param::RelayoutFormat;
443 switch (param().mode) {
444 case Param::Mode::NCHW_NCHW88:
445 // nchw to nchw8c
446 {
447 exec_workspace = TensorLayout(
448 {src[0], round_up(src[1], 8_z), src[2], src[3]}, src.dtype,
449 src.format);
450 exec_src = exec_workspace
451 .reshape(
452 {src[0], div_ceil(src[1], 8_z), 8, src[2],
453 src[3]})
454 .dimshuffle({0, 1, 3, 4, 2});
455 exec_dst = dst;
456 }
457 break;
458 case Param::Mode::NCHW_NCHW4:
459 // nchw to nchw4
460 {
461 const size_t group = param().group;
462 const size_t icpg = src[1] / group;
463 exec_workspace = TensorLayout(
464 {src[0], group * round_up(icpg, 4_z), src[2], src[3]},
465 src.dtype, src.format);
466 exec_src = exec_workspace
467 .reshape(
468 {src[0], group * div_ceil(icpg, 4_z), 4,
469 src[2], src[3]})
470 .dimshuffle({0, 1, 3, 4, 2});
471 exec_dst = dst;
472 }
473 break;
474 case Param::Mode::NCHW_NCHW4_WEIGHT:
475 // nchw to nchw4_weight
476 {
477 if (src.ndim == 4) {
478 exec_workspace = TensorLayout(
479 {round_up(src[0], 4_z), round_up(src[1], 4_z), src[2],
480 src[3]},
481 src.dtype, src.format);
482 exec_src =
483 exec_workspace
484 .reshape(
485 {round_up(src[0], 4_z),
486 div_ceil(src[1], 4_z), 4, src[2], src[3]})
487 .dimshuffle({0, 1, 3, 4, 2});
488 exec_dst = dst;
489 } else if (src.ndim == 5) {
490 exec_workspace = TensorLayout(
491 {src[0], round_up(src[1], 4_z), round_up(src[2], 4_z),
492 src[3], src[4]},
493 src.dtype, src.format);
494 exec_src =
495 exec_workspace

Callers

nothing calls this directly

Calls 6

TensorLayoutClass · 0.85
round_upFunction · 0.85
div_ceilFunction · 0.85
dimshuffleMethod · 0.80
paramFunction · 0.50
reshapeMethod · 0.45

Tested by

no test coverage detected