diff options
Diffstat (limited to 'src/runtime/IWeightsManager.cpp')
| -rw-r--r-- | src/runtime/IWeightsManager.cpp | 70 |
1 files changed, 49 insertions, 21 deletions
diff --git a/src/runtime/IWeightsManager.cpp b/src/runtime/IWeightsManager.cpp index 081cd990f3..96287dcc49 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 * @@ -25,23 +25,27 @@ namespace arm_compute { -IWeightsManager::IWeightsManager() - : _managed_weights(), _managed_weights_parents() +IWeightsManager::IWeightsManager() : _managed_weights(), _managed_counter(), _managed_weights_parents() { } void IWeightsManager::manage(const ITensor *weights, ITransformWeights *parent) { - if(!are_weights_managed(weights)) + 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 // store the parent's link - if(parent != nullptr) + if (parent != nullptr) { - if(_managed_weights_parents.find(weights) == _managed_weights_parents.end()) + if (_managed_weights_parents.find(weights) == _managed_weights_parents.end()) { _managed_weights_parents[weights] = parent; } @@ -54,13 +58,13 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights // Find if I have the same weights with weights transform. If I do, don't run the reshape auto item = _managed_weights.find(weights); - bool perform_run{ true }; - ITensor *weights_tensor{ nullptr }; + bool perform_run{true}; + ITensor *weights_tensor{nullptr}; // Check if I already have the requested transform and I have run the reshape function - for(auto it : item->second) + for (auto it : item->second) { - if(it->is_reshape_run() && (it->uid() == weights_transform->uid())) + if (it->is_reshape_run() && (it->uid() == weights_transform->uid())) { weights_tensor = it->get_weights(); perform_run = false; @@ -68,7 +72,7 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights } } - if(perform_run) + if (perform_run) { weights_transform->run(); weights_tensor = weights_transform->get_weights(); @@ -76,10 +80,10 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights // Check if we can release memory from parent auto parent_item = _managed_weights_parents.find(weights); - if(parent_item != _managed_weights_parents.end()) + if (parent_item != _managed_weights_parents.end()) { int32_t refcount = parent_item->second->decrease_refcount(); - if(refcount == 0) + if (refcount == 0) { parent_item->second->release(); } @@ -87,20 +91,20 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights // Check top level weights. If all the transformations are done // mark the weights as unused - if(_managed_weights_parents.find(weights) == _managed_weights_parents.end()) + if (_managed_weights_parents.find(weights) == _managed_weights_parents.end()) { auto item = _managed_weights.find(weights); bool mark_as_unused = true; - for(auto it : item->second) + for (auto it : item->second) { - if(!it->is_reshape_run()) + if (!it->is_reshape_run()) { mark_as_unused = false; break; } } - if(mark_as_unused) + if (mark_as_unused) { weights->mark_as_unused(); } @@ -118,15 +122,15 @@ ITensor *IWeightsManager::acquire(const ITensor *weights, ITransformWeights *wei { ARM_COMPUTE_ERROR_ON_MSG(!are_weights_managed(weights), "Cannot acquire weights. Weights are not managed"); - ITensor *transformed_weights{ nullptr }; + ITensor *transformed_weights{nullptr}; auto item = _managed_weights.find(weights); // Check if I already have the requested transform. If I do, // increase the refcount of the transformed weights object and // reuse the tensor - for(auto it : item->second) + for (auto it : item->second) { - if(it->uid() == weights_transform->uid()) + if (it->uid() == weights_transform->uid()) { transformed_weights = it->get_weights(); it->increase_refcount(); @@ -134,7 +138,7 @@ ITensor *IWeightsManager::acquire(const ITensor *weights, ITransformWeights *wei } } - if(transformed_weights == nullptr) + if (transformed_weights == nullptr) { transformed_weights = weights_transform->get_weights(); weights_transform->increase_refcount(); @@ -146,4 +150,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::pre_mark_as_unused(const ITensor *weights) +{ + if (weights == nullptr || !are_weights_managed(weights)) + { + return; + } + + _managed_counter[weights].is_unused = true; +} } // namespace arm_compute |