COMPMID-2693 Weights manager to mark the initial weights as unsued

Change-Id: Iaa964e71d0aa80fff0a465127ac3716563bbe85f
Signed-off-by: Michalis Spyrou <michalis.spyrou@arm.com>
Reviewed-on: https://review.mlplatform.org/c/2017
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

2 files changed, 22 insertions, 2 deletions

diff --git a/src/runtime/IWeightsManager.cpp b/src/runtime/IWeightsManager.cpp
index 6dfb925fe6..b367b5f70b 100644
--- a/src/runtime/IWeightsManager.cpp
+++ b/src/runtime/IWeightsManager.cpp
@@ -85,6 +85,27 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights
         }
     }
 
+    // Check top level weights. If all the transformations are done
+    // mark the weights as unused
+    if(_managed_weights_parents.find(weights) == _managed_weights_parents.end())
+    {
+        auto item           = _managed_weights.find(weights);
+        bool mark_as_unused = true;
+        for(auto it : item->second)
+        {
+            if(!it->is_reshape_run())
+            {
+                mark_as_unused = false;
+                break;
+            }
+        }
+
+        if(mark_as_unused)
+        {
+            weights->mark_as_unused();
+        }
+    }
+
     return weights_tensor;
 }
 
diff --git a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
index 7adc3bca9e..ee622f4699 100644
--- a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
+++ b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
@@ -426,7 +426,6 @@ void NEFullyConnectedLayer::prepare()
         {
             if(_weights_manager && _weights_manager->are_weights_managed(_original_weights))
             {
-                cur_weights->mark_as_unused();
                 cur_weights = _weights_manager->run(cur_weights, &_reshape_weights_managed_function);
             }
             else
@@ -455,9 +454,9 @@ void NEFullyConnectedLayer::prepare()
             {
                 _converted_weights_output.allocator()->allocate();
                 _convert_weights.run();
+                cur_weights->mark_as_unused();
             }
 
-            cur_weights->mark_as_unused();
             _are_weights_converted = true;
         }