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

Method exec_internal

dnn/src/cuda/convolution/backward_data/matmul.cpp:97–146  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

95
96template <typename T>
97void ConvolutionBackwardDataImpl::AlgoMatmul::exec_internal(const ExecArgs& args) {
98 auto&& fm = args.filter_meta;
99 size_t N = args.grad_layout->shape[0], IC = fm.icpg,
100 IH = args.grad_layout->shape[2], IW = args.grad_layout->shape[3],
101 OC = fm.ocpg, OH = args.diff_layout->shape[2],
102 OW = args.diff_layout->shape[3], FH = fm.spatial[0], FW = fm.spatial[1],
103 PH = fm.padding[0], PW = fm.padding[1], SH = fm.stride[0], SW = fm.stride[1],
104 DH = fm.dilation[0], DW = fm.dilation[1];
105 auto stream = cuda_stream(args.handle);
106
107 auto config = prepare_sub_opr(args);
108
109 auto&& sizes = matmul_get_workspace_bundle(args.as_fwd_args());
110 sizes.push_back(config.second->get_workspace_in_bytes(
111 config.first[0], config.first[1], config.first[2]));
112 auto wbundle = WorkspaceBundle(args.workspace.raw_ptr, sizes);
113
114 T* diff_t = static_cast<T*>(wbundle.get(0));
115 T* col = static_cast<T*>(wbundle.get(1));
116 {
117 // transpose diff
118 TensorLayout froml({N, OC * OH * OW}, typename DTypeTrait<T>::dtype()),
119 tol(froml);
120 froml.stride[0] = args.diff_layout->stride[0];
121 tol.stride[0] = 1;
122 tol.stride[1] = N;
123 TensorND from(args.diff_tensor->ptr<T>(), froml), to(diff_t, tol);
124 args.handle->relayout_opr()->exec(from, to);
125 }
126 {
127 // take gemm grad
128 TensorLayout Al({OC, IC * FH * FW}, typename DTypeTrait<T>::dtype()),
129 Bl({IC * FH * FW, OH * OW * N}, typename DTypeTrait<T>::dtype()),
130 Cl({OC, OH * OW * N}, typename DTypeTrait<T>::dtype());
131 TensorND A(args.filter_tensor->ptr<T>(), Al), B(col, Bl), C(diff_t, Cl);
132 if (fm.should_flip) {
133 convolution::flip_filter(
134 args.as_fwd_args(), wbundle.get_workspace(2), A.get_ref_ptr());
135 config.second->exec(A, C, B, wbundle.get_workspace(3));
136 } else {
137 config.second->exec(A, C, B, wbundle.get_workspace(2));
138 }
139 }
140 {
141 // col2im
142 convolution::col2im<T>(
143 col, args.grad_tensor->ptr<T>(), N, args.grad_layout->stride[0], IC, IH,
144 IW, FH, FW, OH, OW, PH, PW, SH, SW, DH, DW, stream);
145 }
146}
147
148// vim: syntax=cpp.doxygen

Callers

nothing calls this directly

Calls 13

cuda_streamFunction · 0.85
WorkspaceBundleClass · 0.85
flip_filterFunction · 0.85
prepare_sub_oprFunction · 0.70
as_fwd_argsMethod · 0.45
push_backMethod · 0.45
getMethod · 0.45
execMethod · 0.45
relayout_oprMethod · 0.45
get_workspaceMethod · 0.45

Tested by

no test coverage detected