Avoid releasing weights if they are used by multiple functions

Resolves: COMPMID-4769

Signed-off-by: Georgios Pinitas <georgios.pinitas@arm.com>
Change-Id: Iccadcbd68b0fd84ed3bf212e358a4ea944084a40
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/c/VisualCompute/ComputeLibrary/+/349845
Tested-by: bsgcomp <bsgcomp@arm.com>
Reviewed-by: Giorgio Arena <giorgio.arena@arm.com>
Comments-Addressed: bsgcomp <bsgcomp@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6107
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Reviewed-by: SiCong Li <sicong.li@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

1 files changed, 31 insertions, 2 deletions

diff --git a/src/runtime/IWeightsManager.cpp b/src/runtime/IWeightsManager.cpp
index 081cd990f3..8f27795285 100644
--- a/src/runtime/IWeightsManager.cpp
+++ b/src/runtime/IWeightsManager.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019 Arm Limited.
+ * Copyright (c) 2019, 2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -26,7 +26,7 @@
 namespace arm_compute
 {
 IWeightsManager::IWeightsManager()
-    : _managed_weights(), _managed_weights_parents()
+    : _managed_weights(), _managed_counter(), _managed_weights_parents()
 {
 }
 
@@ -35,6 +35,11 @@ void IWeightsManager::manage(const ITensor *weights, ITransformWeights *parent)
     if(!are_weights_managed(weights))
     {
         _managed_weights[weights];
+        _managed_counter[weights];
+    }
+    else
+    {
+        _managed_counter[weights].counter++;
     }
 
     // In case the weights are an output of a previous reshape function
@@ -146,4 +151,28 @@ ITensor *IWeightsManager::acquire(const ITensor *weights, ITransformWeights *wei
 
     return transformed_weights;
 }
+
+void IWeightsManager::release(const ITensor *weights)
+{
+    if(weights == nullptr || !are_weights_managed(weights))
+    {
+        return;
+    }
+
+    _managed_counter[weights].counter--;
+    if(_managed_counter[weights].counter == 0 && _managed_counter[weights].is_unused)
+    {
+        weights->mark_as_unused();
+    }
+}
+
+void IWeightsManager::mark_as_unused(const ITensor *weights)
+{
+    if(weights == nullptr || !are_weights_managed(weights))
+    {
+        return;
+    }
+
+    _managed_counter[weights].is_unused = true;
+}
 } // namespace arm_compute