From c569dbb438af22f40abeba7f0fe8c23ccb35d14c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AF=85?= <liyin@xiaomi.com>
Date: Wed, 6 Jun 2018 19:22:19 +0800
Subject: [PATCH] Do not preallocate reshape, reuse input tensor instead.

---
 mace/core/tensor.h     | 17 ++++++++++-------
 mace/core/workspace.cc |  3 ++-
 mace/kernels/reshape.h |  3 ++-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/mace/core/tensor.h b/mace/core/tensor.h
index 0090e958..6dd41a4a 100644
--- a/mace/core/tensor.h
+++ b/mace/core/tensor.h
@@ -223,16 +223,19 @@ class Tensor {
     }
   }
 
-  inline void ResizeWithBuffer(const std::vector<index_t> &shape,
-                               BufferBase *buffer) {
-    MACE_CHECK(!has_opencl_image(), "Cannot resize image, use ResizeImage.");
-    shape_ = shape;
-    image_shape_.clear();
-    if (buffer_ != nullptr && is_buffer_owner_) {
+  // Make this tensor reuse other tensor's buffer.
+  // This tensor has the same dtype, shape and image_shape.
+  // It could be reshaped later (with image shape unchanged).
+  inline void ReuseTensorBuffer(const Tensor &other) {
+    if (is_buffer_owner_ && buffer_ != nullptr) {
       delete buffer_;
     }
-    buffer_ = buffer;
     is_buffer_owner_ = false;
+    buffer_ = other.buffer_;
+    allocator_ = other.allocator_;
+    dtype_ = other.dtype_;
+    shape_ = other.shape_;
+    image_shape_ = other.image_shape_;
   }
 
   inline MaceStatus ResizeImage(const std::vector<index_t> &shape,
diff --git a/mace/core/workspace.cc b/mace/core/workspace.cc
index 410c0118..322670de 100644
--- a/mace/core/workspace.cc
+++ b/mace/core/workspace.cc
@@ -176,7 +176,8 @@ MaceStatus Workspace::CreateOutputTensorBuffer(const NetDef &net_def,
     const int op_device =
         ProtoArgHelper::GetOptionalArg<OperatorDef, int>(
             op, "device", static_cast<int>(device_type));
-    if (op_device == device_type && !op.mem_id().empty()) {
+    if (op_device == device_type && !op.mem_id().empty()
+        && op.type() != "Reshape") {
       auto mem_ids = op.mem_id();
       int count = mem_ids.size();
       for (int i = 0; i < count; ++i) {
diff --git a/mace/kernels/reshape.h b/mace/kernels/reshape.h
index 87519bc9..cfa7bb2e 100644
--- a/mace/kernels/reshape.h
+++ b/mace/kernels/reshape.h
@@ -36,7 +36,8 @@ struct ReshapeFunctor {
                   Tensor *output,
                   StatsFuture *future) {
     MACE_UNUSED(future);
-    output->ResizeWithBuffer(out_shape, input->UnderlyingBuffer());
+    output->ReuseTensorBuffer(*input);
+    output->Reshape(out_shape);
 
     return MACE_SUCCESS;
   }
-- 
GitLab