| 82 | |
| 83 | template <typename ftype, typename dtype, typename gtype> |
| 84 | void kern_matmul(const NCBKernParam& param) { |
| 85 | bool is_xcorr = !param.filter_meta.should_flip; |
| 86 | UNPACK_CONV_F32_NCB_KERN_SIZES(param); |
| 87 | auto bundle = get_bundle(param); |
| 88 | bundle.set(param.workspace_ptr); |
| 89 | bool is1X1 = (FH == 1 && FW == 1 && SH == 1 && SW == 1 && PH == 0 && PW == 0); |
| 90 | |
| 91 | typedef void (*Func1)(const gtype*, gtype*, int, int, int, int, int, int, int); |
| 92 | typedef void (*Func2)( |
| 93 | const gtype*, gtype*, int, int, int, int, int, int, int, int, int, int, |
| 94 | int); |
| 95 | Func1 f1 = nullptr; |
| 96 | Func2 f2 = nullptr; |
| 97 | if (is_xcorr) { |
| 98 | f1 = col2img<true>; |
| 99 | f2 = col2img_stride_padding<true>; |
| 100 | } else { |
| 101 | f1 = col2img<false>; |
| 102 | f2 = col2img_stride_padding<false>; |
| 103 | } |
| 104 | ftype* filter = const_cast<ftype*>(param.filter<ftype>()); |
| 105 | TensorND A_src, A_dst; |
| 106 | { |
| 107 | A_src.layout = TensorLayout( |
| 108 | {IC * FH * FW, OC}, |
| 109 | {static_cast<std::ptrdiff_t>(1), |
| 110 | static_cast<std::ptrdiff_t>(IC * FH * FW)}, |
| 111 | param.filter_type); |
| 112 | A_src.reset_ptr(static_cast<void*>(filter)); |
| 113 | A_dst.layout = TensorLayout({IC * FH * FW, OC}, param.filter_type); |
| 114 | A_dst.reset_ptr(static_cast<void*>(bundle.get(2))); |
| 115 | // TODO Should be removed once armv8 convolution support transpose. |
| 116 | get_relayout_opr()->exec(A_src, A_dst, inplace_cpu_handle().get()); |
| 117 | } |
| 118 | TensorND B_, C_; |
| 119 | for (size_t n = 0; n < N; ++n) { |
| 120 | gtype *C_src, *C_dst; |
| 121 | dtype* diff = const_cast<dtype*>(param.diff<dtype>() + n * param.inp_bs); |
| 122 | gtype* grad = param.grad<gtype>() + n * param.out_bs; |
| 123 | if (is1X1) { |
| 124 | C_src = grad; |
| 125 | } else { |
| 126 | C_src = static_cast<gtype*>(bundle.get(0)); |
| 127 | } |
| 128 | { |
| 129 | B_.layout = TensorLayout({OC, IH * IW}, param.diff_type); |
| 130 | B_.reset_ptr(static_cast<void*>(diff)); |
| 131 | C_.layout = TensorLayout({IC * FH * FW, IH * IW}, param.grad_type); |
| 132 | C_.reset_ptr(C_src); |
| 133 | Workspace workspace( |
| 134 | static_cast<dt_byte*>(bundle.get(1)), bundle.get_size(1)); |
| 135 | get_matmul_opr(param)->exec(A_dst, B_, C_, workspace); |
| 136 | } |
| 137 | |
| 138 | if (!is1X1) { |
| 139 | C_dst = grad; |
| 140 | std::memset(C_dst, 0, param.grad_type.size() * IC * OH * OW); |
| 141 | if (PH == 0 && PW == 0 && SH == 1 && SW == 1) { |
nothing calls this directly
no test coverage detected