diff --git a/paddle/fluid/memory/allocation/CMakeLists.txt b/paddle/fluid/memory/allocation/CMakeLists.txt index 84d22ac96ca1ec25cd974a051aac2e4d52b49b59..71cf12ebf02df6b2217c239385584421042f0340 100644 --- a/paddle/fluid/memory/allocation/CMakeLists.txt +++ b/paddle/fluid/memory/allocation/CMakeLists.txt @@ -34,6 +34,7 @@ endif() cc_library(aligned_allocator SRCS aligned_allocator.cc DEPS allocator) cc_library(auto_increment_allocator SRCS auto_increment_allocator.cc DEPS allocator) +cc_library(zero_size_allocator SRCS zero_size_allocator.cc DEPS allocator) cc_library(allocator_facade SRCS allocator_facade.cc DEPS ${AllocatorFacadeDeps} cpu_allocator @@ -42,6 +43,7 @@ cc_library(allocator_facade SRCS allocator_facade.cc DEPS naive_managed_allocator aligned_allocator auto_increment_allocator + zero_size_allocator cuda_device_guard) nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocator_facade) diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc index 32228216461d62123e2819199b9fcd6ed085c3fe..971e7d02c58ed6482b758a48a803d06b5e9fe692 100644 --- a/paddle/fluid/memory/allocation/allocator_facade.cc +++ b/paddle/fluid/memory/allocation/allocator_facade.cc @@ -23,6 +23,7 @@ #include "paddle/fluid/memory/allocation/locked_allocator.h" #include "paddle/fluid/memory/allocation/naive_managed_allocator.h" #include "paddle/fluid/memory/allocation/pinned_allocator.h" +#include "paddle/fluid/memory/allocation/zero_size_allocator.h" #include "paddle/fluid/platform/cuda_device_guard.h" #include "paddle/fluid/platform/gpu_info.h" #include "paddle/fluid/platform/place.h" @@ -118,6 +119,7 @@ class AllocatorFacadePrivate { AllocatorFacadePrivate() { InitCPUAllocator(); InitCUDAAllocator(); + WrapZeroSizeAllocator(); } private: @@ -133,6 +135,13 @@ class AllocatorFacadePrivate { } #endif } + + void WrapZeroSizeAllocator() { + for (auto& pair : allocators_) { + pair.second = + std::make_shared(pair.second, pair.first); + } + } }; // Pimpl. Make interface clean. diff --git a/paddle/fluid/memory/allocation/zero_size_allocator.cc b/paddle/fluid/memory/allocation/zero_size_allocator.cc new file mode 100644 index 0000000000000000000000000000000000000000..e6cf754a4698db4a51630e72e2ad4d8960e09a7d --- /dev/null +++ b/paddle/fluid/memory/allocation/zero_size_allocator.cc @@ -0,0 +1,40 @@ +// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "paddle/fluid/memory/allocation/zero_size_allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +std::unique_ptr ZeroSizeAllocator::Allocate(size_t size, + Allocator::Attr attr) { + if (size == 0) { + return std::unique_ptr(new ZeroSizeAllocation(place_)); + } else { + return underlying_allocator_->Allocate(size, attr); + } +} +std::shared_ptr ZeroSizeAllocator::AllocateShared( + size_t size, Allocator::Attr attr) { + if (size == 0) { + return std::shared_ptr(new ZeroSizeAllocation(place_)); + } else { + return underlying_allocator_->AllocateShared(size, attr); + } +} +bool ZeroSizeAllocator::IsAllocThreadSafe() const { return true; } +} // namespace allocation +} // namespace memory +} // namespace paddle diff --git a/paddle/fluid/memory/allocation/zero_size_allocator.h b/paddle/fluid/memory/allocation/zero_size_allocator.h new file mode 100644 index 0000000000000000000000000000000000000000..62e14b633cc3f36b25f72f3073f1fd9091d6a4e2 --- /dev/null +++ b/paddle/fluid/memory/allocation/zero_size_allocator.h @@ -0,0 +1,48 @@ +// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#pragma once + +#include "paddle/fluid/memory/allocation/allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +class ZeroSizeAllocation : public Allocation { + public: + explicit ZeroSizeAllocation(const platform::Place& p) + : Allocation(nullptr, 0, p) {} +}; + +class ZeroSizeAllocator : public ManagedAllocator { + public: + ZeroSizeAllocator( + const std::shared_ptr& underlying_allocator, + const platform::Place& p) + : underlying_allocator_(underlying_allocator), place_(p) {} + std::unique_ptr Allocate(size_t size, Attr attr) override; + std::shared_ptr AllocateShared(size_t size, Attr attr) override; + bool IsAllocThreadSafe() const override; + + private: + std::shared_ptr underlying_allocator_; + const platform::Place& place_; +}; + +} // namespace allocation +} // namespace memory +} // namespace paddle