From 73df9310f4e94e43597f283307e3cde0653d8bae Mon Sep 17 00:00:00 2001 From: Giorgio Arena Date: Mon, 16 Aug 2021 12:29:27 +0100 Subject: Address comments on avoiding releasing weights if used by multiple functions Signed-off-by: Giorgio Arena Change-Id: I0b59c5326f5fcbc322fbeb864197ea999de6bd56 Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6112 Tested-by: Arm Jenkins Reviewed-by: Gian Marco Iodice Reviewed-by: SiCong Li Comments-Addressed: Arm Jenkins --- arm_compute/runtime/IWeightsManager.h | 4 ++-- src/runtime/CL/functions/CLFullyConnectedLayer.cpp | 8 +++++--- src/runtime/IWeightsManager.cpp | 2 +- src/runtime/NEON/functions/NEFullyConnectedLayer.cpp | 8 +++++--- 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); -- cgit v1.2.1