COMPMID-1973: Implement FFTConvolutionLayer on NEON

Change-Id: I2e667c0411bda0164a616ffe44473a78de6752c9
Signed-off-by: giuros01 <giuseppe.rossini@arm.com>
Reviewed-on: https://review.mlplatform.org/c/1066
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

4 files changed, 76 insertions, 20 deletions

diff --git a/arm_compute/core/NEON/kernels/NEFFTDigitReverseKernel.h b/arm_compute/core/NEON/kernels/NEFFTDigitReverseKernel.h
index 84d55fd8f4..0090c8cb14 100644
--- a/arm_compute/core/NEON/kernels/NEFFTDigitReverseKernel.h
+++ b/arm_compute/core/NEON/kernels/NEFFTDigitReverseKernel.h
@@ -24,6 +24,7 @@
 #ifndef __ARM_COMPUTE_NEFFTDIGITREVERSEKERNEL_H__
 #define __ARM_COMPUTE_NEFFTDIGITREVERSEKERNEL_H__
 
+#include "arm_compute/core/KernelDescriptors.h"
 #include "arm_compute/core/NEON/INEKernel.h"
 
 namespace arm_compute
@@ -53,31 +54,40 @@ public:
     ~NEFFTDigitReverseKernel() = default;
     /** Set the input and output tensors.
      *
-     * @param[in]  input  Source tensor. Data types supported: F32.
-     * @param[out] output Destination tensor. Data type supported: same as @p input
+     * @param[in]  input  Source tensor. Data types supported: F32. Number of channels supported: 1 (real tensor) or 2 (complex tensor).
+     * @param[out] output Destination tensor. Data type supported: same as @p input. Number of channels supported: 2 (complex tensor).
      * @param[in]  idx    Digit reverse index tensor. Data type supported: U32
-     * @param[in]  axis   Axis to perform digit reverse on.
+     * @param[in]  config Kernel configuration.
      */
-    void configure(const ITensor *input, ITensor *output, const ITensor *idx, unsigned int axis);
+    void configure(const ITensor *input, ITensor *output, const ITensor *idx, const FFTDigitReverseKernelInfo &config);
+
     /** Static function to check if given info will lead to a valid configuration of @ref NEFFTDigitReverseKernel
      *
-     * @param[in] input  Source tensor info. Data types supported: F32.
-     * @param[in] output Destination tensor info. Data type supported: same as @p input
+     * @param[in] input  Source tensor info. Data types supported: F32. Number of channels supported: 1 (real tensor) or 2 (complex tensor).
+     * @param[in] output Destination tensor info. Data type supported: same as @p input. Number of channels supported: 2 (complex tensor).
      * @param[in] idx    Digit reverse index tensor info. Data type supported: U32
-     * @param[in] axis   Axis to perform digit reverse on.
+     * @param[in] config Kernel configuration
      *
      * @return a status
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const ITensorInfo *idx, unsigned int axis);
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const ITensorInfo *idx, const FFTDigitReverseKernelInfo &config);
 
     // Inherited methods overridden:
     void run(const Window &window, const ThreadInfo &info) override;
 
 private:
-    const ITensor *_input;
-    ITensor       *_output;
-    const ITensor *_idx;
-    unsigned int   _axis;
+    using NEFFTDigitReverseKernelFunctionPtr = void (NEFFTDigitReverseKernel::*)(const Window &window);
+
+    template <bool is_input_complex, bool is_conj>
+    void digit_reverse_kernel_axis_0(const Window &window);
+
+    template <bool is_input_complex, bool is_conj>
+    void digit_reverse_kernel_axis_1(const Window &window);
+
+    NEFFTDigitReverseKernelFunctionPtr _func;
+    const ITensor                     *_input;
+    ITensor                           *_output;
+    const ITensor                     *_idx;
 };
 } // namespace arm_compute
 #endif /*__ARM_COMPUTE_NEFFTDIGITREVERSEKERNEL_H__ */
diff --git a/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h b/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h
index 8498d3c613..2210980816 100644
--- a/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h
+++ b/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h
@@ -59,15 +59,15 @@ public:
      *
      * @note If the output tensor is nullptr, the FFT will be performed in-place
      *
