From bb04b54e8d429570b83cad39362bd411665585fa Mon Sep 17 00:00:00 2001 From: sneaxiy Date: Wed, 10 Oct 2018 03:43:38 +0000 Subject: [PATCH] add retry_allocator add unittest of retry_allocator --- paddle/fluid/memory/allocation/CMakeLists.txt | 4 + .../memory/allocation/aligned_allocator.h | 3 + .../memory/allocation/retry_allocator.cc | 88 +++++++++++++++ .../fluid/memory/allocation/retry_allocator.h | 93 ++++++++++++++++ .../memory/allocation/retry_allocator_test.cc | 100 ++++++++++++++++++ 5 files changed, 288 insertions(+) create mode 100644 paddle/fluid/memory/allocation/retry_allocator.cc create mode 100644 paddle/fluid/memory/allocation/retry_allocator.h create mode 100644 paddle/fluid/memory/allocation/retry_allocator_test.cc diff --git a/paddle/fluid/memory/allocation/CMakeLists.txt b/paddle/fluid/memory/allocation/CMakeLists.txt index 94dc13ad5f0..664b3460252 100644 --- a/paddle/fluid/memory/allocation/CMakeLists.txt +++ b/paddle/fluid/memory/allocation/CMakeLists.txt @@ -4,6 +4,8 @@ cc_library(best_fit_allocator SRCS best_fit_allocator.cc DEPS allocator) cc_library(locked_allocator SRCS locked_allocator.cc DEPS allocator) nv_library(cuda_allocator SRCS cuda_allocator.cc DEPS allocator cuda_device_guard) +cc_library(retry_allocator SRCS retry_allocator.cc DEPS allocator) + if (WITH_GPU) nv_test(best_fit_allocator_test SRCS best_fit_allocator_test.cc @@ -49,3 +51,5 @@ cc_library(allocator_facade SRCS allocator_facade.cc DEPS cuda_device_guard) nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocator_facade) + +cc_test(retry_allocator_test SRCS retry_allocator_test.cc DEPS retry_allocator naive_managed_allocator best_fit_allocator locked_allocator cpu_allocator) diff --git a/paddle/fluid/memory/allocation/aligned_allocator.h b/paddle/fluid/memory/allocation/aligned_allocator.h index 3a7868f403e..13c69c153a2 100644 --- a/paddle/fluid/memory/allocation/aligned_allocator.h +++ b/paddle/fluid/memory/allocation/aligned_allocator.h @@ -29,6 +29,9 @@ namespace allocation { // NOTE(yy): kAlignment must be 2^N. a `static_assert` should be added. template class AlignedAllocation : public Allocation { + static_assert(kAlignment > 0 && (kAlignment & (kAlignment - 1)) == 0, + "kAlignment must be 2^N"); + public: AlignedAllocation(std::unique_ptr&& underlying_allocation, size_t size) diff --git a/paddle/fluid/memory/allocation/retry_allocator.cc b/paddle/fluid/memory/allocation/retry_allocator.cc new file mode 100644 index 00000000000..ae54ac13ac6 --- /dev/null +++ b/paddle/fluid/memory/allocation/retry_allocator.cc @@ -0,0 +1,88 @@ +// 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/retry_allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +RetryAllocation::~RetryAllocation() { + auto allocator = retry_allocator_.lock(); + { + // release allocation first + if (UNLIKELY(allocator == nullptr)) return; + allocator->underlying_allocator_->Free(underlying_allocation_.release()); + } + + { + // notify all waited allocators + std::lock_guard lock(allocator->mutex_); + allocator->cv_.notify_all(); + } +} + +bool RetryAllocator::IsAllocThreadSafe() const { return true; } + +std::shared_ptr RetryAllocator::AllocateShared( + size_t size, Allocator::Attr attr) { + return std::shared_ptr(Allocate(size, attr)); +} + +std::unique_ptr RetryAllocator::Allocate(size_t size, + Allocator::Attr attr) { + auto alloc_func = [&, this]() { + return new RetryAllocation(underlying_allocator_->Allocate(size, attr), + this->shared_from_this()); + }; + + // In fact, we can unify the code of allocation success and failure + // But it would add lock even when allocation success at the first time + std::unique_ptr ret; + try { + ret.reset(alloc_func()); + } catch (BadAlloc &) { + { + // We can just write allocation retry inside the predicate function of + // wait_until + // But it needs to acquire the lock when executing predicate function + // For better performance, we use loop here + std::exception_ptr ex; + auto end_time = std::chrono::high_resolution_clock::now() + retry_time_; + std::cv_status status; + do { + { + std::unique_lock lock(mutex_); + status = cv_.wait_until(lock, end_time); + } + try { + ret.reset(alloc_func()); + } catch (BadAlloc &) { + ex = std::current_exception(); + } catch (...) { + std::rethrow_exception(std::current_exception()); + } + } while (ret == nullptr && status != std::cv_status::timeout); + + if (ret == nullptr) std::rethrow_exception(ex); + } + } catch (...) { + std::rethrow_exception(std::current_exception()); + } + return ret; +} + +} // namespace allocation +} // namespace memory +} // namespace paddle diff --git a/paddle/fluid/memory/allocation/retry_allocator.h b/paddle/fluid/memory/allocation/retry_allocator.h new file mode 100644 index 00000000000..ef7945e7502 --- /dev/null +++ b/paddle/fluid/memory/allocation/retry_allocator.h @@ -0,0 +1,93 @@ +// 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. + +#pragma once + +#include // NOLINT +#include // NOLINT +#include +#include // NOLINT +#include "paddle/fluid/memory/allocation/allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +class RetryAllocator; + +class RetryAllocation : public Allocation { + public: + RetryAllocation(std::unique_ptr&& underlying_allocation, + const std::shared_ptr& retry_allocator) + : Allocation(underlying_allocation->ptr(), underlying_allocation->size(), + underlying_allocation->place()), + underlying_allocation_(std::move(underlying_allocation)), + retry_allocator_(retry_allocator) {} + + ~RetryAllocation(); + + private: + std::unique_ptr underlying_allocation_; + std::weak_ptr retry_allocator_; +}; + +class RetryAllocator : public ManagedAllocator, + public std::enable_shared_from_this { + private: + RetryAllocator(std::unique_ptr&& allocator, size_t retry_ms) + : underlying_allocator_( + dynamic_cast(allocator.release())), + retry_time_(retry_ms) { + EnforceCheck(); + } + + public: + template + static std::shared_ptr Create(Args... args) { + return std::shared_ptr( + new RetryAllocator(std::forward(args)...)); + } + + bool IsAllocThreadSafe() const override; + + std::unique_ptr Allocate( + size_t size, Allocator::Attr attr = kDefault) override; + + std::shared_ptr AllocateShared( + size_t size, Allocator::Attr attr = kDefault) override; + + private: + void EnforceCheck() { + PADDLE_ENFORCE_NOT_NULL( + underlying_allocator_.get(), + "UnderlyingAllocator of RetryAllocator must be UnmanagedAllocator"); + PADDLE_ENFORCE(underlying_allocator_->IsAllocThreadSafe(), + "UnderlyingAllocator of RetryAllocator must be thread-safe"); + } + + std::unique_ptr underlying_allocator_; + std::chrono::milliseconds retry_time_; + std::mutex mutex_; + std::condition_variable cv_; + + // For debug, We can add an atomic integer to record how many memory sizes are + // waited to allocate + // std::atomic waited_allocate_size_{0}; + + friend class RetryAllocation; +}; + +} // namespace allocation +} // namespace memory +} // namespace paddle diff --git a/paddle/fluid/memory/allocation/retry_allocator_test.cc b/paddle/fluid/memory/allocation/retry_allocator_test.cc new file mode 100644 index 00000000000..c55742c7bef --- /dev/null +++ b/paddle/fluid/memory/allocation/retry_allocator_test.cc @@ -0,0 +1,100 @@ +// 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/retry_allocator.h" +#include +#include // NOLINT +#include // NOLINT +#include // NOLINT +#include // NOLINT +#include +#include "gtest/gtest.h" +#include "paddle/fluid/memory/allocation/best_fit_allocator.h" +#include "paddle/fluid/memory/allocation/cpu_allocator.h" +#include "paddle/fluid/memory/allocation/locked_allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +TEST(RetryAllocator, RetryAllocator) { + CPUAllocator cpu_allocator; + + size_t size = (1 << 20); + auto cpu_allocation = cpu_allocator.Allocate(size); + + std::unique_ptr best_fit_allocator( + new BestFitAllocator(cpu_allocation.get())); + std::unique_ptr locked_allocator( + new LockedAllocator(std::move(best_fit_allocator))); + + size_t thread_num = 32; + size_t sleep_time = 40; + size_t extra_time = 2; + + // Reserve to perform more tests in the future + std::vector> allocators; + { + std::unique_ptr best_fit_allocator( + new BestFitAllocator(cpu_allocation.get())); + std::unique_ptr locked_allocator( + new LockedAllocator(std::move(best_fit_allocator))); + allocators.push_back( + RetryAllocator::Create(std::move(locked_allocator), + (thread_num - 1) * (sleep_time + extra_time))); + } + + for (auto &allocator : allocators) { + std::vector threads(thread_num); + std::vector addresses(threads.size(), nullptr); + + std::mutex mutex; + std::condition_variable cv; + bool flag = false; + + for (size_t i = 0; i < threads.size(); ++i) { + threads[i] = std::thread([&, i]() { + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return flag; }); + } + + auto ret = allocator->Allocate(size - 1); + addresses[i] = ret->ptr(); + std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time)); + }); + } + + { + std::lock_guard lock(mutex); + flag = true; + cv.notify_all(); + } + + for (auto &th : threads) { + th.join(); + } + + void *val = cpu_allocation->ptr(); + bool is_all_equal = std::all_of(addresses.begin(), addresses.end(), + [val](void *p) { return p == val; }); + ASSERT_TRUE(is_all_equal); + } + + cpu_allocator.FreeUniquePtr(std::move(cpu_allocation)); +} + +} // namespace allocation +} // namespace memory +} // namespace paddle -- GitLab