From 45fbf9fd00cc8fa564229bd6e2c20f84bc7cafd1 Mon Sep 17 00:00:00 2001
From: Georgios Pinitas <georgios.pinitas@arm.com>
Date: Fri, 13 Aug 2021 12:07:59 +0100
Subject: 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>
---
 arm_compute/runtime/IWeightsManager.h | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

(limited to 'arm_compute/runtime/IWeightsManager.h')

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
-- 
cgit v1.2.1