From 05fb448bf48e31d723dfd9f4bbf3899ff65f0fba Mon Sep 17 00:00:00 2001 From: giuros01 Date: Tue, 26 Mar 2019 17:44:40 +0000 Subject: COMPMID-1963: Implement FFT (2D) on NEON Change-Id: I3b564be8d7949e00c6544071ef62dd51de838c96 Signed-off-by: giuros01 Reviewed-on: https://review.mlplatform.org/c/1048 Comments-Addressed: Arm Jenkins Tested-by: Arm Jenkins Reviewed-by: Georgios Pinitas --- src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) (limited to 'src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp') diff --git a/src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp b/src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp index 845fcef4f3..b2ffb01e99 100644 --- a/src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp +++ b/src/core/NEON/kernels/NEFFTDigitReverseKernel.cpp @@ -37,7 +37,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, c { ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 2, DataType::F32); ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(idx, 1, DataType::U32); - ARM_COMPUTE_RETURN_ERROR_ON(axis != 0); + ARM_COMPUTE_RETURN_ERROR_ON(axis > 1); // Checks performed when output is configured if((output != nullptr) && (output->total_size() != 0)) @@ -96,15 +96,24 @@ void NEFFTDigitReverseKernel::run(const Window &window, const ThreadInfo &info) Iterator out(_output, window); const size_t element_size = _input->info()->element_size(); + // Pointers to the buffers + const size_t offset = _input->info()->offset_first_element_in_bytes(); + auto *idx_ptr = reinterpret_cast(_idx->buffer()); + uint8_t *input_ptr = offset + _input->buffer(); + + // Strides + const size_t stride_x = _input->info()->strides_in_bytes()[0]; + const size_t stride_y = _input->info()->strides_in_bytes()[1]; + const size_t stride_z = _input->info()->strides_in_bytes()[2]; + const size_t stride_w = _input->info()->strides_in_bytes()[3]; + execute_window_loop(window, [&](const Coordinates & id) { - unsigned int in_index_1d = *reinterpret_cast(_idx->ptr_to_element(Coordinates(id.x()))); - - auto reverse_id = id; + unsigned int in_index_1d = idx_ptr[id[_axis]]; + auto reverse_id = id; reverse_id.set(_axis, in_index_1d); - memcpy(out.ptr(), _input->ptr_to_element(reverse_id), 2 * element_size); - + memcpy(out.ptr(), input_ptr + reverse_id.x() * stride_x + reverse_id.y() * stride_y + reverse_id.z() * stride_z + reverse_id[3] * stride_w, element_size); }, out); -- cgit v1.2.1