COMPMID-1605: API alignment for the MemoryManager with ARMNN

Change-Id: Iac6a95ba7f388e65b7f1c8865c3e9bf289b233ea
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/155490
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Tested-by: bsgcomp <bsgcomp@arm.com>

15 files changed, 111 insertions, 117 deletions

diff --git a/arm_compute/graph/GraphContext.h b/arm_compute/graph/GraphContext.h
index b77eb1404d..21ba6df785 100644
--- a/arm_compute/graph/GraphContext.h
+++ b/arm_compute/graph/GraphContext.h
@@ -42,6 +42,7 @@ struct MemoryManagerContext
     std::shared_ptr<arm_compute::IMemoryManager> intra_mm    = { nullptr };             /**< Intra-function memory manager */
     std::shared_ptr<arm_compute::IMemoryManager> cross_mm    = { nullptr };             /**< Cross-function memory manager */
     std::shared_ptr<arm_compute::IMemoryGroup>   cross_group = { nullptr };             /**< Cross-function memory group */
+    IAllocator                                  *allocator   = { nullptr };             /**< Backend allocator to use */
 };
 
 /** Graph context **/
diff --git a/arm_compute/runtime/IMemoryManager.h b/arm_compute/runtime/IMemoryManager.h
index 00aa566a50..9280b309e0 100644
--- a/arm_compute/runtime/IMemoryManager.h
+++ b/arm_compute/runtime/IMemoryManager.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017 ARM Limited.
+ * Copyright (c) 2017-2018 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -31,6 +31,8 @@
 
 namespace arm_compute
 {
+// Forward declarations
+class IAllocator;
 class IMemoryGroup;
 
 /** Memory manager interface to handle allocations of backing memory */
@@ -49,8 +51,19 @@ public:
      * @return The pool manager
      */
     virtual IPoolManager *pool_manager() = 0;
-    /** Finalize memory manager */
-    virtual void finalize() = 0;
+    /** Populates the pool manager with the given number of pools
+     *
+     * @pre Pool manager must be empty
+     *
+     * @param[in] allocator Allocator to use for the backing allocations
+     * @param[in] num_pools Number of pools to create
+     */
+    virtual void populate(IAllocator &allocator, size_t num_pools) = 0;
+    /** Clears the pool manager
+     *
+     * @pre All pools must be unoccupied
+     */
+    virtual void clear() = 0;
 };
 } // arm_compute
 #endif /*__ARM_COMPUTE_IMEMORYMANAGER_H__ */
diff --git a/arm_compute/runtime/IPoolManager.h b/arm_compute/runtime/IPoolManager.h
index 4cc3c07b5f..6f5af3d499 100644
--- a/arm_compute/runtime/IPoolManager.h
+++ b/arm_compute/runtime/IPoolManager.h
@@ -53,6 +53,16 @@ public:
      * @param[in] pool Pool to be managed
      */
     virtual void register_pool(std::unique_ptr<IMemoryPool> pool) = 0;
+    /** Releases a free pool from the managed pools
+     *
+     * @return The released pool in case a free pool existed else nullptr
+     */
+    virtual std::unique_ptr<IMemoryPool> release_pool() = 0;
+    /** Clears all pools managed by the pool manager
+     *
+     *  @pre All pools must be unoccupied
+     */
+    virtual void clear_pools() = 0;
     /** Returns the total number of pools managed by the pool manager
      *
      * @return Number of managed pools
diff --git a/arm_compute/runtime/MemoryManagerOnDemand.h b/arm_compute/runtime/MemoryManagerOnDemand.h
index ad4b831e1f..43f557effc 100644
--- a/arm_compute/runtime/MemoryManagerOnDemand.h
+++ b/arm_compute/runtime/MemoryManagerOnDemand.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017 ARM Limited.
+ * Copyright (c) 2017-2018 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -21,23 +21,19 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */
-#ifndef __ARM_COMPUTE_MEMORYMANAGERONDEMAND_H__
-#define __ARM_COMPUTE_MEMORYMANAGERONDEMAND_H__
+#ifndef __ARM_COMPUTE_MEMORY_MANAGER_ON_DEMAND_H__
+#define __ARM_COMPUTE_MEMORY_MANAGER_ON_DEMAND_H__
 
 #include "arm_compute/runtime/IMemoryManager.h"
 
-#include "IAllocator.h"
 #include "arm_compute/runtime/ILifetimeManager.h"
 #include "arm_compute/runtime/IMemoryGroup.h"
 #include "arm_compute/runtime/IPoolManager.h"
 
 #include <memory>
