From 247e22c7317386c21b35cccfaed4567b93401f6e Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Thu, 22 Mar 2018 14:07:31 +0800
Subject: [PATCH] Refactor reshape cpu kernel without memory copy.

---
 mace/core/tensor.h     | 12 ++++++++++++
 mace/kernels/reshape.h |  4 +---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/mace/core/tensor.h b/mace/core/tensor.h
index 1d7d5deb..53ac3c2e 100644
--- a/mace/core/tensor.h
+++ b/mace/core/tensor.h
@@ -178,6 +178,18 @@ 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_) {
+      delete buffer_;
+    }
+    buffer_ = buffer;
+    is_buffer_owner_ = false;
+  }
+
   inline void ResizeImage(const std::vector<index_t> &shape,
                           const std::vector<size_t> &image_shape) {
     shape_ = shape;
diff --git a/mace/kernels/reshape.h b/mace/kernels/reshape.h
index 14e56078..ddcd0dba 100644
--- a/mace/kernels/reshape.h
+++ b/mace/kernels/reshape.h
@@ -21,9 +21,7 @@ struct ReshapeFunctor {
                   const std::vector<index_t> &out_shape,
                   Tensor *output,
                   StatsFuture *future) {
-    output->Resize(out_shape);
-    // TODO(liuqi): copy on write to avoid this copy.
-    output->CopyBytes(input->raw_data(), input->size() * sizeof(T));
+    output->ResizeWithBuffer(out_shape, input->UnderlyingBuffer());
   }
 };
 
-- 
GitLab