Address comments on avoiding releasing weights if used by multiple functions

Signed-off-by: Giorgio Arena <giorgio.arena@arm.com>
Change-Id: I0b59c5326f5fcbc322fbeb864197ea999de6bd56
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6112
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Reviewed-by: SiCong Li <sicong.li@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

4 files changed, 13 insertions, 9 deletions

diff --git a/arm_compute/runtime/IWeightsManager.h b/arm_compute/runtime/IWeightsManager.h
index db39a71314..3b97d696bb 100644
--- a/arm_compute/runtime/IWeightsManager.h
+++ b/arm_compute/runtime/IWeightsManager.h
@@ -81,11 +81,11 @@ public:
      * @param[in] weights Weights to release
      */
     void release(const ITensor *weights);
-    /** Marks weights as unused
+    /** Pre-mark the weights as unused. The weights tensor will get marked as unused only when the counter goes to 0
      *
      * @param weights Weights to mark unused
      */
-    void mark_as_unused(const ITensor *weights);
+    void pre_mark_as_unused(const ITensor *weights);
 
 private:
     struct CounterElement
diff --git a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
index 6836c12022..4f9759c590 100644
--- a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
+++ b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp
@@ -82,7 +82,7 @@ void CLFullyConnectedLayer::configure(const CLCompileContext &compile_context, c
 
     if(_impl->weights_manager != nullptr)
     {
-        _impl->weights_manager->manage(weights);
+        _impl->weights_manager->manage(_impl->original_weights);
     }
 
     if(!_impl->is_prepared)
@@ -125,11 +125,13 @@ void CLFullyConnectedLayer::prepare()
         // 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
+            // Ensure that b gets marked as unused (memory released) only after the last function which uses b also finishes its prepare
+            // This is for cases where multiple functions share the same b (weights)
+            // Therefore when a function marks original b as unused, we pre-mark it in weights manager, and mark it back to used so that it doesn't get released before its last reference
             const ITensor *original_b = _impl->original_weights;
             if(!original_b->is_used())
             {
-                _impl->weights_manager->mark_as_unused(original_b);
+                _impl->weights_manager->pre_mark_as_unused(original_b);
             }
             _impl->original_weights->mark_as_used();
             _impl->weights_manager->release(_impl->original_weights);
diff --git a/src/runtime/IWeightsManager.cpp b/src/runtime/IWeightsManager.cpp
index 8f27795285..373c50c73d 100644
--- a/src/runtime/IWeightsManager.cpp
+++ b/src/runtime/IWeightsManager.cpp
@@ -166,7 +166,7 @@ void IWeightsManager::release(const ITensor *weights)
     }
 }
 
-void IWeightsManager::mark_as_unused(const ITensor *weights)
+void IWeightsManager::pre_mark_as_unused(const ITensor *weights)
 {
     if(weights == nullptr || !are_weights_managed(weights))
     {
diff --git a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
index 84bc004ec4..cb7e2dc7ec 100644
--- a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
+++ b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
@@ -78,7 +78,7 @@ void NEFullyConnectedLayer::configure(const ITensor *input, const ITensor *weigh
 
     if(_impl->weights_manager != nullptr)
     {
-        _impl->weights_manager->manage(weights);
+        _impl->weights_manager->manage(_impl->original_weights);
     }
 
     _impl->aux_mem_req = _impl->op->workspace();
@@ -113,11 +113,13 @@ void NEFullyConnectedLayer::prepare()
         // 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
+            // Ensure that b gets marked as unused (memory released) only after the last function which uses b also finishes its prepare
+            // This is for cases where multiple functions share the same b (weights)
+            // Therefore when a function marks original b as unused, we pre-mark it in weights manager, and mark it back to used so that it doesn't get released before its last reference
             const ITensor *original_b = _impl->original_weights;
             if(!original_b->is_used())
             {
-                _impl->weights_manager->mark_as_unused(original_b);
+                _impl->weights_manager->pre_mark_as_unused(original_b);
             }
             _impl->original_weights->mark_as_used();
             _impl->weights_manager->release(_impl->original_weights);