-#include <set>
 
 namespace arm_compute
 {
-class IAllocator;
-
 /** On-demand memory manager */
 class MemoryManagerOnDemand : public IMemoryManager
 {
@@ -52,33 +48,16 @@ public:
     MemoryManagerOnDemand(MemoryManagerOnDemand &&) = default;
     /** Allow instances of this class to be moved */
     MemoryManagerOnDemand &operator=(MemoryManagerOnDemand &&) = default;
-    /** Sets the number of pools to create
-     *
-     * @param[in] num_pools Number of pools
-     */
-    void set_num_pools(unsigned int num_pools);
-    /** Sets the allocator to be used for configuring the pools
-     *
-     * @param[in] allocator Allocator to use
-     */
-    void set_allocator(IAllocator *allocator);
-    /** Checks if the memory manager has been finalized
-     *
-     * @return True if the memory manager has been finalized else false
-     */
-    bool is_finalized() const;
 
     // Inherited methods overridden:
     ILifetimeManager *lifetime_manager() override;
     IPoolManager     *pool_manager() override;
-    void              finalize() override;
+    void populate(IAllocator &allocator, size_t num_pools) override;
+    void clear() override;
 
 private:
     std::shared_ptr<ILifetimeManager> _lifetime_mgr; /**< Lifetime manager */
     std::shared_ptr<IPoolManager>     _pool_mgr;     /**< Memory pool manager */
-    IAllocator                       *_allocator;    /**< Allocator used for backend allocations */
-    bool                              _is_finalized; /**< Flag that notes if the memory manager has been finalized */
-    unsigned int                      _num_pools;    /**< Number of pools to create */
 };
 } // arm_compute
-#endif /*__ARM_COMPUTE_MEMORYMANAGERONDEMAND_H__ */
+#endif /*__ARM_COMPUTE_MEMORY_MANAGER_ON_DEMAND_H__ */
diff --git a/arm_compute/runtime/PoolManager.h b/arm_compute/runtime/PoolManager.h
index 509080532a..4f0f5384e0 100644
--- a/arm_compute/runtime/PoolManager.h
+++ b/arm_compute/runtime/PoolManager.h
@@ -56,7 +56,9 @@ public:
     IMemoryPool *lock_pool() override;
     void unlock_pool(IMemoryPool *pool) override;
     void register_pool(std::unique_ptr<IMemoryPool> pool) override;
-    size_t num_pools() const override;
+    std::unique_ptr<IMemoryPool> release_pool() override;
+    void                         clear_pools() override;
+    size_t                       num_pools() const override;
 
 private:
     std::list<std::unique_ptr<IMemoryPool>> _free_pools;     /**< List of free pools */
diff --git a/docs/01_library.dox b/docs/01_library.dox
index bd4b300d7d..189602046b 100644
--- a/docs/01_library.dox
+++ b/docs/01_library.dox
@@ -342,8 +342,6 @@ Requesting backing memory for a specific group can be done using @ref IMemoryGro
 - @ref ILifetimeManager that keeps track of the lifetime of the registered objects of the memory groups and given an @ref IAllocator creates an appropriate memory pool that fulfils the memory requirements of all the registered memory groups.
 - @ref IPoolManager that safely manages the registered memory pools.
 
-@note @ref IMemoryManager::finalize should be called once the configuration of all the memory groups, kernels and functions is done, so that the memory manager can allocate the appropriate backing memory.
-
 @note @ref BlobLifetimeManager is currently implemented which models the memory requirements as a vector of distinct memory blobs.
 
 @subsection S4_7_2_working_with_memory_manager Working with the Memory Manager
@@ -385,11 +383,9 @@ tmp3.allocator()->allocate();       // Flag that the lifetime of object tmp3 has
 
 @warning The configuration step should be done sequentially by a single thread so that all the lifetimes are captured correclty.
 
-When configuration of all the operations is finished then the memory manager have to be finalized:
+When configuration of all the operations is finished then the memory manager have to be populated:
 @code{.cpp}
-mm->set_allocator(&allocator); // Set allocator to use
-mm->set_set_num_pools(2);      // Set number of pools to create in case parallel operations can be run
-mm->finalize();                // Finalize memory manager (Object lifetime check, Memory pool creation etc)
+mm->populate(&allocator), 2 /* num_pools */); // Populate memory manager pools
 @endcode
 
 Finally, during execution of the pipeline the memory of the appropriate memory group should be requested before running:
