| 33 | |
| 34 | template <typename T> |
| 35 | void TileBackwardImpl::exec_internal( |
| 36 | _megdnn_tensor_in diff_, _megdnn_tensor_out grad_, |
| 37 | _megdnn_workspace workspace) { |
| 38 | TensorShape grad, diff, times; |
| 39 | simplify_shape(grad_.layout, diff_.layout, param().times, grad, diff, times); |
| 40 | auto stream = cuda_stream(this->handle()); |
| 41 | auto nr_reduces = count_not_ones_in_shape(times); |
| 42 | auto dtype = diff_.layout.dtype; |
| 43 | if (nr_reduces == 0) { |
| 44 | cuda_check(cudaMemcpyAsync( |
| 45 | grad_.raw_ptr(), diff_.raw_ptr(), sizeof(T) * diff.total_nr_elems(), |
| 46 | cudaMemcpyDeviceToDevice, stream)); |
| 47 | } else { |
| 48 | auto ndim = times.ndim; |
| 49 | WorkspaceBundle workspaces( |
| 50 | workspace.raw_ptr, |
| 51 | {diff.total_nr_elems() * sizeof(T), diff.total_nr_elems() * sizeof(T)}); |
| 52 | auto workspace0 = static_cast<T*>(workspaces.get(0)); |
| 53 | auto workspace1 = static_cast<T*>(workspaces.get(1)); |
| 54 | |
| 55 | T *current, *next; |
| 56 | size_t state; |
| 57 | |
| 58 | init_tile_repeat_state( |
| 59 | diff_.ptr<T>(), grad_.ptr<T>(), workspace0, workspace1, current, next, |
| 60 | state, nr_reduces); |
| 61 | TensorND reduce_src, reduce_dst; |
| 62 | for (size_t j = 0; j < ndim; ++j) { |
| 63 | size_t i = j + 1; |
| 64 | if (times.shape[j] != 1) { |
| 65 | // m = sshape[0]*...*sshape[i-2] |
| 66 | auto m = std::accumulate( |
| 67 | grad.shape, grad.shape + j, 1_z, SafeMultiplies<size_t>()); |
| 68 | // n = sshape[i-1]*dshape[i]*... |
| 69 | auto n = std::accumulate( |
| 70 | diff.shape + i, diff.shape + ndim, 1_z, |
| 71 | SafeMultiplies<size_t>()) * |
| 72 | grad.shape[j]; |
| 73 | // forward is repeat (m, n) to (m*times, n) |
| 74 | // backward is reduce (m, times, n) to (m, 1, n) |
| 75 | m_opr->param().axis = 1; |
| 76 | /* |
| 77 | TensorND reduce_src(current, TensorShape{m, times[j], n}); |
| 78 | TensorND reduce_dst(next, TensorShape{m, 1u, n}); |
| 79 | */ |
| 80 | |
| 81 | reduce_src.reset_ptr(current); |
| 82 | reduce_src.layout = TensorLayout(TensorShape{m, times[j], n}, dtype); |
| 83 | reduce_dst.reset_ptr(next); |
| 84 | reduce_dst.layout = TensorLayout(TensorShape{m, 1u, n}, dtype); |
| 85 | m_opr->exec(reduce_src, reduce_dst, Workspace()); |
| 86 | update_tile_repeat_state( |
| 87 | diff_.ptr<T>(), grad_.ptr<T>(), workspace0, workspace1, current, |
| 88 | next, state, nr_reduces); |
| 89 | } |
| 90 | } |
| 91 | megdnn_assert_internal(current == grad_.ptr<T>()); |
| 92 | megdnn_assert_internal(next == nullptr); |
nothing calls this directly
no test coverage detected