3 files changed, 13 insertions, 5 deletions
diff --git a/arm_compute/runtime/NEON/functions/NEConvolutionLayer.h b/arm_compute/runtime/NEON/functions/NEConvolutionLayer.h
index 1bd7e6a95f..8e040b3055 100644
--- a/arm_compute/runtime/NEON/functions/NEConvolutionLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEConvolutionLayer.h
@@ -34,6 +34,7 @@
 #include "arm_compute/core/NEON/kernels/NEIm2ColKernel.h"
 #include "arm_compute/core/NEON/kernels/NEWeightsReshapeKernel.h"
 #include "arm_compute/core/Types.h"
+#include "arm_compute/runtime/MemoryGroup.h"
 #include "arm_compute/runtime/Tensor.h"
 
 namespace arm_compute
@@ -48,7 +49,7 @@ class NEConvolutionLayerReshapeWeights : public IFunction
 {
 public:
     /** Constructor */
-    NEConvolutionLayerReshapeWeights();
+    NEConvolutionLayerReshapeWeights(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
     /** Set the input and output tensors.
      *
      * @param[in]  weights      Weights tensor. Weights are 4D tensor with dimensions [kernel_x, kernel_y, IFM, OFM]. Data type supported: QS8/QS16/F32.
@@ -62,6 +63,7 @@ public:
     void run() override;
 
 private:
+    MemoryGroup              _memory_group;
     NEWeightsReshapeKernel   _weights_reshape_kernel;
     NEGEMMTranspose1xWKernel _weights_transposed_kernel;
     Tensor                   _weights_reshaped;
@@ -79,7 +81,7 @@ class NEConvolutionLayer : public IFunction
 {
 public:
     /** Constructor */
-    NEConvolutionLayer();
+    NEConvolutionLayer(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
     /** Set the input and output tensors.
      *
      * @param[in]  input        Source tensor. 3 lower dimensions represent a single input [width, height, IFM],
@@ -98,6 +100,7 @@ public:
     void run() override;
 
 private:
+    MemoryGroup                      _memory_group;
     NEIm2ColKernel                   _input_im2col_kernel;
     NEGEMMInterleave4x4Kernel        _input_interleave_kernel;
     NEConvolutionLayerReshapeWeights _reshape_weights;
diff --git a/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h b/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
index 08099b8539..463a7d53e3 100644
--- a/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h
@@ -32,6 +32,7 @@
 #include "arm_compute/core/NEON/kernels/NEGEMMTranspose1xWKernel.h"
 #include "arm_compute/core/NEON/kernels/NEIm2ColKernel.h"
 #include "arm_compute/core/NEON/kernels/NETransposeKernel.h"
+#include "arm_compute/runtime/MemoryGroup.h"
 #include "arm_compute/runtime/Tensor.h"
 
 namespace arm_compute
@@ -47,7 +48,7 @@ class NEFullyConnectedLayerReshapeWeights : public IFunction
 {
 public:
     /** Constructor */
-    NEFullyConnectedLayerReshapeWeights();
+    NEFullyConnectedLayerReshapeWeights(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
     /** Set the input and output tensors.
      *
      * @param[in]  input               Weights tensor. The weights must be 2 dimensional. Data types supported: QS8/QS16/F32.
@@ -61,6 +62,7 @@ public:
     void run() override;
 
 private:
+    MemoryGroup              _memory_group;
     NETransposeKernel        _transpose_kernel;
     NEGEMMTranspose1xWKernel _transpose1xW_kernel;
     Tensor                   _transpose_output;
@@ -81,7 +83,7 @@ class NEFullyConnectedLayer : public IFunction
 {
 public:
     /** Constructor */
-    NEFullyConnectedLayer();
+    NEFullyConnectedLayer(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
     /** Set the input and output tensors.
      *
      * @param[in]  input                Source tensor. Data type supported: QS8/QS16/F32.
@@ -97,6 +99,7 @@ public:
     void run() override;
 
 private:
+    MemoryGroup                         _memory_group;
     NEIm2ColKernel                      _im2col_kernel;
     NEFullyConnectedLayerReshapeWeights _reshape_weights_kernel;
     NEGEMMInterleave4x4Kernel           _interleave4x4_kernel;
diff --git a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
index 01402aee63..a265f70043 100644
--- a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
+++ b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
@@ -27,6 +27,7 @@
 #include "arm_compute/core/NEON/kernels/NEFillBorderKernel.h"
 #include "arm_compute/core/NEON/kernels/NESoftmaxLayerKernel.h"
 #include "arm_compute/runtime/IFunction.h"
+#include "arm_compute/runtime/MemoryGroup.h"
 #include "arm_compute/runtime/Tensor.h"
 
 namespace arm_compute
@@ -47,7 +48,7 @@ class NESoftmaxLayer : public IFunction
 {
 public:
     /** Constructor */
-    NESoftmaxLayer();
+    NESoftmaxLayer(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
     /** Set the input and output tensors.
      *
      * @param[in]  input  Source tensor. Data types supported: QS8/QS16/F16/F32.
@@ -59,6 +60,7 @@ public:
     void run() override;
 
 private:
+    MemoryGroup                 _memory_group;
     NELogits1DMaxKernel         _max_kernel;
     NELogits1DShiftExpSumKernel _shift_exp_sum_kernel;
     NELogits1DNormKernel        _norm_kernel;