diff --git a/paddle/pten/api/lib/utils/allocator.h b/paddle/pten/api/lib/utils/allocator.h index ab7a0fcef289f59abd7f8c1429432ca719c4e7bf..4f5a810e400ce471ae05b4fef8a6b25a3465ab23 100644 --- a/paddle/pten/api/lib/utils/allocator.h +++ b/paddle/pten/api/lib/utils/allocator.h @@ -38,6 +38,8 @@ 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_; diff --git a/paddle/pten/api/lib/utils/storage.cc b/paddle/pten/api/lib/utils/storage.cc index ba26e7f600d60c07df1179fdd1a4ff84d2383b5c..9ee1b9e5b7f92ee4a709ffcb2eafafb6c6bca4ac 100644 --- a/paddle/pten/api/lib/utils/storage.cc +++ b/paddle/pten/api/lib/utils/storage.cc @@ -20,13 +20,15 @@ namespace experimental { ExternalStorage::ExternalStorage(void* ptr, size_t size, const paddle::platform::Place& place) - : pten::Storage(pten::Allocation(ptr, place)), size_(size) {} + : pten::Storage( + std::make_shared(ptr, size, place)), + size_(size) {} ExternalStorage::ExternalStorage(const pten::intrusive_ptr& root, size_t delta, size_t size) - : Storage(pten::Allocation(static_cast(root->data()) + delta, - root->place())), + : Storage(std::make_shared( + static_cast(root->data()) + delta, size, root->place())), size_(size) { PADDLE_ENFORCE_LE(static_cast(delta + size), root->size(), diff --git a/paddle/pten/api/lib/utils/storage.h b/paddle/pten/api/lib/utils/storage.h index e98c5a82fedddfc0bbb4e404033b910e6f154b2e..132ffe165eec633a28bc8c2b2a036e9432100047 100644 --- a/paddle/pten/api/lib/utils/storage.h +++ b/paddle/pten/api/lib/utils/storage.h @@ -35,13 +35,18 @@ class ExternalStorage : public pten::Storage { } void Clear() override { - data_.Clear(); + data_ = nullptr; size_ = 0; + offset_ = 0; } size_t size() const noexcept override { return size_; } const paddle::platform::Place& place() const override { - return data_.place(); + PADDLE_ENFORCE_NOT_NULL( + data_, + paddle::platform::errors::Unavailable( + "Unable to visit place as data_ has not been initialized yet.")); + return data_->place(); } bool OwnsMemory() const noexcept override { return false; } @@ -54,74 +59,61 @@ class SharedStorage : public pten::Storage { explicit SharedStorage( const std::shared_ptr& allocation, size_t offset) - : allocation_(allocation) { + : Storage(allocation) { CHECK(allocation); - data_ = pten::Allocation( - reinterpret_cast(reinterpret_cast(allocation->ptr()) + - offset), - allocation->place()); + place_ = 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) { - data_ = pten::Allocation(nullptr, place); + place_ = 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 { - ResetAllocation(paddle::memory::AllocShared(place(), n), 0); + this->Clear(); + data_ = paddle::memory::AllocShared(place(), n); + size_ = n; } + static const char* name() { return "SharedStorage"; } + void Clear() override { - data_.Clear(); + data_ = nullptr; size_ = 0; } size_t size() const noexcept override { return size_; } - const paddle::platform::Place& place() const override { - return data_.place(); - } + const paddle::platform::Place& place() const override { return place_; } bool OwnsMemory() const noexcept override { return false; } const std::shared_ptr& GetAllocation() { - return allocation_; + return data_; } // Temporary method: For compatible with fluid Tensor and improve performance void ResetAllocation(std::shared_ptr allocation, size_t offset) { - allocation_ = allocation; - data_ = pten::Allocation( - reinterpret_cast(reinterpret_cast(allocation->ptr()) + - offset), - allocation->place()); + data_ = allocation; 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) { - data_ = pten::Allocation(nullptr, place); + place_ = place; } // Temporary method: For compatible with fluid Tensor and improve performance - void Reset() { - if (allocation_ != nullptr) { - allocation_.reset(); - } - data_.Clear(); - size_ = 0; - } + void Reset() { this->Clear(); } private: + Place place_; int64_t size_{0}; - std::shared_ptr allocation_; }; class TensorStorage : public paddle::memory::allocation::Allocation { diff --git a/paddle/pten/core/allocator.h b/paddle/pten/core/allocator.h index f03c591b1db2061dc19ed7c40d6a45325e134309..74455be1368343306e3d0a4bc7d1581ab2b210c0 100644 --- a/paddle/pten/core/allocator.h +++ b/paddle/pten/core/allocator.h @@ -134,9 +134,12 @@ 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& a, size_t n) { diff --git a/paddle/pten/core/storage.cc b/paddle/pten/core/storage.cc index 5cac122b7dee611e54f90743b4b5e283ebe7c9ac..f7c7f6873410179575050ccdc88e7002aeb936c7 100644 --- a/paddle/pten/core/storage.cc +++ b/paddle/pten/core/storage.cc @@ -17,8 +17,8 @@ limitations under the License. */ namespace pten { void TensorStorage::Realloc(size_t size) { - data_.Clear(); - data_ = Allocate(alloc_, size); + this->Clear(); + data_ = paddle::memory::AllocShared(alloc_->place(), size); size_ = size; } diff --git a/paddle/pten/core/storage.h b/paddle/pten/core/storage.h index ef9e22a0804e73f57cd8b3a7936adb7de02effdb..bc652f52c1f3b2e339757dce77ef6a63ec3560db 100644 --- a/paddle/pten/core/storage.h +++ b/paddle/pten/core/storage.h @@ -21,6 +21,7 @@ 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" @@ -35,14 +36,32 @@ class Storage : public intrusive_ref_counter { Storage() = default; Storage(const Storage&) = delete; - explicit Storage(Allocation&& data) : data_(std::move(data)) {} + /* --------- shared_ptr -------- */ + // Initialize a Storage with unique Allocation + explicit Storage(std::shared_ptr&& data) + : data_(std::move(data)) {} - virtual ~Storage() = default; + // Initialize a Storage shareing Allocation with another storage + explicit Storage(const std::shared_ptr& data) + : data_(data) {} + + void* data() const { + return data_ ? reinterpret_cast( + reinterpret_cast(data_->ptr()) + offset_) + : nullptr; + } + + const std::shared_ptr data_shared() const { + return data_; + } - /// \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 ReallocShared(size_t n) { + PADDLE_THROW(paddle::platform::errors::Unimplemented( + "ReallocShared has not been overrided by the current Storage")); + } + /* --------- shared_ptr -------- */ + + virtual ~Storage() = default; virtual void Clear() = 0; @@ -52,7 +71,8 @@ class Storage : public intrusive_ref_counter { virtual void Realloc(size_t n) = 0; protected: - Allocation data_; + size_t offset_{0}; + std::shared_ptr data_; }; class TensorStorage : public Storage { @@ -60,23 +80,38 @@ class TensorStorage : public Storage { using Place = paddle::platform::Place; explicit TensorStorage(const std::shared_ptr& a) : alloc_(a) {} + TensorStorage(const std::shared_ptr& a, size_t size) - : Storage(Allocate(a, size)), alloc_(a), size_(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; ~TensorStorage() = default; static const char* name() { return "TensorStorage"; } - void Realloc(size_t size) override; - size_t size() const noexcept override { return size_; } - void Clear() override { - data_.Clear(); - size_ = 0; + 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(); } - const Place& place() const override { return data_.place(); } bool OwnsMemory() const noexcept override { return true; } const std::shared_ptr& allocator() const noexcept { return alloc_; diff --git a/paddle/pten/tests/core/allocator.h b/paddle/pten/tests/core/allocator.h index 4a5a9b1690923d65f199de94ed7284284a7db15b..094c0e8437d9811709972b3482703bbefdc4e54e 100644 --- a/paddle/pten/tests/core/allocator.h +++ b/paddle/pten/tests/core/allocator.h @@ -44,8 +44,12 @@ class FancyAllocator : public pten::Allocator { Allocation Allocate(size_t bytes_size) override { void* data = ::operator new(bytes_size); - return Allocation(data, data, &Delete, paddle::platform::CPUPlace()); + return Allocation(data, data, &Delete, place()); } + + const paddle::platform::Place& place() override { return place_; } + + paddle::platform::Place place_ = paddle::platform::CPUPlace(); }; template