@@ -422,10 +418,8 @@ CLConvolutionLayer conv1(mm), conv2(mm);
 conv1.configure(...);
 conv2.configure(...);
 
-// Finalize memory manager
-mm->set_allocator(&allocator); // Set allocator to use
-mm->set_set_num_pools(1);      // Set number of pools to create in case parallel operations can be run
-mm->finalize();                // Finalize memory manager (Object lifetime check, Memory pool creation etc)
+// Populate memory manager
+mm->populate(&allocator), 1 /* num_pools */); // Populate memory manager pools
 
 // Run layers (Memory will be recycled for internal buffers for conv1 and conv2
 conv1.run();
diff --git a/examples/neon_cnn.cpp b/examples/neon_cnn.cpp
index 1df81256b9..6f26af7af4 100644
--- a/examples/neon_cnn.cpp
+++ b/examples/neon_cnn.cpp
@@ -205,28 +205,11 @@ public:
 
         /* -----------------------End: [Allocate tensors] */
 
-        // Finalize layers memory manager
+        // Populate the layers manager. (Validity checks, memory allocations etc)
+        mm_layers->populate(allocator, 1 /* num_pools */);
 
-        // Set allocator that the memory manager will use
-        mm_layers->set_allocator(&allocator);
-
-        // Number of pools that the manager will create. This specifies how many layers you want to run in parallel
-        mm_layers->set_num_pools(1);
-
-        // Finalize the manager. (Validity checks, memory allocations etc)
-        mm_layers->finalize();
-
-        // Finalize transitions memory manager
-
-        // Set allocator that the memory manager will use
-        mm_transitions->set_allocator(&allocator);
-
-        // Number of pools that the manager will create. This specifies how many models we can run in parallel.
-        // Setting to 2 as we need one for the input and one for the output at any given time
-        mm_transitions->set_num_pools(2);
-
-        // Finalize the manager. (Validity checks, memory allocations etc)
-        mm_transitions->finalize();
+        // Populate the transitions manager. (Validity checks, memory allocations etc)
+        mm_transitions->populate(allocator, 2 /* num_pools */);
 
         return true;
     }
diff --git a/src/graph/GraphContext.cpp b/src/graph/GraphContext.cpp
index 5f33ed3537..037b40b68b 100644
--- a/src/graph/GraphContext.cpp
+++ b/src/graph/GraphContext.cpp
@@ -25,6 +25,7 @@
 
 #include "arm_compute/graph.h"
 #include "arm_compute/graph/Utils.h"
+#include "arm_compute/graph/backends/BackendRegistry.h"
 
 namespace arm_compute
 {
@@ -75,17 +76,20 @@ std::map<Target, MemoryManagerContext> &GraphContext::memory_managers()
 
 void GraphContext::finalize()
 {
+    const size_t num_pools = 1;
     for(auto &mm_obj : _memory_managers)
     {
+        ARM_COMPUTE_ERROR_ON(!mm_obj.second.allocator);
+
         // Finalize intra layer memory manager
         if(mm_obj.second.intra_mm != nullptr)
         {
-            mm_obj.second.intra_mm->finalize();
+            mm_obj.second.intra_mm->populate(*mm_obj.second.allocator, num_pools);
         }
         // Finalize cross layer memory manager
         if(mm_obj.second.cross_mm != nullptr)
         {
-            mm_obj.second.cross_mm->finalize();
+            mm_obj.second.cross_mm->populate(*mm_obj.second.allocator, num_pools);
         }
     }
 }
diff --git a/src/graph/backends/CL/CLDeviceBackend.cpp b/src/graph/backends/CL/CLDeviceBackend.cpp
index f35daf4ae5..ae7f0a50b3 100644
--- a/src/graph/backends/CL/CLDeviceBackend.cpp
+++ b/src/graph/backends/CL/CLDeviceBackend.cpp
@@ -127,6 +127,7 @@ void CLDeviceBackend::setup_backend_context(GraphContext &ctx)
         mm_ctx.intra_mm    = create_memory_manager(MemoryManagerAffinity::Buffer);
         mm_ctx.cross_mm    = create_memory_manager(MemoryManagerAffinity::Buffer);
         mm_ctx.cross_group = std::make_shared<CLMemoryGroup>(mm_ctx.cross_mm);
+        mm_ctx.allocator   = _allocator.get();
 
         ctx.insert_memory_management_ctx(std::move(mm_ctx));
     }
