diff options
Diffstat (limited to 'src/gpu/cl/kernels/ClTransposeKernel.cpp')
| -rw-r--r-- | src/gpu/cl/kernels/ClTransposeKernel.cpp | 46 |
1 files changed, 35 insertions, 11 deletions
diff --git a/src/gpu/cl/kernels/ClTransposeKernel.cpp b/src/gpu/cl/kernels/ClTransposeKernel.cpp index 6eb2bf81c0..f95a215107 100644 --- a/src/gpu/cl/kernels/ClTransposeKernel.cpp +++ b/src/gpu/cl/kernels/ClTransposeKernel.cpp @@ -58,14 +58,37 @@ void ClTransposeKernel::configure(const CLCompileContext &compile_context, const const TensorShape dst_shape = misc::shape_calculator::compute_transposed_shape(*src); auto_init_if_empty(*dst, src->clone()->set_tensor_shape(dst_shape)); + // Explicitly set the tensor shape to preserve dimensions + dst->set_tensor_shape(dst_shape); + ARM_COMPUTE_ERROR_THROW_ON(ClTransposeKernel::validate(src, dst)); auto padding_info = get_padding_info({src, dst}); - // Create kernel - const unsigned int vec_size_x = adjust_vec_size(max_cl_vector_width / src->element_size(), src->dimension(0)); - const int vec_size_x_leftovers = src->dimension(0) % vec_size_x; - const unsigned int vec_size_y = adjust_vec_size(max_cl_vector_width / src->element_size(), src->dimension(1)); - const int vec_size_y_leftovers = src->dimension(1) % vec_size_y; + unsigned int vec_size_x; + unsigned int vec_size_y; + + // Set the optimal tile size for each data type without register spilling + switch (src->element_size()) + { + case 1: + vec_size_x = adjust_vec_size(8, src->dimension(0)); + vec_size_y = adjust_vec_size(16, src->dimension(1)); + break; + case 2: + vec_size_x = adjust_vec_size(8, src->dimension(0)); + vec_size_y = adjust_vec_size(8, src->dimension(1)); + break; + case 4: + vec_size_x = adjust_vec_size(4, src->dimension(0)); + vec_size_y = adjust_vec_size(8, src->dimension(1)); + break; + default: + ARM_COMPUTE_ERROR("Unsupported data type"); + break; + } + + const int vec_size_x_leftovers = src->dimension(0) % vec_size_x; + const int vec_size_y_leftovers = src->dimension(1) % vec_size_y; CLBuildOptions build_opts; build_opts.add_option("-DDATA_TYPE_IN_BYTES=" + support::cpp11::to_string(src->element_size())); @@ -78,7 +101,7 @@ void ClTransposeKernel::configure(const CLCompileContext &compile_context, const // Configure kernel window Window win = calculate_max_window(*src, Steps(vec_size_x, vec_size_y)); - ICLKernel::configure_internal(win, cl::NDRange(2, 8)); + ICLKernel::configure_internal(win); ARM_COMPUTE_ERROR_ON(has_padding_changed(padding_info)); } @@ -87,7 +110,6 @@ Status ClTransposeKernel::validate(const ITensorInfo *src, const ITensorInfo *ds ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src, dst); ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(src); ARM_COMPUTE_RETURN_ERROR_ON(src->data_type() == DataType::UNKNOWN); - ARM_COMPUTE_RETURN_ERROR_ON_MSG(src->num_dimensions() > 2, "Transpose up to 2-D src tensor is supported"); // Validate configured dst if (dst->total_size() != 0) @@ -112,15 +134,17 @@ void ClTransposeKernel::run_op(ITensorPack &tensors, const Window &window, cl::C utils::cast::polymorphic_downcast<const ICLTensor *>(tensors.get_const_tensor(TensorType::ACL_SRC)); auto dst = utils::cast::polymorphic_downcast<ICLTensor *>(tensors.get_tensor(TensorType::ACL_DST)); - Window slice = window.first_slice_window_2D(); + // Collapse dimensions higher than width and height into the batch dimension + Window collapsed = window.collapse_if_possible(ICLKernel::window(), Window::DimZ); + Window slice = collapsed.first_slice_window_3D(); do { unsigned int idx = 0; - add_2D_tensor_argument(idx, src, slice); - add_2D_tensor_argument(idx, dst, slice); + add_3D_tensor_argument(idx, src, slice); + add_3D_tensor_argument(idx, dst, slice); enqueue(queue, *this, slice, lws_hint()); - } while (window.slide_window_slice_2D(slice)); + } while (collapsed.slide_window_slice_3D(slice)); } } // namespace kernels } // namespace opencl |