| 90 | } |
| 91 | |
| 92 | bool RelayoutForwardImpl::Param::try_copy_2d(bool cross_dev) { |
| 93 | TensorLayout lsrc = m_src.layout, ldst = m_dst.layout; |
| 94 | |
| 95 | if (lsrc.ndim > 2 || ldst.ndim > 2) |
| 96 | return false; |
| 97 | if (ldst.dtype.is_low_bit()) |
| 98 | return false; |
| 99 | |
| 100 | if (ldst.ndim == 1 && lsrc.ndim == 1) { |
| 101 | megdnn_assert(ldst.stride[0] != 1 || lsrc.stride[0] != 1); |
| 102 | if (lsrc.stride[0] < 1 || ldst.stride[0] < 1 || !cross_dev) |
| 103 | // test case: src=16x128x128(49152, 384, 3), dst=16x128x128(16384, 128, 1) |
| 104 | // for both src and dst are one-dimensional, and one of them are not |
| 105 | // contiguous, the relayout opr will call cudaMemcpy2DAsync, and the |
| 106 | // bandwidth=5GiB/s. it is better to call copy_general, the |
| 107 | // bandwidth=100GiB/s. call cudaMemcpy2DAsync when cross_dev, OR return |
| 108 | // false and call copy_general. |
| 109 | return false; |
| 110 | // extend to ndim == 2 |
| 111 | megdnn_assert(ldst.shape[0] == lsrc.shape[0]); |
| 112 | ldst.ndim = lsrc.ndim = 2; |
| 113 | ldst.shape[1] = lsrc.shape[1] = 1; |
| 114 | ldst.stride[1] = lsrc.stride[1] = 1; |
| 115 | } else if (ldst.ndim < 2) { |
| 116 | if (!expand_dim2(ldst, lsrc)) |
| 117 | return false; |
| 118 | } else if (lsrc.ndim < 2) { |
| 119 | if (!expand_dim2(lsrc, ldst)) |
| 120 | return false; |
| 121 | } |
| 122 | if (ldst.stride[1] != 1 || lsrc.stride[1] != 1 || ldst.shape[0] != lsrc.shape[0] || |
| 123 | ldst.shape[1] != lsrc.shape[1] || |
| 124 | ldst.stride[0] < static_cast<ptrdiff_t>(ldst.shape[1]) || |
| 125 | lsrc.stride[0] < static_cast<ptrdiff_t>(ldst.shape[1])) |
| 126 | return false; |
| 127 | |
| 128 | auto dsize = dtype_size(); |
| 129 | cuda_check(cudaMemcpy2DAsync( |
| 130 | m_dst.raw_ptr(), ldst.stride[0] * dsize, m_src.raw_ptr(), |
| 131 | lsrc.stride[0] * dsize, ldst.shape[1] * dsize, ldst.shape[0], |
| 132 | cudaMemcpyDeviceToDevice, m_opr->stream())); |
| 133 | |
| 134 | return true; |
| 135 | }; |
| 136 | |
| 137 | bool RelayoutForwardImpl::Param::try_copy_last_contig() { |
| 138 | if (m_dst.layout.dtype.is_low_bit()) |
no test coverage detected