[neurun] Revise MemoryAllocator (#3386)

- MemoryBlock -> Block
- base_offset -> offset
- Remove mem_id concept
- Introduce graph::operand::Index instead of uint32_t
- Apply above changes to symbols of MemoryAllocator
- Apply above changes to tests of MemoryAllocator
- Append Logging

Signed-off-by: Yongseop Kim <yons.kim@samsung.com>

diff --git a/runtimes/neurun/src/backend/cpu/MemoryAllocator.cc b/runtimes/neurun/src/backend/cpu/MemoryAllocator.cc
index 41b36bd..889f94c 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/MemoryAllocator.cc
+++ b/runtimes/neurun/src/backend/cpu/MemoryAllocator.cc
@@ -15,6 +15,7 @@
  */
 
 #include "MemoryAllocator.h"
+#include "logging.h"
 #include <cassert>
 
 namespace neurun
@@ -30,36 +31,32 @@ BumpAllocator::~BumpAllocator()
     delete[] _base;
 }
 
-uint32_t BumpAllocator::allocate(size_t size)
+Block BumpAllocator::allocate(const graph::operand::Index &index, size_t size)
 {
   assert(size != 0);
 
-  MemoryBlock blk{_curr_pos, size};
-  _mem_blk_map[_mem_idx] = blk;
-  _curr_pos += size;
-  return _mem_idx++;
+  Block blk{_pos, size};
+  _pos += size;
+
+  VERBOSE(BP_ALLOC) << "alloc(#" << index.value() << "): " << blk.offset << ", " << blk.size
+                    << std::endl;
+
+  return blk;
 }
 
 void BumpAllocator::finalize()
 {
-  assert(!_base && !_reserved_size && _curr_pos != 0);
-  _reserved_size = _curr_pos;
-  _base = new uint8_t[_reserved_size];
-}
+  assert(!_base && _pos != 0);
 
-void BumpAllocator::free(uint32_t)
-{
-  assert(_base && _reserved_size > 0);
+  VERBOSE(BP_ALLOC) << "final position: " << _pos << std::endl;
 
-  // DO NOTHING
-  // In the case of this BumpAllocator, ignore the case like where reallocations are needed
+  _base = new uint8_t[_pos];
 }
 
-MemoryBlock BumpAllocator::getMemoryBlock(uint32_t mem_id)
+void BumpAllocator::free(const graph::operand::Index &index)
 {
-  assert(_mem_blk_map.size() > 0);
-
-  return _mem_blk_map[mem_id];
+  VERBOSE(BP_ALLOC) << "free(#" << index.value() << "): "
+                    << "NOTHING does" << std::endl;
 }
 
 } // namespace cpu

diff --git a/runtimes/neurun/src/backend/cpu/MemoryAllocator.h b/runtimes/neurun/src/backend/cpu/MemoryAllocator.h
index e191839..d949cf7 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/MemoryAllocator.h
+++ b/runtimes/neurun/src/backend/cpu/MemoryAllocator.h
@@ -17,7 +17,7 @@
 #ifndef __NEURUN_BACKEND_CPU_MEMORY_ALLOCATOR_H__
 #define __NEURUN_BACKEND_CPU_MEMORY_ALLOCATOR_H__
 
-#include <unordered_map>
+#include "graph/operand/Index.h"
 
 namespace neurun
 {
@@ -26,38 +26,33 @@ namespace backend
 namespace cpu
 {
 
-struct MemoryBlock
+struct Block
 {
-  uint32_t base_offset;
+  uint32_t offset;
   uint32_t size;
 };
 
 struct IMemoryAllocator
 {
   virtual ~IMemoryAllocator() = default;
-  virtual uint32_t allocate(size_t) = 0;
-  virtual void free(uint32_t) = 0;
+  virtual Block allocate(const graph::operand::Index &, size_t) = 0;
+  virtual void free(const graph::operand::Index &) = 0;
   virtual void finalize() = 0;
   virtual uint8_t *base() const = 0;
-  virtual MemoryBlock getMemoryBlock(uint32_t) = 0;
 };
 
 class BumpAllocator : public IMemoryAllocator
 {
 public:
   virtual ~BumpAllocator() override;
-  virtual uint32_t allocate(size_t size) override;
-  virtual void free(uint32_t mem_id) override;
+  virtual Block allocate(const graph::operand::Index &index, size_t size) override;
+  virtual void free(const graph::operand::Index &index) override;
   virtual void finalize() override;
   virtual uint8_t *base() const override { return _base; }
-  virtual MemoryBlock getMemoryBlock(uint32_t mem_id) override;
 
 private:
   uint8_t *_base = nullptr;
-  uint32_t _reserved_size = 0;
-  uint32_t _curr_pos = 0;
-  uint32_t _mem_idx = 0;
-  std::unordered_map<uint32_t, MemoryBlock> _mem_blk_map;
+  uint32_t _pos = 0;
 };
 
 } // namespace cpu

diff --git a/runtimes/neurun/src/backend/cpu/TensorBuilder.cc b/runtimes/neurun/src/backend/cpu/TensorBuilder.cc
index 5eee633..4bb21b6 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/TensorBuilder.cc
+++ b/runtimes/neurun/src/backend/cpu/TensorBuilder.cc
@@ -56,8 +56,8 @@ void TensorBuilder::prepare(void)
     // If we do not make tensor here currently, stages would cause segment fault
 
     const auto size = info.total_size(); // NOTE This size may not be accurate
-    auto mem_id = _mem_alloc->allocate(size);
-    _tensor_mem_map[ind] = mem_id;
+    auto mem_blk = _mem_alloc->allocate(ind, size);
+    _tensor_mem_map[ind] = mem_blk;
   }
   assert(_tensor_info_map.size() == _tensor_mem_map.size());
 