@@ -195,8 +196,6 @@ std::shared_ptr<arm_compute::IMemoryManager> CLDeviceBackend::create_memory_mana
     auto pool_mgr     = std::make_shared<PoolManager>();
     auto mm           = std::make_shared<MemoryManagerOnDemand>(lifetime_mgr, pool_mgr);
 
-    mm->set_allocator(_allocator.get());
-
     return mm;
 }
 } // namespace backends
diff --git a/src/graph/backends/GLES/GCDeviceBackend.cpp b/src/graph/backends/GLES/GCDeviceBackend.cpp
index ec3cf4f21e..5f0bf3f263 100644
--- a/src/graph/backends/GLES/GCDeviceBackend.cpp
+++ b/src/graph/backends/GLES/GCDeviceBackend.cpp
@@ -86,6 +86,7 @@ void GCDeviceBackend::setup_backend_context(GraphContext &ctx)
         mm_ctx.intra_mm    = create_memory_manager(MemoryManagerAffinity::Buffer);
         mm_ctx.cross_mm    = create_memory_manager(MemoryManagerAffinity::Buffer);
         mm_ctx.cross_group = std::make_shared<GCMemoryGroup>(mm_ctx.cross_mm);
+        mm_ctx.allocator   = &_allocator;
 
         ctx.insert_memory_management_ctx(std::move(mm_ctx));
     }
@@ -151,8 +152,6 @@ std::shared_ptr<arm_compute::IMemoryManager> GCDeviceBackend::create_memory_mana
     auto pool_mgr     = std::make_shared<PoolManager>();
     auto mm           = std::make_shared<MemoryManagerOnDemand>(lifetime_mgr, pool_mgr);
 
-    mm->set_allocator(&_allocator);
-
     return mm;
 }
 } // namespace backends
diff --git a/src/graph/backends/NEON/NEDeviceBackend.cpp b/src/graph/backends/NEON/NEDeviceBackend.cpp
index 5fc44d0c68..23ced2fef3 100644
--- a/src/graph/backends/NEON/NEDeviceBackend.cpp
+++ b/src/graph/backends/NEON/NEDeviceBackend.cpp
@@ -86,6 +86,7 @@ void NEDeviceBackend::setup_backend_context(GraphContext &ctx)
         mm_ctx.intra_mm    = create_memory_manager(MemoryManagerAffinity::Offset);
         mm_ctx.cross_mm    = create_memory_manager(MemoryManagerAffinity::Offset);
         mm_ctx.cross_group = std::make_shared<MemoryGroup>(mm_ctx.cross_mm);
+        mm_ctx.allocator   = &_allocator;
 
         ctx.insert_memory_management_ctx(std::move(mm_ctx));
     }
@@ -156,8 +157,6 @@ std::shared_ptr<arm_compute::IMemoryManager> NEDeviceBackend::create_memory_mana
     auto pool_mgr = std::make_shared<PoolManager>();
     auto mm       = std::make_shared<MemoryManagerOnDemand>(lifetime_mgr, pool_mgr);
 
-    mm->set_allocator(&_allocator);
-
     return mm;
 }
 } // namespace backends
diff --git a/src/runtime/MemoryManagerOnDemand.cpp b/src/runtime/MemoryManagerOnDemand.cpp
index 4dfa28bddb..d9803a8caf 100644
--- a/src/runtime/MemoryManagerOnDemand.cpp
+++ b/src/runtime/MemoryManagerOnDemand.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, 2017 ARM Limited.
+ * Copyright (c) 2016-2018 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -29,33 +29,15 @@
 
 #include <memory>
 
