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, 19 insertions, 1 deletions

diff --git a/arm_compute/runtime/IWeightsManager.h b/arm_compute/runtime/IWeightsManager.h
index 12aa1da8a7..db39a71314 100644
--- a/arm_compute/runtime/IWeightsManager.h
+++ b/arm_compute/runtime/IWeightsManager.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019 Arm Limited.
+ * Copyright (c) 2019, 2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -76,9 +76,27 @@ public:
      * @return True if the weights tensor is managed else false
      */
     bool are_weights_managed(const ITensor *weights);
+    /** Release weights refcount and mark as unused if reaches 0
+     *
+     * @param[in] weights Weights to release
+     */
+    void release(const ITensor *weights);
+    /** Marks weights as unused
+     *
+     * @param weights Weights to mark unused
+     */
+    void mark_as_unused(const ITensor *weights);
+
+private:
+    struct CounterElement
+    {
+        bool             is_unused{ false };
+        std::atomic<int> counter{ 1 };
+    };
 
 private:
     std::map<const ITensor *, std::vector<ITransformWeights *>> _managed_weights;
+    std::map<const ITensor *, CounterElement>                   _managed_counter;
     std::map<const ITensor *, ITransformWeights *>              _managed_weights_parents;
 };
 } // arm_compute