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>

6 files changed, 94 insertions, 4 deletions

diff --git a/arm_compute/core/ITensor.h b/arm_compute/core/ITensor.h
index 131ee205ea..32b93576bd 100644
--- a/arm_compute/core/ITensor.h
+++ b/arm_compute/core/ITensor.h
@@ -90,6 +90,8 @@ public:
     bool is_used() const;
     /** Marks a tensor as unused */
     void mark_as_unused() const;
+    /** Marks a tensor as used */
+    void mark_as_used() const;
 
 private:
     mutable bool _is_used = { true }; /**< Flag that marks if the tensor is used or not */
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
diff --git a/src/core/ITensor.cpp b/src/core/ITensor.cpp
index e263596333..9ed9b3f2e9 100644
--- a/src/core/ITensor.cpp
+++ b/src/core/ITensor.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -169,4 +169,9 @@ void ITensor::mark_as_unused() const
 {
     _is_used = false;
 }
+
+void ITensor::mark_as_used() const
+{
+    _is_used = true;
+}
 } // namespace arm_compute
\ No newline at end of file
diff --git a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
index ae10cd23b1..c719a667a7 100644
--- a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
+++ b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
@@ -80,6 +80,11 @@ void CLFullyConnectedLayer::configure(const CLCompileContext &compile_context, c
 
     _impl->op->configure(compile_context, input->info(), weights->info(), (biases != nullptr) ? biases->info() : nullptr, output->info(), fc_info);
 
+    if(_impl->weights_manager != nullptr)
+    {
+        _impl->weights_manager->manage(weights);
+    }
+
     _impl->aux_mem_req = _impl->op->workspace();
     _impl->run_pack    = { { ACL_SRC_0, input }, { ACL_SRC_1, weights }, { ACL_SRC_2, biases }, { ACL_DST, output } };
     _impl->workspace   = manage_workspace<CLTensor>(_impl->aux_mem_req, _impl->memory_group, _impl->run_pack, _impl->run_pack);
@@ -108,6 +113,19 @@ void CLFullyConnectedLayer::prepare()
         // Release temporary tensors that are only used in prepare stage
         release_temporaries<CLTensor>(_impl->aux_mem_req, _impl->workspace);
         _impl->is_prepared = true;
+
+        // Handle weights managed infrastructure
+        if(_impl->weights_manager != nullptr && _impl->weights_manager->are_weights_managed(_impl->original_weights))
+        {
+            // If function marks b as unused ensure that all prepare stages are done before releasing
+            const ITensor *original_b = _impl->original_weights;
+            if(!original_b->is_used())
+            {
+                _impl->weights_manager->mark_as_unused(original_b);
+            }
+            _impl->original_weights->mark_as_used();
+            _impl->weights_manager->release(_impl->original_weights);
+        }
     }
 }
 } // namespace arm_compute
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
diff --git a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
index 504200e9ce..84bc004ec4 100644
--- a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
+++ b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
@@ -76,6 +76,11 @@ void NEFullyConnectedLayer::configure(const ITensor *input, const ITensor *weigh
 
     _impl->op->configure(input->info(), weights->info(), (biases != nullptr) ? biases->info() : nullptr, output->info(), fc_info);
 
+    if(_impl->weights_manager != nullptr)
+    {
+        _impl->weights_manager->manage(weights);
+    }
+
     _impl->aux_mem_req = _impl->op->workspace();
     _impl->run_pack    = { { ACL_SRC_0, input }, { ACL_SRC_1, weights }, { ACL_SRC_2, biases }, { ACL_DST, output } };
     _impl->workspace   = manage_workspace<Tensor>(_impl->aux_mem_req, _impl->memory_group, _impl->run_pack, _impl->run_pack);
@@ -104,6 +109,19 @@ void NEFullyConnectedLayer::prepare()
         // Release temporary tensors that are only used in prepare stage
         release_temporaries<Tensor>(_impl->aux_mem_req, _impl->workspace);
         _impl->is_prepared = true;
+
+        // Handle weights managed infrastructure
+        if(_impl->weights_manager != nullptr && _impl->weights_manager->are_weights_managed(_impl->original_weights))
+        {
+            // If function marks b as unused ensure that all prepare stages are done before releasing
+            const ITensor *original_b = _impl->original_weights;
+            if(!original_b->is_used())
+            {
+                _impl->weights_manager->mark_as_unused(original_b);
+            }
+            _impl->original_weights->mark_as_used();
+            _impl->weights_manager->release(_impl->original_weights);
+        }
     }
 }
 } // namespace arm_compute