diff options
Diffstat (limited to 'arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h')
| -rw-r--r-- | arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h | 19 |
1 files changed, 12 insertions, 7 deletions
diff --git a/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h b/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h index a4c4be6f35..8498d3c613 100644 --- a/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h +++ b/arm_compute/core/NEON/kernels/NEFFTRadixStageKernel.h @@ -24,10 +24,10 @@ #ifndef __ARM_COMPUTE_NEFFTRADIXSTAGEKERNEL_H__ #define __ARM_COMPUTE_NEFFTRADIXSTAGEKERNEL_H__ -#include "arm_compute/core/NEON/INEKernel.h" - #include "arm_compute/core/KernelDescriptors.h" +#include "arm_compute/core/NEON/INEKernel.h" +#include <arm_neon.h> #include <set> namespace arm_compute @@ -87,12 +87,17 @@ private: ITensor *_output; bool _run_in_place; unsigned int _Nx; + unsigned int _axis; + unsigned int _radix; + + void set_radix_stage_axis0(const FFTRadixStageKernelInfo &config); + void set_radix_stage_axis1(const FFTRadixStageKernelInfo &config); - template <bool first_stage> - void set_radix_stage_fun(unsigned int radix); + using FFTFunctionPointerAxis0 = std::function<void(float *, float *, unsigned int, unsigned int, const float32x2_t &, unsigned int)>; + using FFTFunctionPointerAxis1 = std::function<void(float *, float *, unsigned int, unsigned int, const float32x2_t &, unsigned int, unsigned int)>; - using FFTFunctionPointerInPlace = std::function<void(float *, float *, unsigned int, unsigned int)>; - FFTFunctionPointerInPlace _func; + FFTFunctionPointerAxis0 _func_0; + FFTFunctionPointerAxis1 _func_1; }; } // namespace arm_compute -#endif /*__ARM_COMPUTE_NEFFTKERNEL_H__ */ +#endif /*__ARM_COMPUTE_NEFFTRADIXSTAGEKERNEL_H__ */ |