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

Function kern_matmul

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

Source from the content-addressed store, hash-verified

82
83template <typename ftype, typename dtype, typename gtype>
84void 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) {

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