3 files changed, 19 insertions, 7 deletions
diff --git a/arm_compute/runtime/CL/functions/CLFlattenLayer.h b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
index f5f4ff554f..ffe06aa610 100644
--- a/arm_compute/runtime/CL/functions/CLFlattenLayer.h
+++ b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
@@ -25,7 +25,8 @@
 #define ARM_COMPUTE_CLFLATTENLAYER_H
 
 #include "arm_compute/core/Types.h"
-#include "arm_compute/runtime/CL/ICLSimpleFunction.h"
+#include "arm_compute/runtime/CL/functions/CLReshapeLayer.h"
+#include "arm_compute/runtime/IFunction.h"
 
 namespace arm_compute
 {
@@ -38,7 +39,7 @@ class ITensorInfo;
 * -# @ref CLFlattenLayerKernel
 *
 */
-class CLFlattenLayer : public ICLSimpleFunction
+class CLFlattenLayer : public IFunction
 {
 public:
     /** Initialise the kernel's input and output.
@@ -68,6 +69,12 @@ public:
      * @return a status
      */
     static Status validate(const ITensorInfo *input, const ITensorInfo *output);
+
+    // Inherited methods overridden:
+    void run() override;
+
+private:
+    CLReshapeLayer _reshape{};
 };
 } // namespace arm_compute
 
diff --git a/arm_compute/runtime/NEON/functions/NEFlattenLayer.h b/arm_compute/runtime/NEON/functions/NEFlattenLayer.h
index 9f0d5226de..1104aac77f 100644
--- a/arm_compute/runtime/NEON/functions/NEFlattenLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEFlattenLayer.h
@@ -25,7 +25,8 @@
 #define ARM_COMPUTE_NEFLATTENLAYER_H
 
 #include "arm_compute/core/Types.h"
-#include "arm_compute/runtime/NEON/INESimpleFunctionNoBorder.h"
+#include "arm_compute/runtime/IFunction.h"
+#include "arm_compute/runtime/NEON/functions/NEReshapeLayer.h"
 
 namespace arm_compute
 {
@@ -33,7 +34,7 @@ class ITensor;
 class ITensorInfo;
 
 /** Basic function to execute flatten layer kernel. */
-class NEFlattenLayer : public INESimpleFunctionNoBorder
+class NEFlattenLayer : public IFunction
 {
 public:
     /** Initialise the kernel's input and output.
@@ -54,6 +55,12 @@ public:
      * @return a status
      */
     static Status validate(const ITensorInfo *input, const ITensorInfo *output);
+
+    // Inherited methods overridden:
+    void run() override;
+
+private:
+    NEReshapeLayer _reshape{};
 };
 } // namespace arm_compute
 
diff --git a/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h b/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
index 0a7748a94b..8bf3e95f81 100644
--- a/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
@@ -35,8 +35,6 @@
 
 namespace arm_compute
 {
-class NEFlattenLayerKernel;
-
 /** Basic function to reshape the weights of Fully Connected layer with NEON. This function calls the following kernels:
  *
  * @note  The fully connected layer accepts "weights" tensors only with 2 dimensions.
@@ -181,7 +179,7 @@ private:
 
     MemoryGroup                                                         _memory_group;
     IWeightsManager                                                    *_weights_manager;
-    std::unique_ptr<NEFlattenLayerKernel>                               _flatten_kernel;
+    NEFlattenLayer                                                      _flatten;
     NEConvertFullyConnectedWeights                                      _convert_weights;
     weights_transformations::NEConvertFullyConnectedWeightsManaged      _convert_weights_managed;
     NEFullyConnectedLayerReshapeWeights                                 _reshape_weights_function;