COMPMID-719: NEWinogradLayer reordering using NEPermute.
Input reordering from NCHW to NHWC
Output reordering from NHWC to NCHW
Weights reordering from [Ofm x Ifm x Height x Width] to [Height x Width x Ifm x Ofm]
Change-Id: I85aabedb1f9c13700bc4919eb3130f4d4bd0b465
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/113631
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>
diff --git a/src/core/CPP/kernels/CPPPermuteKernel.cpp b/src/core/CPP/kernels/CPPPermuteKernel.cpp
index 4b137b0..80b0aba 100644
--- a/src/core/CPP/kernels/CPPPermuteKernel.cpp
+++ b/src/core/CPP/kernels/CPPPermuteKernel.cpp
@@ -51,13 +51,18 @@
DataType::U32, DataType::S32,
DataType::F16, DataType::F32);
ARM_COMPUTE_RETURN_ERROR_ON_MSG(input->num_dimensions() < 3, "Invalid input size!");
- ARM_COMPUTE_RETURN_ERROR_ON_MSG(perm.num_dimensions() != 3 && ((perm[0] != 2 && perm[1] != 0 && perm[2] != 1) || (perm[0] != 1 && perm[1] != 2 && perm[2] != 0)),
- "Only [2, 0, 1] and [1, 2, 0] permutation is supported");
+ ARM_COMPUTE_RETURN_ERROR_ON_MSG(
+ (perm.num_dimensions() != 3 && ((perm[0] != 2 && perm[1] != 0 && perm[2] != 1) || (perm[0] != 1 && perm[1] != 2 && perm[2] != 0))) && (perm.num_dimensions() != 4 && ((perm[0] != 2 && perm[1] != 0
+ && perm[2] != 1)
+ || (perm[0] != 1 && perm[1] != 2 && perm[2] != 0))),
+ "Only [2, 0, 1],[1, 2, 0] and [3, 2, 0, 1] permutation is supported");
+
+ const TensorShape output_shape = get_output_shape(input, perm);
// Validate configured output
if(output->total_size() != 0)
{
- ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DIMENSIONS(output->tensor_shape(), get_output_shape(input, perm));
+ ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DIMENSIONS(output->tensor_shape(), output_shape);
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_FIXED_POINT(input, output);
}
@@ -72,11 +77,13 @@
const int output_stride_x = _output->info()->strides_in_bytes().x();
const int output_stride_y = _output->info()->strides_in_bytes().y();
const int output_stride_z = _output->info()->strides_in_bytes().z();
+ const int output_stride_w = _output->info()->strides_in_bytes()[3];
Window window_out(window);
window_out.set(Window::DimX, Window::Dimension(0, 0, 0));
window_out.set(Window::DimY, Window::Dimension(0, 0, 0));
window_out.set(Window::DimZ, Window::Dimension(0, 0, 0));
+ window_out.set(3, Window::Dimension(0, 0, 0));
// Create iterators
Iterator in(_input, window);
@@ -87,21 +94,35 @@
{
execute_window_loop(window, [&](const Coordinates & id)
{
- const int idx = id.y() * output_stride_z + id.x() * output_stride_y + id.z() * output_stride_x;
+ const int idx = id[3] * output_stride_w + id.y() * output_stride_z + id.x() * output_stride_y + id.z() * output_stride_x;
*(reinterpret_cast<T *>(out.ptr() + idx)) = *(reinterpret_cast<const T *>(in.ptr()));
},
in, out);
}
// Run [1, 2, 0] permute
- else
+ else if(_perm[0] == 1 && _perm[1] == 2 && _perm[2] == 0)
{
execute_window_loop(window, [&](const Coordinates & id)
{
- const int idx = id.x() * output_stride_z + id.z() * output_stride_y + id.y() * output_stride_x;
+ const int idx = id[3] * output_stride_w + id.x() * output_stride_z + id.z() * output_stride_y + id.y() * output_stride_x;
*(reinterpret_cast<T *>(out.ptr() + idx)) = *(reinterpret_cast<const T *>(in.ptr()));
},
in, out);
}
+ // Run [3, 2, 0, 1] permute
+ else if(_perm[0] == 3 && _perm[1] == 2 && _perm[2] == 0 && _perm[3] == 1)
+ {
+ execute_window_loop(window, [&](const Coordinates & id)
+ {
+ const int idx = id[3] * output_stride_x + id[2] * output_stride_y + id[0] * output_stride_z + id[1] * output_stride_w;
+ *(reinterpret_cast<T *>(out.ptr() + idx)) = *(reinterpret_cast<const T *>(in.ptr()));
+ },
+ in, out);
+ }
+ else
+ {
+ ARM_COMPUTE_ERROR("Not supported.");
+ }
}
CPPPermuteKernel::CPPPermuteKernel()
@@ -112,9 +133,9 @@
void CPPPermuteKernel::configure(const ITensor *input, ITensor *output, const PermutationVector &perm)
{
ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
-
+ const TensorShape output_shape = get_output_shape(input->info(), perm);
// Output auto inizialitation if not yet initialized
- auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(get_output_shape(input->info(), perm)));
+ auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape));
// Perform validation step
ARM_COMPUTE_ERROR_THROW_ON(validate_arguments(input->info(), output->info(), perm));
diff --git a/src/core/NEON/kernels/NEWinogradLayerKernel.cpp b/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
index fe63336..eaf77e6 100644
--- a/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
@@ -108,30 +108,25 @@
}
NEWinogradLayerKernel::NEWinogradLayerKernel()
- : _convolver(nullptr), _output(nullptr)
+ : _convolver(nullptr)
{
}
-void NEWinogradLayerKernel::configure(ITensor *output, Winograd3x3F32 *convolver)
+void NEWinogradLayerKernel::configure(Winograd3x3F32 *convolver)
{
- ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::F32);
+ ARM_COMPUTE_ERROR_ON_NULLPTR(convolver);
_convolver = convolver;
- Window win = calculate_max_window(*output->info());
+ Window win;
+ win.set(Window::DimX, Window::Dimension(0, 15, 1));
INEKernel::configure(win);
}
void NEWinogradLayerKernel::run(const Window &window, const ThreadInfo &info)
{
- ARM_COMPUTE_UNUSED(window);
ARM_COMPUTE_UNUSED(info);
ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
- ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window);
- ARM_COMPUTE_ERROR_ON(info.num_threads < 1);
- const size_t tid = info.thread_id;
- const size_t num_threads = std::min(info.num_threads, 16);
- const size_t num_gemms_per_thread = 16 / num_threads;
- const size_t first_gemm = tid * num_gemms_per_thread;
- const size_t last_gemm = (tid == (num_threads - 1)) ? 15 : first_gemm + num_gemms_per_thread - 1;
+ const size_t first_gemm = window.x().start();
+ const size_t last_gemm = window.x().end();
_convolver->_pimpl->convolver.execute(first_gemm, last_gemm);
}
} // namespace arm_compute
diff --git a/src/runtime/NEON/functions/NEWinogradLayer.cpp b/src/runtime/NEON/functions/NEWinogradLayer.cpp
index 3251de4..800153e 100644
--- a/src/runtime/NEON/functions/NEWinogradLayer.cpp
+++ b/src/runtime/NEON/functions/NEWinogradLayer.cpp
@@ -43,7 +43,8 @@
namespace arm_compute
{
NEWinogradLayer::NEWinogradLayer(std::shared_ptr<IMemoryManager> memory_manager)
- : _memory_group(std::move(memory_manager)), _winograd_kernel(), _weights_workspace(), _workspace(), _kernel_storage(), _input(), _weights(), _output(), _reshaped_kernel(false), _conv()
+ : _memory_group(std::move(memory_manager)), _winograd_kernel(), _permute_input(), _permute_weights(), _permute_output(), _workspace(), _kernel_storage(), _input_nhwc(), _output_nhwc(),
+ _weights_hwio(), _input(), _weights(), _output(), _reshaped_kernel(false), _conv()
{
} /* arm_compute */
@@ -71,9 +72,8 @@
ARM_COMPUTE_ERROR_ON_MSG(stride_y != 1 || stride_x != 1, "Winograd layer only supports unit strides.");
// Get convolved dimensions
- auto padding = PADDING_VALID;
- const int in_channels = input->info()->dimension(2);
-
+ auto padding = PADDING_VALID;
+ const int in_channels = input->info()->dimension(2);
const int out_channels = output->info()->dimension(2);
const int weights_width = weights->info()->dimension(0);
const int weights_height = weights->info()->dimension(1);
@@ -88,25 +88,45 @@
_memory_group.manage(&_kernel_storage);
// Get workbench size and allocate memory
+
constexpr size_t wspace_alignment = 64;
const size_t ws_size = NEWinogradLayerKernel::get_working_space_size(in_shape, kernel_shape, padding);
_workspace.allocator()->init(TensorInfo(TensorShape{ (ws_size + wspace_alignment - 1) }, 1, DataType::U8));
_memory_group.manage(&_workspace);
-
- // Workspace for weights transform
- const size_t weights_transform_size = NEWinogradLayerKernel::get_kernel_transform_working_size(kernel_shape);
- _weights_workspace.allocator()->init(TensorInfo(TensorShape{ (weights_transform_size + wspace_alignment - 1) }, 1, DataType::U8));
- _memory_group.manage(&_weights_workspace);
-
+ _memory_group.manage(&_input_nhwc);
_kernel_storage.allocator()->allocate();
_workspace.allocator()->allocate();
- _weights_workspace.allocator()->allocate();
// Create Winograd operator object
_conv = support::cpp14::make_unique<Winograd3x3F32>(kernel_shape, in_shape, padding, _kernel_storage.buffer());
// Configure the kernel, padding not needed so it's safe to call configure after allocare
- _winograd_kernel.configure(output, _conv.get());
+ _winograd_kernel.configure(_conv.get());
+
+ // Re-order a weight tensor from [Output feature map x Input feature map x Height x Width] to [Height x Width x Input feature map x Output feature map]
+ switch(weights->info()->num_dimensions())
+ {
+ case 3:
+ {
+ _permute_weights.configure(weights, &_weights_hwio, PermutationVector(2U, 0U, 1U));
+ break;
+ }
+ case 4:
+ {
+ _permute_weights.configure(weights, &_weights_hwio, PermutationVector(3U, 2U, 0U, 1U));
+ break;
+ }
+ default:
+ {
+ ARM_COMPUTE_ERROR("Not supported.");
+ break;
+ }
+ }
+ // configure the kernel to transform the input tensor from NCHW -> NHWC
+ _permute_input.configure(input, &_input_nhwc, PermutationVector(2U, 0U, 1U));
+
+ _weights_hwio.allocator()->allocate();
+ _input_nhwc.allocator()->allocate();
}
void NEWinogradLayer::run()
@@ -115,29 +135,42 @@
_memory_group.acquire();
if(!_reshaped_kernel)
{
- _conv->transform_weights(reinterpret_cast<const float *>(_weights->buffer()), reinterpret_cast<float *>(_weights_workspace.buffer()));
_reshaped_kernel = true;
+ _permute_weights.run();
+ _conv->transform_weights(reinterpret_cast<const float *>(_weights_hwio.buffer()), nullptr);
}
const Tensor4DShape in_shape(internal_get_input_shape(_input));
auto padding = PADDING_VALID;
//Bring channels to the front as Winograd code expects the tensor to be in the format NHWC
- _conv->nchw2nhwc(in_shape, padding, _workspace.buffer(), reinterpret_cast<const float *>(_input->buffer()));
+ _permute_input.run();
//Get ptrs into the workspace
std::pair<void *, void *> nhwc_ptrs = _conv->get_nhwc_ptrs(in_shape, padding, _workspace.buffer());
//Setup matrices ptrs and transfor the input tensor to the appropriate form before running GEMM.
- _conv->reshape_input(in_shape, padding, nhwc_ptrs.second, _workspace.buffer());
+ _conv->reshape_input(in_shape, padding, reinterpret_cast<float *>(_input_nhwc.buffer()), _workspace.buffer());
//Run 16 GEMMs in multiple threads, each kernel runs one or more GEMMs
- NEScheduler::get().schedule(&_winograd_kernel, Window::DimY);
+ NEScheduler::get().schedule(&_winograd_kernel, Window::DimX);
//Transform the output to the appropriate form
_conv->reshape_output(in_shape, padding, nhwc_ptrs.first);
- //Transform back to NCHW
- _conv->nhwc2nchw(in_shape, padding, _workspace.buffer(), reinterpret_cast<float *>(_output->buffer()));
+ const unsigned int out_width = _output->info()->dimension(0);
+ const unsigned int out_height = _output->info()->dimension(1);
+ const unsigned int out_channels = _output->info()->dimension(2);
+ const unsigned int out_batches = _output->info()->dimension(3);
+
+ // We create a temporary tensor with the results in the workspace so that the we can run a function to reorder from NHWC -> NCHW
+ Tensor output_nhwc;
+ TensorInfo info(TensorShape(out_channels, out_width, out_height, out_batches), 1, _output->info()->data_type());
+ output_nhwc.allocator()->init(info);
+ output_nhwc.allocator()->import_memory(Memory(static_cast<uint8_t *>(nhwc_ptrs.first)));
+
+ // Reorder the convoluted output to ACL's ordering NCHW
+ _permute_output.configure(&output_nhwc, _output, PermutationVector(1U, 2U, 0U));
+ _permute_output.run();
_memory_group.release();
#else /* __aarch64__ */