From 713c05dd60e30a83a1403f2fcff9ae97a74589d5 Mon Sep 17 00:00:00 2001
From: Zeng Jinle <32832641+sneaxiy@users.noreply.github.com>
Date: Mon, 9 Sep 2019 20:28:03 +0800
Subject: [PATCH] refine tensor.mutable_data, test=develop (#19680)

---
 paddle/fluid/framework/tensor.cc      |  2 ++
 paddle/fluid/framework/tensor_test.cc | 17 ++++++++++++++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/framework/tensor.cc b/paddle/fluid/framework/tensor.cc
index f8bf9c5dc78..7b39c5359e8 100644
--- a/paddle/fluid/framework/tensor.cc
+++ b/paddle/fluid/framework/tensor.cc
@@ -51,6 +51,8 @@ void* Tensor::mutable_data(platform::Place place, proto::VarType::Type type,
   /* some versions of boost::variant don't have operator!= */
   if (holder_ == nullptr || !(holder_->place() == place) ||
       holder_->size() < size + offset_) {
+    // Reset holder first before re-allocate to save memory
+    holder_.reset();
     holder_ = memory::AllocShared(place, size);
     offset_ = 0;
   }
diff --git a/paddle/fluid/framework/tensor_test.cc b/paddle/fluid/framework/tensor_test.cc
index 238af9bafeb..f94c0c532bd 100644
--- a/paddle/fluid/framework/tensor_test.cc
+++ b/paddle/fluid/framework/tensor_test.cc
@@ -57,23 +57,28 @@ TEST(Tensor, MutableData) {
     // initialization
     p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                         platform::CPUPlace());
+    auto p1_holder = src_tensor.Holder();
     EXPECT_NE(p1, nullptr);
     // set src_tensor a new dim with large size
     // momery is supposed to be re-allocated
     p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 4}),
                                         platform::CPUPlace());
     EXPECT_NE(p2, nullptr);
-    EXPECT_NE(p1, p2);
+    auto p2_holder1 = src_tensor.Holder();
+    EXPECT_NE(p1_holder.get(), p2_holder1.get());
     // set src_tensor a new dim with same size
     // momery block is supposed to be unchanged
     p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
                                         platform::CPUPlace());
-    EXPECT_EQ(p1, p2);
+    auto p2_holder2 = src_tensor.Holder();
+    EXPECT_EQ(p2_holder1.get(), p2_holder2.get());
     // set src_tensor a new dim with smaller size
     // momery block is supposed to be unchanged
     p2 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2}),
                                         platform::CPUPlace());
+    auto p2_holder3 = src_tensor.Holder();
     EXPECT_EQ(p1, p2);
+    EXPECT_EQ(p2_holder2.get(), p2_holder3.get());
 
     float* p3 = nullptr;
     float* p4 = nullptr;
@@ -82,14 +87,18 @@ TEST(Tensor, MutableData) {
     auto* tmp = src_tensor.mutable_data<uint8_t>(framework::make_ddim({2, 2}),
                                                  platform::CPUPlace());
     p3 = reinterpret_cast<float*>(tmp);
+    auto p3_holder1 = src_tensor.Holder();
     EXPECT_EQ(p1, p3);
+    EXPECT_EQ(p2_holder3.get(), p3_holder1.get());
 
     // set src_tensor a different type but bigger size.
     // memory block is supposed to be changed.
     auto* tmp2 = src_tensor.mutable_data<double>(
         framework::make_ddim({2, 2, 3}), platform::CPUPlace());
+    auto p3_holder2 = src_tensor.Holder();
     p4 = reinterpret_cast<float*>(tmp2);
     EXPECT_NE(p1, p4);
+    EXPECT_NE(p3_holder1.get(), p3_holder2.get());
   }
   // Not sure if it's desired, but currently, Tensor type can be changed.
   {
@@ -113,13 +122,15 @@ TEST(Tensor, MutableData) {
     // initialization
     p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                         platform::CUDAPlace());
+    auto p1_holder = src_tensor.Holder();
     EXPECT_NE(p1, nullptr);
     // set src_tensor a new dim with large size
     // momery is supposed to be re-allocated
     p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 1024}),
                                         platform::CUDAPlace());
+    auto p2_holder = src_tensor.Holder();
     EXPECT_NE(p2, nullptr);
-    EXPECT_NE(p1, p2);
+    EXPECT_NE(p1_holder.get(), p2_holder.get());
     // set src_tensor a new dim with same size
     // momery block is supposed to be unchanged
     p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
-- 
GitLab