[onert] Reallocated tensor when size becomes larger (#13658)

This commit updates dynamic allocator usage on basic tensor to reallocate when shape is changed to larger size. It will improve performance and reduce memory usage when tensor size is lesser by dynamic shape on inference than prepare phase. ONE-DCO-1.0-Signed-off-by: Hyeongseok Oh <[email protected]>
Samsung · Aug 14, 2024 · c31c280 · c31c280
1 parent ab9a924
commit c31c280
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 31 deletions.
diff --git a/runtime/onert/core/include/backend/basic/Tensor.h b/runtime/onert/core/include/backend/basic/Tensor.h
@@ -41,8 +41,8 @@ class Tensor : public IPortableTensor
 
 public:
   Tensor(const ir::OperandInfo &info, DynamicMemoryManager *dynamic_mem_mgr)
-    : IPortableTensor(info), _layout(ir::Layout::NHWC), _buffer(nullptr), _num_references(0),
-      _dynamic_mem_mgr(dynamic_mem_mgr), _allocator(nullptr)
+    : IPortableTensor(info), _layout(ir::Layout::NHWC), _buffer(nullptr), _size(info.total_size()),
+      _num_references(0), _dynamic_mem_mgr(dynamic_mem_mgr), _allocator(nullptr)
   {
     // DO NOTHING
   }
@@ -128,6 +128,7 @@ class Tensor : public IPortableTensor
 protected:
   const ir::Layout _layout;
   uint8_t *_buffer;
+  size_t _size;
   int32_t _num_references;
   DynamicMemoryManager *_dynamic_mem_mgr;
 

diff --git a/runtime/onert/core/src/backend/basic/Tensor.cc b/runtime/onert/core/src/backend/basic/Tensor.cc
@@ -32,40 +32,23 @@ void Tensor::setShape(const ir::Shape &new_shape) { _info.shape(new_shape); }
 
 bool Tensor::applyShape(const ir::Shape &new_shape)
 {
-  bool previously_dynamic = is_dynamic();
+  if (_buffer != nullptr && new_shape == _info.shape())
+    return true;
 
-  auto allocTensorMem = [&]() {
-    auto capacity = total_size();
-    assert(_dynamic_mem_mgr);
-    auto alloc = _dynamic_mem_mgr->allocate(this, capacity);
-    setBuffer(alloc);
-  };
-
-  if (!previously_dynamic || buffer() == nullptr)
+  // Always set shape - when buffer with same or larger size was already allocated, shape could
+  // differ
+  _info.shape(new_shape);
+  set_dynamic();
+  if (_buffer == nullptr || _size < _info.total_size())
   {
-    // Always set shape - when buffer with same size was already allocated, shape could differ
-    setShape(new_shape);
-    set_dynamic();
-    allocTensorMem();
-  }
-  else
-  {
-    auto previous_size = total_size();
-    auto new_size = new_shape.num_elements() * ir::sizeOfDataType(data_type());
-    if (previous_size != new_size)
-    {
-      assert(_dynamic_mem_mgr);
+    assert(_dynamic_mem_mgr);
+    if (_allocator)
       _dynamic_mem_mgr->deallocate(this);
 
-      setShape(new_shape);
-      set_dynamic();
-      allocTensorMem();
-    }
-    else
-    { // when buffer with same size was already allocated, shape could differ
-      setShape(new_shape);
-    }
+    _size = _info.total_size();
+    setBuffer(_dynamic_mem_mgr->allocate(this, _size));
   }
+
   return true;
 }