| 7 | namespace cuda { |
| 8 | |
| 9 | void TransposeForwardImpl::exec( |
| 10 | _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) { |
| 11 | check_exec(src.layout, dst.layout, workspace.size); |
| 12 | if (src.layout.dtype == dtype::Float32()) { |
| 13 | auto handle = concrete_handle(this->handle()); |
| 14 | cublas_check(cublasSgeam( |
| 15 | handle->cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_T, dst.layout.shape[1], |
| 16 | dst.layout.shape[0], handle->one_device(), src.ptr<dt_float32>(), |
| 17 | src.layout.stride[0], handle->zero_device(), src.ptr<dt_float32>(), |
| 18 | src.layout.stride[0], dst.ptr<dt_float32>(), dst.layout.stride[0])); |
| 19 | } else { |
| 20 | auto stream = cuda_stream(handle()); |
| 21 | #define cb(DType) \ |
| 22 | if (src.layout.dtype.enumv() == DTypeTrait<DType>::enumv) { \ |
| 23 | using T = typename DTypeTrait<DType>::ctype; \ |
| 24 | transpose<T>( \ |
| 25 | src.ptr<T>(), dst.ptr<T>(), src.layout.shape[0], src.layout.shape[1], \ |
| 26 | src.layout.stride[0], dst.layout.stride[0], stream); \ |
| 27 | } |
| 28 | MEGDNN_FOREACH_COMPUTING_DTYPE(cb) |
| 29 | #undef cb |
| 30 | } |
| 31 | } |
| 32 | |
| 33 | } // namespace cuda |
| 34 | } // namespace megdnn |
nothing calls this directly
no test coverage detected