-     * @param[in,out] input  Source tensor. Data types supported: F32.
-     * @param[out]    output Destination tensor. Data type supported: same as @p input
+     * @param[in,out] input  Source tensor. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[out]    output Destination tensor. Data type supported: same as @p input. Number of channels supported: same as @p input.
      * @param[in]     config FFT descriptor metadata.
      */
     void configure(ITensor *input, ITensor *output, const FFTRadixStageKernelInfo &config);
     /** Static function to check if given info will lead to a valid configuration of @ref NEFFTRadixStageKernel
      *
-     * @param[in] input  Source tensor info. Data types supported: F32.
-     * @param[in] output Destination tensor info. Data type supported: same as @p input
+     * @param[in] input  Source tensor info. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[in] output Destination tensor info. Data type supported: same as @p input. Number of channels supported: same as @p input.
      * @param[in] config FFT descriptor metadata.
      *
      * @return a status
diff --git a/arm_compute/core/NEON/kernels/NEFFTScaleKernel.h b/arm_compute/core/NEON/kernels/NEFFTScaleKernel.h
index 5a19af7e62..51e6d5ab20 100644
--- a/arm_compute/core/NEON/kernels/NEFFTScaleKernel.h
+++ b/arm_compute/core/NEON/kernels/NEFFTScaleKernel.h
@@ -55,15 +55,15 @@ public:
     ~NEFFTScaleKernel() = default;
     /** Set the input and output tensors.
      *
-     * @param[in,out] input  Source tensor. Data types supported: F32.
-     * @param[out]    output Destination tensor. Data type supported: same as @p input
+     * @param[in,out] input  Source tensor. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[out]    output Destination tensor. Data type supported: same as @p input. Number of channels supported: 1 (real tensor) or 2 (complex tensor).
      * @param[in]     config Kernel configuration
      */
     void configure(ITensor *input, ITensor *output, const FFTScaleKernelInfo &config);
     /** Static function to check if given info will lead to a valid configuration of @ref NEFFTScaleKernel
      *
-     * @param[in] input  Source tensor info. Data types supported: F32.
-     * @param[in] output Destination tensor info. Data type supported: same as @p input
+     * @param[in] input  Source tensor info. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[in] output Destination tensor info. Data type supported: same as @p input. Number of channels supported: 1 (real tensor) or 2 (complex tensor).
      * @param[in] config Kernel configuration
      *
      * @return a status
diff --git a/arm_compute/core/NEON/kernels/NEPixelWiseMultiplicationKernel.h b/arm_compute/core/NEON/kernels/NEPixelWiseMultiplicationKernel.h
index 2a8e36b1de..cbb961f235 100644
--- a/arm_compute/core/NEON/kernels/NEPixelWiseMultiplicationKernel.h
+++ b/arm_compute/core/NEON/kernels/NEPixelWiseMultiplicationKernel.h
@@ -128,5 +128,51 @@ private:
     float          _scale;
     int            _scale_exponent;
 };
+
+/** Interface for the complex pixelwise multiplication kernel. */
+class NEComplexPixelWiseMultiplicationKernel : public INEKernel
+{
+public:
+    const char *name() const override
+    {
+        return "NEComplexPixelWiseMultiplicationKernel";
+    }
+    /** Default constructor.*/
+    NEComplexPixelWiseMultiplicationKernel();
+    /** Prevent instances of this class from being copied (As this class contains pointers) */
+    NEComplexPixelWiseMultiplicationKernel(const NEComplexPixelWiseMultiplicationKernel &) = delete;
+    /** Prevent instances of this class from being copied (As this class contains pointers) */
+    NEComplexPixelWiseMultiplicationKernel &operator=(const NEComplexPixelWiseMultiplicationKernel &) = delete;
+    /** Allow instances of this class to be moved */
+    NEComplexPixelWiseMultiplicationKernel(NEComplexPixelWiseMultiplicationKernel &&) = default;
+    /** Allow instances of this class to be moved */
+    NEComplexPixelWiseMultiplicationKernel &operator=(NEComplexPixelWiseMultiplicationKernel &&) = default;
+    /** Initialise the kernel's input, output and border mode.
+     *
+     * @param[in]  input1 An input tensor. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[in]  input2 An input tensor. Data types supported: same as @p input1. Number of channels supported: same as @p input1.
+     * @param[out] output The output tensor, Data types supported: same as @p input1.  Number of channels supported: same as @p input1.
+     */
+    void configure(const ITensor *input1, const ITensor *input2, ITensor *output);
+    /** Static function to check if given info will lead to a valid configuration of @ref NEComplexPixelWiseMultiplicationKernel
+     *
+     * @param[in] input1 An input tensor info. Data types supported: F32. Number of channels supported: 2 (complex tensor).
+     * @param[in] input2 An input tensor info. Data types supported: same as @p input1. Number of channels supported: same as @p input1.
+     * @param[in] output The output tensor info. Data types supported: same as @p input1. Number of channels supported: same as @p input1.
+     *
+     * @return a status
+     */
+    static Status validate(const ITensorInfo *input1, const ITensorInfo *input2, const ITensorInfo *output);
+
+    // Inherited methods overridden:
+    void run(const Window &window, const ThreadInfo &info) override;
+    BorderSize border_size() const override;
+
+private:
+    const ITensor *_input1;
+    const ITensor *_input2;
+    ITensor       *_output;
+};
+
 } // namespace arm_compute
 #endif /*__ARM_COMPUTE_NEPIXELWISEMULTIPLICATIONKERNEL_H__ */