-using namespace arm_compute;
-
+namespace arm_compute
+{
 MemoryManagerOnDemand::MemoryManagerOnDemand(std::shared_ptr<ILifetimeManager> lifetime_manager, std::shared_ptr<IPoolManager> pool_manager)
-    : _lifetime_mgr(std::move(lifetime_manager)), _pool_mgr(std::move(pool_manager)), _allocator(nullptr), _is_finalized(false), _num_pools(1)
+    : _lifetime_mgr(std::move(lifetime_manager)), _pool_mgr(std::move(pool_manager))
 {
     ARM_COMPUTE_ERROR_ON_MSG(!_lifetime_mgr, "Lifetime manager not specified correctly!");
     ARM_COMPUTE_ERROR_ON_MSG(!_pool_mgr, "Pool manager not specified correctly!");
 }
 
-bool MemoryManagerOnDemand::is_finalized() const
-{
-    return _is_finalized;
-}
-
-void MemoryManagerOnDemand::set_num_pools(unsigned int num_pools)
-{
-    ARM_COMPUTE_ERROR_ON(num_pools == 0);
-    _num_pools = num_pools;
-}
-
-void MemoryManagerOnDemand::set_allocator(IAllocator *allocator)
-{
-    ARM_COMPUTE_ERROR_ON_MSG(is_finalized(), "Memory manager is already finalized!");
-    ARM_COMPUTE_ERROR_ON(allocator == nullptr);
-    _allocator = allocator;
-}
-
 ILifetimeManager *MemoryManagerOnDemand::lifetime_manager()
 {
     return _lifetime_mgr.get();
@@ -66,23 +48,26 @@ IPoolManager *MemoryManagerOnDemand::pool_manager()
     return _pool_mgr.get();
 }
 
-void MemoryManagerOnDemand::finalize()
+void MemoryManagerOnDemand::populate(arm_compute::IAllocator &allocator, size_t num_pools)
 {
-    ARM_COMPUTE_ERROR_ON_MSG(is_finalized(), "Memory manager is already finalized!");
     ARM_COMPUTE_ERROR_ON(!_lifetime_mgr);
     ARM_COMPUTE_ERROR_ON(!_pool_mgr);
-    ARM_COMPUTE_ERROR_ON_MSG(!_lifetime_mgr->are_all_finalized(), "All the objects have not been finalized! ");
-    ARM_COMPUTE_ERROR_ON(_allocator == nullptr);
+    ARM_COMPUTE_ERROR_ON_MSG(!_lifetime_mgr->are_all_finalized(), "All the objects have not been finalized!");
+    ARM_COMPUTE_ERROR_ON_MSG(_pool_mgr->num_pools() != 0, "Pool manager already contains pools!");
 
     // Create pools
-    auto pool_template = _lifetime_mgr->create_pool(_allocator);
-    for(int i = _num_pools; i > 1; --i)
+    auto pool_template = _lifetime_mgr->create_pool(&allocator);
+    for(int i = num_pools; i > 1; --i)
     {
         auto pool = pool_template->duplicate();
         _pool_mgr->register_pool(std::move(pool));
     }
     _pool_mgr->register_pool(std::move(pool_template));
+}
 
-    // Set finalized to true
-    _is_finalized = true;
+void MemoryManagerOnDemand::clear()
+{
+    ARM_COMPUTE_ERROR_ON_MSG(!_pool_mgr, "Pool manager not specified correctly!");
+    _pool_mgr->clear_pools();
 }
+} //namespace arm_compute
diff --git a/src/runtime/PoolManager.cpp b/src/runtime/PoolManager.cpp
index 293241d6f4..5ec2ce92c1 100644
--- a/src/runtime/PoolManager.cpp
+++ b/src/runtime/PoolManager.cpp
@@ -73,6 +73,36 @@ void PoolManager::register_pool(std::unique_ptr<IMemoryPool> pool)
     _sem = arm_compute::support::cpp14::make_unique<arm_compute::Semaphore>(_free_pools.size());
 }
 
