1 files changed, 33 insertions, 8 deletions
diff --git a/arm_compute/runtime/CL/functions/CLFlattenLayer.h b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
index 98cf49af48..182e97dc67 100644
--- a/arm_compute/runtime/CL/functions/CLFlattenLayer.h
+++ b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2020 ARM Limited.
+ * Copyright (c) 2017-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -25,22 +25,40 @@
 #define ARM_COMPUTE_CLFLATTENLAYER_H
 
 #include "arm_compute/core/Types.h"
-#include "arm_compute/runtime/CL/ICLSimpleFunction.h"
+#include "arm_compute/runtime/IFunction.h"
 
+#include <memory>
 namespace arm_compute
 {
+class CLCompileContext;
 class ICLTensor;
+class ITensorInfo;
 
-/** Basic function to execute flatten. This function calls the following OpenCL kernel:
-*
-* -# @ref CLFlattenLayerKernel
-*
-*/
-class CLFlattenLayer : public ICLSimpleFunction
+/** Basic function to execute flatten */
+class CLFlattenLayer : public IFunction
 {
 public:
+    CLFlattenLayer();
+    /** Destructor */
+    ~CLFlattenLayer();
+    /** Prevent instances of this class from being copied (As this class contains pointers) */
+    CLFlattenLayer(const CLFlattenLayer &) = delete;
+    /** Default move constructor */
+    CLFlattenLayer(CLFlattenLayer &&);
+    /** Prevent instances of this class from being copied (As this class contains pointers) */
+    CLFlattenLayer &operator=(const CLFlattenLayer &) = delete;
+    /** Default move assignment operator */
+    CLFlattenLayer &operator=(CLFlattenLayer &&);
     /** Initialise the kernel's input and output.
      *
+     * Valid data layouts:
+     * - All
+     *
+     * Valid data type configurations:
+     * |src            |dst            |
+     * |:--------------|:--------------|
+     * |All            |All            |
+     *
      * @param[in]  input  First input tensor to flatten with at least 3 dimensions.
      *                    The dimensions above the third will be interpreted as batches. Data types supported: All.
      * @param[out] output Output tensor with shape [w*h*d, input_batches] where:
@@ -66,6 +84,13 @@ public:
      * @return a status
      */
     static Status validate(const ITensorInfo *input, const ITensorInfo *output);
+
+    // Inherited methods overridden:
+    void run() override;
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> _impl;
 };
 } // namespace arm_compute