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

Function kern_matmul_nchw44

dnn/src/fallback/convolution/algos.cpp:590–660  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

588/* ===================== Matrix mul nchw44 algo ===================== */
589namespace {
590void kern_matmul_nchw44(const NCBKernParam& param) {
591 bool is_xcorr = !param.filter_meta.should_flip;
592 UNPACK_CONV_F32_NCB_KERN_SIZES(param);
593 auto bundle = get_bundle(param);
594 bundle.set(param.workspace_ptr);
595 bool is1X1 = (FH == 1 && FW == 1 && SH == 1 && SW == 1 && PH == 0 && PW == 0);
596
597 typedef void (*Func1)(const float*, float*, int, int, int, int, int, int, int);
598 typedef void (*Func2)(
599 const float*, float*, int, int, int, int, int, int, int, int, int, int,
600 int);
601 Func1 f1 = nullptr;
602 Func2 f2 = nullptr;
603 if (is_xcorr) {
604 f1 = col2img_nchw44<true>;
605 f2 = col2img_stride_padding_nchw44<true>;
606 } else {
607 f1 = col2img_nchw44<false>;
608 f2 = col2img_stride_padding_nchw44<false>;
609 }
610 float* filter = const_cast<float*>(param.filter<float>());
611 TensorND A_src, A_dst;
612 {
613 A_src.layout = TensorLayout(
614 {IC / 4 * FH * FW, OC / 4, 4, 4},
615 {
616 static_cast<std::ptrdiff_t>(16),
617 static_cast<std::ptrdiff_t>(IC * FH * FW * 4),
618 static_cast<std::ptrdiff_t>(1),
619 static_cast<std::ptrdiff_t>(4),
620 },
621 param.filter_type);
622 A_src.reset_ptr(static_cast<void*>(filter));
623 A_dst.layout =
624 TensorLayout({IC / 4 * FH * FW, OC / 4, 4, 4}, param.filter_type);
625 A_dst.reset_ptr(static_cast<void*>(bundle.get(2)));
626 // TODO Should be removed once armv8 convolution support transpose.
627 get_relayout_opr()->exec(A_src, A_dst, inplace_cpu_handle().get());
628 }
629 TensorND B_, C_;
630 for (size_t n = 0; n < N; ++n) {
631 float *C_src, *C_dst;
632 float* diff = const_cast<float*>(param.diff<float>() + n * param.inp_bs);
633 float* grad = param.grad<float>() + n * param.out_bs;
634 if (is1X1) {
635 C_src = grad;
636 } else {
637 C_src = static_cast<float*>(bundle.get(0));
638 }
639 {
640 B_.layout = TensorLayout({OC / 4, IH * IW, 4}, param.diff_type);
641 B_.reset_ptr(static_cast<void*>(diff));
642 C_.layout = TensorLayout({IC / 4 * FH * FW, IH * IW, 4}, param.grad_type);
643 C_.reset_ptr(C_src);
644 Workspace workspace(
645 static_cast<dt_byte*>(bundle.get(1)), bundle.get_size(1));
646 auto matmul_opr = get_matmul_opr(param);
647 matmul_opr->exec(A_dst, B_, C_, workspace);

Callers

nothing calls this directly

Calls 10

TensorLayoutClass · 0.85
get_relayout_oprFunction · 0.85
get_matmul_oprFunction · 0.85
reset_ptrMethod · 0.80
get_bundleFunction · 0.70
setMethod · 0.45
getMethod · 0.45
execMethod · 0.45
get_sizeMethod · 0.45
sizeMethod · 0.45

Tested by

no test coverage detected