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>
diff --git a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
index 84bc004..cb7e2dc 100644
--- a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
+++ b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp
@@ -78,7 +78,7 @@
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 @@
// 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);