+std::unique_ptr<IMemoryPool> PoolManager::release_pool()
+{
+    std::lock_guard<arm_compute::Mutex> lock(_mtx);
+    ARM_COMPUTE_ERROR_ON_MSG(!_occupied_pools.empty(), "All pools should be free in order to release one!");
+
+    if(!_free_pools.empty())
+    {
+        std::unique_ptr<IMemoryPool> pool = std::move(_free_pools.front());
+        ARM_COMPUTE_ERROR_ON(_free_pools.front() != nullptr);
+        _free_pools.pop_front();
+
+        // Update semaphore
+        _sem = arm_compute::support::cpp14::make_unique<arm_compute::Semaphore>(_free_pools.size());
+
+        return pool;
+    }
+
+    return nullptr;
+}
+
+void PoolManager::clear_pools()
+{
+    std::lock_guard<arm_compute::Mutex> lock(_mtx);
+    ARM_COMPUTE_ERROR_ON_MSG(!_occupied_pools.empty(), "All pools should be free in order to clear the PoolManager!");
+    _free_pools.clear();
+
+    // Update semaphore
+    _sem = nullptr;
+}
+
 size_t PoolManager::num_pools() const
 {
     std::lock_guard<arm_compute::Mutex> lock(_mtx);
diff --git a/tests/validation/NEON/UNIT/MemoryManager.cpp b/tests/validation/NEON/UNIT/MemoryManager.cpp
index 27b05ed55a..431254cf94 100644
--- a/tests/validation/NEON/UNIT/MemoryManager.cpp
+++ b/tests/validation/NEON/UNIT/MemoryManager.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017 ARM Limited.
+ * Copyright (c) 2017-2018 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -74,10 +74,8 @@ TEST_CASE(BlobMemoryManagerSimpleWithinFunctionLevel, framework::DatasetMode::AL
     ARM_COMPUTE_EXPECT(!dst.info()->is_resizable(), framework::LogLevel::ERRORS);
 
     // Finalize memory manager
-    mm->set_allocator(&allocator);
-    mm->set_num_pools(1);
-    mm->finalize();
-    ARM_COMPUTE_EXPECT(mm->is_finalized(), framework::LogLevel::ERRORS);
+    mm->populate(allocator, 1 /* num_pools */);
+    ARM_COMPUTE_EXPECT(mm->pool_manager()->num_pools() == 1, framework::LogLevel::ERRORS);
     ARM_COMPUTE_EXPECT(mm->lifetime_manager()->are_all_finalized(), framework::LogLevel::ERRORS);
 
     // Fill tensors
@@ -86,6 +84,10 @@ TEST_CASE(BlobMemoryManagerSimpleWithinFunctionLevel, framework::DatasetMode::AL
     // Compute functions
     norm_layer_1.run();
     norm_layer_2.run();
+
+    // Clear manager
+    mm->clear();
+    ARM_COMPUTE_EXPECT(mm->pool_manager()->num_pools() == 0, framework::LogLevel::ERRORS);
 }
 
 TEST_SUITE_END()
diff --git a/tests/validation/fixtures/UNIT/MemoryManagerFixture.h b/tests/validation/fixtures/UNIT/MemoryManagerFixture.h
index d8e2b0b427..01f9092da5 100644
--- a/tests/validation/fixtures/UNIT/MemoryManagerFixture.h
+++ b/tests/validation/fixtures/UNIT/MemoryManagerFixture.h
@@ -99,11 +99,9 @@ protected:
         dst.allocator()->allocate();
 
         // Finalize memory manager
-        mm->set_allocator(&_allocator);
-        mm->set_num_pools(1);
-        mm->finalize();
-        ARM_COMPUTE_EXPECT(mm->is_finalized(), framework::LogLevel::ERRORS);
+        mm->populate(_allocator, 1 /* num_pools */);
         ARM_COMPUTE_EXPECT(mm->lifetime_manager()->are_all_finalized(), framework::LogLevel::ERRORS);
+        ARM_COMPUTE_EXPECT(mm->pool_manager()->num_pools() == 1, framework::LogLevel::ERRORS);
 
         // Fill tensors
         fill(AccessorType(src), 0);
@@ -207,11 +205,9 @@ protected:
         dst.allocator()->allocate();
 
         // Finalize memory manager
-        mm->set_allocator(&allocator);
-        mm->set_num_pools(1);
-        mm->finalize();
-        ARM_COMPUTE_EXPECT(mm->is_finalized(), framework::LogLevel::ERRORS);
+        mm->populate(_allocator, 1 /* num_pools */);
         ARM_COMPUTE_EXPECT(mm->lifetime_manager()->are_all_finalized(), framework::LogLevel::ERRORS);
+        ARM_COMPUTE_EXPECT(mm->pool_manager()->num_pools() == 1, framework::LogLevel::ERRORS);
 
         // Fill tensors (1st iteration)
         fill(AccessorType(src), 0);
@@ -343,11 +339,9 @@ protected:
         dst.allocator()->allocate();
 
         // Finalize memory manager
-        mm->set_allocator(&allocator);
-        mm->set_num_pools(1);
-        mm->finalize();
-        ARM_COMPUTE_EXPECT(mm->is_finalized(), framework::LogLevel::ERRORS);
+        mm->populate(_allocator, 1 /* num_pools */);
         ARM_COMPUTE_EXPECT(mm->lifetime_manager()->are_all_finalized(), framework::LogLevel::ERRORS);
+        ARM_COMPUTE_EXPECT(mm->pool_manager()->num_pools() == 1, framework::LogLevel::ERRORS);
 
         // Fill tensors (1st iteration)
         fill(AccessorType(src), 0);