COMPMID-1451: Fix allocation of weights in Deconvolution

Change-Id: If3ca0b034a7448df1e5349b51a2b124f1b4e99c1
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/153956
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Tested-by: bsgcomp <bsgcomp@arm.com>

1 files changed, 15 insertions, 5 deletions

diff --git a/src/runtime/CL/functions/CLDeconvolutionLayer.cpp b/src/runtime/CL/functions/CLDeconvolutionLayer.cpp
index 951d1ec4f0..bbf4e66adb 100644
--- a/src/runtime/CL/functions/CLDeconvolutionLayer.cpp
+++ b/src/runtime/CL/functions/CLDeconvolutionLayer.cpp
@@ -42,7 +42,7 @@ CLDeconvolutionLayer::CLDeconvolutionLayer(std::shared_ptr<IMemoryManager> memor
       _conv_f(),
       _flip_weights(),
       _scaled_output(),
-      _weights(),
+      _original_weights(nullptr),
       _weights_flipped(),
       _is_prepared(false)
 {
@@ -120,7 +120,7 @@ void CLDeconvolutionLayer::configure(ICLTensor *input, ICLTensor *weights, const
     const size_t idx_w = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
     const size_t idx_h = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
 
-    _weights = weights;
+    _original_weights = weights;
     _weights_flipped.allocator()->init(weights->info()->clone()->set_data_layout(data_layout));
     _flip_weights.configure(weights, &_weights_flipped);
 
@@ -138,7 +138,6 @@ void CLDeconvolutionLayer::configure(ICLTensor *input, ICLTensor *weights, const
     _is_prepared = false;
 
     _memory_group.manage(&_scaled_output);
-    _memory_group.manage(&_weights_flipped);
 
     // Find the upsampled dimensions and the padding needed for the convolution with stride 1 in order to match output shape
     unsigned int      padx            = 0;
@@ -175,14 +174,25 @@ void CLDeconvolutionLayer::prepare()
 {
     if(!_is_prepared)
     {
+        ARM_COMPUTE_ERROR_ON(!_original_weights->is_used());
+
+        // Run weights flipping and mark original weights tensor as unused
         _weights_flipped.allocator()->allocate();
         _weights_flipped.map(true);
-        _weights->map(CLScheduler::get().queue(), true);
+        _original_weights->map(CLScheduler::get().queue(), true);
         CPPScheduler::get().schedule(&_flip_weights, Window::DimZ);
         _weights_flipped.unmap();
-        _weights->unmap(CLScheduler::get().queue());
+        _original_weights->unmap(CLScheduler::get().queue());
+        _original_weights->mark_as_unused();
+
+        // Prepare convolution
         _conv_f.prepare();
 
+        if(!_weights_flipped.is_used())
+        {
+            _weights_flipped.allocator()->free();
+        }
+
         _is_prepared = true;
     }
 }