src/armnn/memory/PoolManager.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // See LICENSE file in the project root for full license information.
   4 //
   5 #include "IMemoryPool.hpp"
   6 #include "PoolManager.hpp"
   7
   8 #include "boost/assert.hpp"
   9 #include "boost/polymorphic_cast.hpp"
  10
  11 #include <algorithm>
  12
  13 namespace armnn
  14 {
  15
  16 PoolManager::PoolManager()
  17         : m_FreePools()
  18         , m_OccupiedPools()
  19         , m_Semaphore()
  20         , m_Mutex()
  21 {}
  22
  23 arm_compute::IMemoryPool *PoolManager::lock_pool()
  24 {
  25     BOOST_ASSERT_MSG(!(m_FreePools.empty() && m_OccupiedPools.empty()), "Haven't setup any pools");
  26
  27     m_Semaphore->wait();
  28     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  29
  30     BOOST_ASSERT_MSG(!m_FreePools.empty(), "Empty pool must exist as semaphore has been signalled");
  31     m_OccupiedPools.splice(std::begin(m_OccupiedPools), m_FreePools, std::begin(m_FreePools));
  32
  33     return m_OccupiedPools.front().get();
  34 }
  35
  36 void PoolManager::unlock_pool(arm_compute::IMemoryPool *pool)
  37 {
  38     BOOST_ASSERT_MSG(!(m_FreePools.empty() && m_OccupiedPools.empty()), "Haven't setup any pools!");
  39
  40     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  41
  42     auto it = std::find_if(
  43             std::begin(m_OccupiedPools),
  44             std::end(m_OccupiedPools),
  45             [pool](const std::unique_ptr<arm_compute::IMemoryPool> &poolIterator)
  46             {
  47                 return poolIterator.get() == pool;
  48             }
  49     );
  50
  51     BOOST_ASSERT_MSG(it != std::end(m_OccupiedPools), "Pool to be unlocked couldn't be found");
  52     m_FreePools.splice(std::begin(m_FreePools), m_OccupiedPools, it);
  53     m_Semaphore->signal();
  54 }
  55
  56 void PoolManager::register_pool(std::unique_ptr<arm_compute::IMemoryPool> pool)
  57 {
  58     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  59     BOOST_ASSERT_MSG(m_OccupiedPools.empty(), "All pools should be free in order to register a new one");
  60
  61     // Set pool
  62     m_FreePools.push_front(std::move(pool));
  63
  64     // Update semaphore
  65     m_Semaphore = std::make_unique<arm_compute::Semaphore>(m_FreePools.size());
  66 }
  67
  68 size_t PoolManager::num_pools() const
  69 {
  70     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  71
  72     return m_FreePools.size() + m_OccupiedPools.size();
  73 }
  74
  75 void PoolManager::AllocatePools()
  76 {
  77     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  78
  79     for (auto& pool : m_FreePools)
  80     {
  81         boost::polymorphic_downcast<IMemoryPool*>(pool.get())->AllocatePool();
  82     }
  83
  84     for (auto& pool : m_OccupiedPools)
  85     {
  86         boost::polymorphic_downcast<IMemoryPool*>(pool.get())->AllocatePool();
  87     }
  88 }
  89
  90 void PoolManager::ReleasePools()
  91 {
  92     std::lock_guard<arm_compute::Mutex> lock(m_Mutex);
  93
  94     for (auto& pool : m_FreePools)
  95     {
  96         boost::polymorphic_downcast<IMemoryPool*>(pool.get())->ReleasePool();
  97     }
  98
  99     for (auto& pool : m_OccupiedPools)
 100     {
 101         boost::polymorphic_downcast<IMemoryPool*>(pool.get())->ReleasePool();
 102     }
 103 }
 104
 105 } //namespace armnn