| 8 | namespace cuda { |
| 9 | |
| 10 | void ResizeBackwardImpl::exec( |
| 11 | _megdnn_tensor_in diff, _megdnn_tensor_out grad, _megdnn_workspace workspace) { |
| 12 | check_exec(diff.layout, grad.layout, workspace.size); |
| 13 | auto stream = cuda_stream(this->handle()); |
| 14 | bool is_nhwc = param().format == param::Resize::Format::NHWC; |
| 15 | size_t N, C, IH, IW, OH, OW; |
| 16 | if (is_nhwc) { |
| 17 | if (param().imode != Param::InterpolationMode::LINEAR && |
| 18 | is_nhwc_contig_wc(grad.layout)) { |
| 19 | megdnn_assert( |
| 20 | 0, |
| 21 | "unsupport mode in resizeBackward, only support param().imode = " |
| 22 | "LINEAR"); |
| 23 | } |
| 24 | N = grad.layout.shape[0]; |
| 25 | C = grad.layout.shape[3]; |
| 26 | IH = grad.layout.shape[1]; |
| 27 | IW = grad.layout.shape[2]; |
| 28 | OH = diff.layout.shape[1]; |
| 29 | OW = diff.layout.shape[2]; |
| 30 | } else { |
| 31 | N = grad.layout.shape[0], C = grad.layout.shape[1], IH = grad.layout.shape[2], |
| 32 | IW = grad.layout.shape[3], OH = diff.layout.shape[2], OW = diff.layout.shape[3]; |
| 33 | } |
| 34 | size_t max_batch_x_channel = max_batch_x_channel_size(); |
| 35 | size_t max_batch_size = max_batch_x_channel / C; |
| 36 | while (N > 0) { |
| 37 | size_t curr_batch_size = N > max_batch_size ? max_batch_size : N; |
| 38 | switch (grad.layout.dtype.enumv()) { |
| 39 | #define cb(_t) \ |
| 40 | case DTypeTrait<_t>::enumv: { \ |
| 41 | typedef DTypeTrait<_t>::ctype ct; \ |
| 42 | ct* diff_ptr = diff.ptr<ct>(); \ |
| 43 | ct* grad_ptr = grad.ptr<ct>(); \ |
| 44 | resize::backward_data_proxy( \ |
| 45 | is_nhwc, resize::get_imode(param().imode), diff_ptr, grad_ptr, \ |
| 46 | curr_batch_size, C, IH, IW, OH, OW, stream); \ |
| 47 | if (N <= max_batch_size) { \ |
| 48 | return; \ |
| 49 | } else { \ |
| 50 | N -= max_batch_size; \ |
| 51 | diff_ptr += curr_batch_size * diff.layout.stride[0]; \ |
| 52 | grad_ptr += curr_batch_size * grad.layout.stride[0]; \ |
| 53 | } \ |
| 54 | break; \ |
| 55 | } |
| 56 | cb(megdnn::dtype::Float32); |
| 57 | DNN_INC_FLOAT16(cb(megdnn::dtype::Float16)); |
| 58 | default: |
| 59 | megdnn_throw(ssprintf( |
| 60 | "unsupported dtype: %s in resize backward", |
| 61 | grad.layout.dtype.name())); |
| 62 | } |
| 63 | #undef cb |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | size_t ResizeBackwardImpl::get_workspace_in_bytes( |