diff --git a/runtime/onert/backend/train/MemoryManager.cc b/runtime/onert/backend/train/MemoryManager.cc index fd156fea231..64a665dd620 100644 --- a/runtime/onert/backend/train/MemoryManager.cc +++ b/runtime/onert/backend/train/MemoryManager.cc @@ -93,6 +93,42 @@ uint8_t *DisposableMemoryManager::getBuffer(const DisposableTensorIndex &ind) co return _mem_alloc->base() + mem_blk.offset; } +LayerScopeMemoryManager::LayerScopeMemoryManager() : _mem_planner{createMemoryPlanner()} +{ + // DO NOTHING +} + +basic::IMemoryPlanner *LayerScopeMemoryManager::createMemoryPlanner() +{ + auto planner_id = util::getConfigString(util::config::CPU_MEMORY_PLANNER); + return MemoryPlannerFactory::get().create(planner_id); +} + +void LayerScopeMemoryManager::allocate(void) +{ + _mem_alloc = std::make_shared(_mem_planner->capacity()); + assert(_mem_alloc->base()); +} + +uint8_t *LayerScopeMemoryManager::getBuffer(const LayerScopeTensorIndex &ind) const +{ + assert(_mem_planner->memory_plans().find(ind) != _mem_planner->memory_plans().end()); + const auto &mem_blk = _mem_planner->memory_plans().at(ind); + return _mem_alloc->base() + mem_blk.offset; +} + +void LayerScopeMemoryManager::deallocate(void) { _mem_alloc->release(); } + +void LayerScopeMemoryManager::claimPlan(const LayerScopeTensorIndex &ind, uint32_t size) +{ + _mem_planner->claim(ind, size); +} + +void LayerScopeMemoryManager::releasePlan(const LayerScopeTensorIndex &ind) +{ + _mem_planner->release(ind); +} + } // namespace train } // namespace backend } // namespace onert diff --git a/runtime/onert/backend/train/MemoryManager.h b/runtime/onert/backend/train/MemoryManager.h index 98e840bf7f7..8333c838bce 100644 --- a/runtime/onert/backend/train/MemoryManager.h +++ b/runtime/onert/backend/train/MemoryManager.h @@ -20,6 +20,7 @@ #include #include "DisposableTensorIndex.h" +#include "LayerScopeTensorIndex.h" namespace onert { @@ -67,7 +68,25 @@ class DisposableMemoryManager std::shared_ptr _mem_alloc; }; -// TODO: Add LayerScopeMemoryManager using MemoryPlannerFactory +class LayerScopeMemoryManager +{ +public: + LayerScopeMemoryManager(); + + void allocate(void); + uint8_t *getBuffer(const LayerScopeTensorIndex &ind) const; + void deallocate(void); + + void claimPlan(const LayerScopeTensorIndex &ind, uint32_t size); + void releasePlan(const LayerScopeTensorIndex &ind); + +private: + basic::IMemoryPlanner *createMemoryPlanner(); + +private: + std::shared_ptr> _mem_planner; + std::shared_ptr _mem_alloc; +}; } // namespace train } // namespace backend