Revert "[Unify Tensors PR #1] Replaced pten::Allocation with...

Revert "[Unify Tensors PR #1] Replaced pten::Allocation with shared_ptr<memory::Allocation> for Storage (#38301)"

This reverts commit 42cf2bee.

paddle/pten/api/lib/utils/allocator.h

@@ -38,8 +38,6 @@ class DefaultAllocator : public pten::Allocator {
     return Allocation(ptr, a.release(), &Delete, place_);
   }
 
-  const paddle::platform::Place& place() override { return place_; }
-
  private:
   paddle::platform::Place place_;
   static paddle::memory::Allocator::AllocationDeleter deleter_;

paddle/pten/api/lib/utils/storage.cc

@@ -20,15 +20,13 @@ namespace experimental {
 ExternalStorage::ExternalStorage(void* ptr,
                                  size_t size,
                                  const paddle::platform::Place& place)
-    : pten::Storage(
-          std::make_shared<paddle::memory::Allocation>(ptr, size, place)),
-      size_(size) {}
+    : pten::Storage(pten::Allocation(ptr, place)), size_(size) {}
 
 ExternalStorage::ExternalStorage(const pten::intrusive_ptr<pten::Storage>& root,
                                  size_t delta,
                                  size_t size)
-    : Storage(std::make_shared<paddle::memory::Allocation>(
-          static_cast<uint8_t*>(root->data()) + delta, size, root->place())),
+    : Storage(pten::Allocation(static_cast<uint8_t*>(root->data()) + delta,
+                               root->place())),
       size_(size) {
   PADDLE_ENFORCE_LE(static_cast<size_t>(delta + size),
                     root->size(),

paddle/pten/api/lib/utils/storage.h

@@ -35,18 +35,13 @@ class ExternalStorage : public pten::Storage {
   }
 
   void Clear() override {
-    data_ = nullptr;
+    data_.Clear();
     size_ = 0;
-    offset_ = 0;
   }
 
   size_t size() const noexcept override { return size_; }
   const paddle::platform::Place& place() const override {
-    PADDLE_ENFORCE_NOT_NULL(
-        data_,
-        paddle::platform::errors::Unavailable(
-            "Unable to visit place as data_ has not been initialized yet."));
-    return data_->place();
+    return data_.place();
   }
   bool OwnsMemory() const noexcept override { return false; }
 
@@ -59,61 +54,74 @@ class SharedStorage : public pten::Storage {
   explicit SharedStorage(
       const std::shared_ptr<paddle::memory::Allocation>& allocation,
       size_t offset)
-      : Storage(allocation) {
+      : allocation_(allocation) {
     CHECK(allocation);
-    place_ = allocation->place();
+    data_ = pten::Allocation(
+        reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(allocation->ptr()) +
+                                offset),
+        allocation->place());
     size_ = allocation->size();
-    offset_ = offset;
   }
 
   // In order to be compatible with the original Tensor design and execution
   // system, we need to allow the uninitialized SharedStorage to exist,
   // and it can be removed after the compatibility phase is over in the future
   explicit SharedStorage(const paddle::platform::Place& place) {
-    place_ = place;
+    data_ = pten::Allocation(nullptr, place);
   }
 
+  static const char* name() { return "SharedStorage"; }
+
+  // In order to be compatible with the original Tensor design and execution
+  // system, we need to allow the SharedStorage realloc,
+  // and it can be removed after the compatibility phase is over in the future
   void Realloc(size_t n) override {
-    this->Clear();
-    data_ = paddle::memory::AllocShared(place(), n);
-    size_ = n;
+    ResetAllocation(paddle::memory::AllocShared(place(), n), 0);
   }
 
-  static const char* name() { return "SharedStorage"; }
-
   void Clear() override {
-    data_ = nullptr;
+    data_.Clear();
     size_ = 0;
   }
 
   size_t size() const noexcept override { return size_; }
-  const paddle::platform::Place& place() const override { return place_; }
+  const paddle::platform::Place& place() const override {
+    return data_.place();
+  }
   bool OwnsMemory() const noexcept override { return false; }
 
   const std::shared_ptr<paddle::memory::Allocation>& GetAllocation() {
-    return data_;
+    return allocation_;
   }
 
   // Temporary method: For compatible with fluid Tensor and improve performance
   void ResetAllocation(std::shared_ptr<paddle::memory::Allocation> allocation,
                        size_t offset) {
-    data_ = allocation;
+    allocation_ = allocation;
+    data_ = pten::Allocation(
+        reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(allocation->ptr()) +
+                                offset),
+        allocation->place());
     size_ = allocation->size();
-    place_ = allocation->place();
-    offset_ = offset;
   }
 
   // Temporary method: For compatible with fluid Tensor and improve performance
   void ResetAllocationPlace(const paddle::platform::Place& place) {
-    place_ = place;
+    data_ = pten::Allocation(nullptr, place);
   }
 
   // Temporary method: For compatible with fluid Tensor and improve performance
-  void Reset() { this->Clear(); }
+  void Reset() {
+    if (allocation_ != nullptr) {
+      allocation_.reset();
+    }
+    data_.Clear();
+    size_ = 0;
+  }
 
  private:
-  Place place_;
   int64_t size_{0};
+  std::shared_ptr<paddle::memory::Allocation> allocation_;
 };
 
 class TensorStorage : public paddle::memory::allocation::Allocation {

paddle/pten/core/allocator.h

@@ -134,12 +134,9 @@ inline void swap(Allocation& a, Allocation& b) noexcept {
 /// mainly used for general data structures such as Tensor. The raw
 /// allocator is more universal and efficient.
 class Allocator {
-  using Place = paddle::platform::Place;
-
  public:
   virtual ~Allocator() = default;
   virtual Allocation Allocate(size_t bytes_size) = 0;
-  virtual const Place& place() = 0;
 };
 
 inline Allocation Allocate(const std::shared_ptr<Allocator>& a, size_t n) {

paddle/pten/core/storage.cc

@@ -17,8 +17,8 @@ limitations under the License. */
 namespace pten {
 
 void TensorStorage::Realloc(size_t size) {
-  this->Clear();
-  data_ = paddle::memory::AllocShared(alloc_->place(), size);
+  data_.Clear();
+  data_ = Allocate(alloc_, size);
   size_ = size;
 }
 

paddle/pten/core/storage.h

@@ -21,7 +21,6 @@ limitations under the License. */
 #include "paddle/pten/core/utils/intrusive_ref_counter.h"
 #include "paddle/pten/core/utils/type_info.h"
 
-#include "paddle/fluid/memory/memory.h"
 #include "paddle/fluid/platform/place.h"
 #include "paddle/pten/core/allocator.h"
 
@@ -36,33 +35,15 @@ class Storage : public intrusive_ref_counter<Storage> {
   Storage() = default;
   Storage(const Storage&) = delete;
 
-  /* --------- shared_ptr<Allocation> -------- */
-  // Initialize a Storage with unique Allocation
-  explicit Storage(std::shared_ptr<paddle::memory::Allocation>&& data)
-      : data_(std::move(data)) {}
-
-  // Initialize a Storage shareing Allocation with another storage
-  explicit Storage(const std::shared_ptr<paddle::memory::Allocation>& data)
-      : data_(data) {}
-
-  void* data() const {
-    return data_ ? reinterpret_cast<void*>(
-                       reinterpret_cast<uintptr_t>(data_->ptr()) + offset_)
-                 : nullptr;
-  }
-
-  const std::shared_ptr<paddle::memory::Allocation> data_shared() const {
-    return data_;
-  }
-
-  virtual void ReallocShared(size_t n) {
-    PADDLE_THROW(paddle::platform::errors::Unimplemented(
-        "ReallocShared has not been overrided by the current Storage"));
-  }
-  /* --------- shared_ptr<Allocation> -------- */
+  explicit Storage(Allocation&& data) : data_(std::move(data)) {}
 
   virtual ~Storage() = default;
 
+  /// \brief Get the mutable data pointer of the storage.
+  /// This function is set to inline to improve performance.
+  /// \return The mutable data pointer of the storage.
+  void* data() const noexcept { return data_.operator->(); }
+
   virtual void Clear() = 0;
 
   virtual size_t size() const = 0;
@@ -71,8 +52,7 @@ class Storage : public intrusive_ref_counter<Storage> {
   virtual void Realloc(size_t n) = 0;
 
  protected:
-  size_t offset_{0};
-  std::shared_ptr<paddle::memory::Allocation> data_;
+  Allocation data_;
 };
 
 class TensorStorage : public Storage {
@@ -80,38 +60,23 @@ class TensorStorage : public Storage {
   using Place = paddle::platform::Place;
 
   explicit TensorStorage(const std::shared_ptr<Allocator>& a) : alloc_(a) {}
-
   TensorStorage(const std::shared_ptr<Allocator>& a, size_t size)
-      : Storage(paddle::memory::AllocShared(a->place(), size)), alloc_(a) {
-    size_ = data_->size();
-  }
-
-  void Clear() override {
-    data_ = nullptr;
-    size_ = 0;
-    offset_ = 0;
-  }
-
-  void Realloc(size_t size) override;
+      : Storage(Allocate(a, size)), alloc_(a), size_(size) {}
 
   ~TensorStorage() = default;
 
   static const char* name() { return "TensorStorage"; }
 
+  void Realloc(size_t size) override;
+
   size_t size() const noexcept override { return size_; }
 
-  const Place& place() const override {
-    if (!data_ && !alloc_) {
-      PADDLE_THROW(paddle::platform::errors::Unimplemented(
-          "Unable to visit place: either data_ or alloc_ has to be initialized "
-          "first."));
-    }
-    if (data_) {
-      return data_->place();
-    }
-    return alloc_->place();
+  void Clear() override {
+    data_.Clear();
+    size_ = 0;
   }
 
+  const Place& place() const override { return data_.place(); }
   bool OwnsMemory() const noexcept override { return true; }
   const std::shared_ptr<Allocator>& allocator() const noexcept {
     return alloc_;

paddle/pten/tests/core/allocator.h

@@ -44,12 +44,8 @@ class FancyAllocator : public pten::Allocator {
 
   Allocation Allocate(size_t bytes_size) override {
     void* data = ::operator new(bytes_size);
-    return Allocation(data, data, &Delete, place());
+    return Allocation(data, data, &Delete, paddle::platform::CPUPlace());
   }
-
-  const paddle::platform::Place& place() override { return place_; }
-
-  paddle::platform::Place place_ = paddle::platform::CPUPlace();
 };
 
 template <typename T>