@@ -73,10 +73,9 @@ void TensorBuilder::prepare(void)
   for (auto &entry : _tensor_mem_map)
   {
     auto ind = entry.first;
-    auto mem_id = entry.second;
-    auto mem_blk = _mem_alloc->getMemoryBlock(mem_id);
+    auto mem_blk = entry.second;
     auto &tensor = _tensors[ind];
-    tensor->setBuffer(_mem_alloc->base() + mem_blk.base_offset);
+    tensor->setBuffer(_mem_alloc->base() + mem_blk.offset);
   }
 }
 

diff --git a/runtimes/neurun/src/backend/cpu/TensorBuilder.h b/runtimes/neurun/src/backend/cpu/TensorBuilder.h
index 72cf425..0c126d2 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/TensorBuilder.h
+++ b/runtimes/neurun/src/backend/cpu/TensorBuilder.h
@@ -22,6 +22,7 @@
 #include "backend/interface/ITensorBuilder.h"
 #include "backend/cpu/operand/Tensor.h"
 #include "graph/operand/Index.h"
+#include "MemoryAllocator.h"
 
 namespace neurun
 {
@@ -30,8 +31,6 @@ namespace backend
 namespace cpu
 {
 
-struct IMemoryAllocator;
-
 class TensorBuilder : public ITensorBuilder
 {
 public:
@@ -53,7 +52,7 @@ public:
 private:
   std::unordered_map<graph::operand::Index, ::arm_compute::TensorInfo> _tensor_info_map;
   std::unordered_map<graph::operand::Index, std::shared_ptr<operand::Tensor>> _tensors;
-  std::unordered_map<graph::operand::Index, uint32_t> _tensor_mem_map;
+  std::unordered_map<graph::operand::Index, Block> _tensor_mem_map;
   std::shared_ptr<IMemoryAllocator> _mem_alloc;
 };
 

diff --git a/runtimes/neurun/test/backend/cpu/MemoryAllocator.cc b/runtimes/neurun/test/backend/cpu/MemoryAllocator.cc
index 5a115d5..9d29e79 100644 (file)
--- a/runtimes/neurun/test/backend/cpu/MemoryAllocator.cc
+++ b/runtimes/neurun/test/backend/cpu/MemoryAllocator.cc
@@ -17,28 +17,22 @@
 #include <gtest/gtest.h>
 
 #include "backend/cpu/MemoryAllocator.h"
+#include "graph/operand/Index.h"
 
 TEST(BumpAllocator, allocate_test)
 {
   ::neurun::backend::cpu::BumpAllocator allocator;
 
-  size_t mem_sz0 = 10;
-  auto mem_id0 = allocator.allocate(mem_sz0);
-  auto mem_blk0 = allocator.getMemoryBlock(mem_id0);
-  ASSERT_EQ(mem_blk0.base_offset, 0);
-  ASSERT_EQ(mem_blk0.size, mem_sz0);
+  auto allocate = [&allocator](uint32_t index, size_t size, uint32_t expected_offset) {
+    ::neurun::graph::operand::Index mem_idx(index);
+    auto mem_blk = allocator.allocate(mem_idx, size);
+    ASSERT_EQ(mem_blk.offset, expected_offset);
+    ASSERT_EQ(mem_blk.size, size);
+  };
 
-  size_t mem_sz1 = 20;
-  auto mem_id1 = allocator.allocate(mem_sz1);
-  auto mem_blk1 = allocator.getMemoryBlock(mem_id1);
-  ASSERT_EQ(mem_blk1.base_offset, mem_sz0);
-  ASSERT_EQ(mem_blk1.size, mem_sz1);
-
-  size_t mem_sz2 = 30;
-  auto mem_id2 = allocator.allocate(mem_sz2);
-  auto mem_blk2 = allocator.getMemoryBlock(mem_id2);
-  ASSERT_EQ(mem_blk2.base_offset, mem_sz0 + mem_sz1);
-  ASSERT_EQ(mem_blk2.size, mem_sz2);
+  allocate(0, 10, 0);
+  allocate(1, 20, 10);
+  allocate(2, 30, 30);
 
   ASSERT_EQ(allocator.base(), nullptr);
 }
@@ -47,14 +41,14 @@ TEST(BumpAllocator, finalize_test)
 {
   ::neurun::backend::cpu::BumpAllocator allocator;
 
-  size_t mem_sz0 = 10;
-  allocator.allocate(mem_sz0);
-
-  size_t mem_sz1 = 20;
-  allocator.allocate(mem_sz1);
+  auto allocate = [&allocator](uint32_t index, size_t size) {
+    ::neurun::graph::operand::Index mem_idx(index);
+    auto mem_blk = allocator.allocate(mem_idx, size);
+  };
 
-  size_t mem_sz2 = 30;
-  allocator.allocate(mem_sz2);
+  allocate(0, 10);
+  allocate(1, 20);
+  allocate(2, 30);
 
   allocator.finalize();