Fix memory lifetime issue

B matrix was getting released while was still needed for a Native GEMM
execution

Resolves: COMPMID-4705

Signed-off-by: Georgios Pinitas <georgios.pinitas@arm.com>
Change-Id: I1e490abbe1feabd451d10fa785de21e725f1d6e0
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6001
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

2 files changed, 2 insertions, 31 deletions

diff --git a/src/runtime/CL/functions/CLGEMMLowpMatrixMultiplyCore.cpp b/src/runtime/CL/functions/CLGEMMLowpMatrixMultiplyCore.cpp
index bd31d47b4f..1ae2dfbad6 100644
--- a/src/runtime/CL/functions/CLGEMMLowpMatrixMultiplyCore.cpp
+++ b/src/runtime/CL/functions/CLGEMMLowpMatrixMultiplyCore.cpp
@@ -112,34 +112,8 @@ void CLGEMMLowpMatrixMultiplyCore::prepare()
     {
         _impl->op->prepare(_impl->run_pack);
 
-        auto has_reshape = std::find_if(_impl->aux_mem_req.begin(),
-                                        _impl->aux_mem_req.end(),
-                                        [](const MemoryInfo & m) -> bool { return m.lifetime == MemoryLifetime::Persistent; });
-
-        if(has_reshape != std::end(_impl->aux_mem_req))
-        {
-            _impl->b->mark_as_unused();
-        }
-        else
-        {
-            // Pack the B matrix to be used as the underlying GEMM performs no reshapes
-            _impl->run_pack.add_const_tensor(ACL_SRC_1, _impl->b);
-        }
-
         // Release temporary tensors that are only used in prepare stage
-        for(auto &ws : _impl->workspace_tensors)
-        {
-            const int slot = ws.slot;
-            for(auto &m : _impl->aux_mem_req)
-            {
-                if(m.slot == slot && m.lifetime == MemoryLifetime::Prepare)
-                {
-                    auto tensor = ws.tensor.get();
-                    tensor->allocator()->free();
-                    break;
-                }
-            }
-        }
+        release_temporaries(_impl->aux_mem_req, _impl->workspace_tensors);
 
         _impl->is_prepared = true;
     }
diff --git a/src/runtime/gpu/cl/operators/ClGemmLowpMatrixMultiplyCore.cpp b/src/runtime/gpu/cl/operators/ClGemmLowpMatrixMultiplyCore.cpp
index 82047ad2a5..64c8743f13 100644
--- a/src/runtime/gpu/cl/operators/ClGemmLowpMatrixMultiplyCore.cpp
+++ b/src/runtime/gpu/cl/operators/ClGemmLowpMatrixMultiplyCore.cpp
@@ -335,8 +335,6 @@ void ClGemmLowpMatrixMultiplyCore::configure(const CLCompileContext &compile_con
         {
             _run_output_stage = true;
 
-            // _memory_group.manage(&_mm_result_s32);
-
             if(_is_gemm_reshaped)
             {
                 _mm_reshaped_only_rhs_kernel->configure(compile_context, a, matrix_b, &_mm_result_s32, gemm_kernel_info);
@@ -726,12 +724,11 @@ void ClGemmLowpMatrixMultiplyCore::prepare(ITensorPack &tensors)
         {
             ITensorPack convert_to_qs8_pack = { { ACL_SRC, b }, { ACL_DST, rhs_qasymm8.get() } };
             CLScheduler::get().enqueue_op(*_weights_to_qasymm8, convert_to_qs8_pack, false);
+            b->mark_as_unused();
         }
 
         if(_is_gemm_reshaped && _reshape_b_only_on_first_run)
         {
-            ARM_COMPUTE_ERROR_ON(!b->is_used());
-
             // Run reshape kernel and mark original weights tensor as unused
             ITensorPack mtx_b_pack =
             {