sneaxiy
6 years ago
parent
6d8771b55c
commit
2a639d5c2a
@ -0,0 +1,74 @@
|
||||
// Copyright (c) 2019 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 <array>
|
||||
#include <deque>
|
||||
#include "paddle/fluid/platform/enforce.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace framework {
|
||||
|
||||
template <typename T, size_t N>
|
||||
class SmallStack {
|
||||
static_assert(N > 0, "N must be larger than 0");
|
||||
|
||||
public:
|
||||
inline void push(const T& item) {
|
||||
if (size_ < N) {
|
||||
head_[size_] = item;
|
||||
} else {
|
||||
tail_.emplace_back(item);
|
||||
}
|
||||
++size_;
|
||||
}
|
||||
|
||||
inline void pop() {
|
||||
PADDLE_ENFORCE(!empty(), "Try to pop element from empty stack.");
|
||||
if (size_ > N) {
|
||||
tail_.pop_back();
|
||||
}
|
||||
--size_;
|
||||
}
|
||||
|
||||
inline const T& top() const {
|
||||
PADDLE_ENFORCE(!empty(), "Try to get top element of empty stack.");
|
||||
return size_ <= N ? head_[size_ - 1] : tail_.back();
|
||||
}
|
||||
|
||||
inline T& top() {
|
||||
PADDLE_ENFORCE(!empty(), "Try to get top element of empty stack.");
|
||||
return size_ <= N ? head_[size_ - 1] : tail_.back();
|
||||
}
|
||||
|
||||
inline bool empty() const { return size_ == 0; }
|
||||
|
||||
inline size_t size() const { return size_; }
|
||||
|
||||
// This API can only be used in unittest
|
||||
T& operator[](size_t i) { return i < N ? head_[i] : tail_[i - N]; }
|
||||
|
||||
const T& operator[](size_t i) const {
|
||||
return i < N ? head_[i] : tail_[i - N];
|
||||
}
|
||||
|
||||
private:
|
||||
T head_[N];
|
||||
std::deque<T> tail_;
|
||||
size_t size_;
|
||||
};
|
||||
|
||||
} // namespace framework
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,145 @@
|
||||
// Copyright (c) 2019 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/multi_bin_buffered_allocator.h"
|
||||
#include <algorithm>
|
||||
#include <limits>
|
||||
#include "paddle/fluid/platform/lock_guard_ptr.h"
|
||||
|
||||
DEFINE_double(tolerant_times, 2,
|
||||
"Tolerant memory size times of buffered_allocator");
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
static void CheckAndModifyMemoryDivisionPlan(
|
||||
std::vector<size_t> *division_plan) {
|
||||
// Check whether the division plan is strictly sorted
|
||||
bool is_strictly_sorted = true;
|
||||
for (size_t i = 1; i < division_plan->size(); ++i) {
|
||||
if ((*division_plan)[i - 1] >= (*division_plan)[i]) {
|
||||
is_strictly_sorted = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
PADDLE_ENFORCE(is_strictly_sorted, "Divison plan must be stricted sorted");
|
||||
|
||||
// Insert 0 and remove MAX to disivion plan for clean binary searching code
|
||||
if (division_plan->empty() || division_plan->front() != 0) {
|
||||
division_plan->insert(division_plan->begin(), 0);
|
||||
}
|
||||
|
||||
constexpr auto kSizeTypeMax = std::numeric_limits<size_t>::max();
|
||||
if (division_plan->back() == kSizeTypeMax) {
|
||||
division_plan->pop_back();
|
||||
}
|
||||
|
||||
PADDLE_ENFORCE(division_plan->size() >= 1, "Division plan cannot be empty");
|
||||
}
|
||||
|
||||
static std::vector<size_t> GetDefaultDivisionPlan() {
|
||||
std::vector<size_t> plan;
|
||||
for (size_t i = 0; i < sizeof(size_t) * 8; ++i) {
|
||||
plan.push_back(static_cast<size_t>(1) << i);
|
||||
}
|
||||
return plan;
|
||||
}
|
||||
|
||||
inline static size_t FindDivisionPlanBinIndex(const std::vector<size_t> &bins,
|
||||
size_t size) {
|
||||
return static_cast<size_t>(std::upper_bound(bins.begin(), bins.end(), size) -
|
||||
bins.begin() - 1);
|
||||
}
|
||||
|
||||
inline static size_t TolerantUpperSize(size_t size) {
|
||||
return static_cast<size_t>(size * FLAGS_tolerant_times);
|
||||
}
|
||||
|
||||
MultiBinBufferedAllocator::MultiBinBufferedAllocator(
|
||||
std::shared_ptr<Allocator> underlying_allocator)
|
||||
: MultiBinBufferedAllocator(std::move(underlying_allocator),
|
||||
GetDefaultDivisionPlan()) {}
|
||||
|
||||
MultiBinBufferedAllocator::MultiBinBufferedAllocator(
|
||||
std::shared_ptr<Allocator> underlying_allocator,
|
||||
const std::vector<size_t> &division_plan)
|
||||
: underlying_allocator_(std::move(underlying_allocator)),
|
||||
division_plan_(division_plan) {
|
||||
CheckAndModifyMemoryDivisionPlan(&division_plan_);
|
||||
allocations_.resize(division_plan_.size());
|
||||
mtx_.resize(division_plan_.size());
|
||||
if (underlying_allocator_->IsAllocThreadSafe()) {
|
||||
for (auto &mtx : mtx_) {
|
||||
mtx.reset(new std::mutex());
|
||||
}
|
||||
}
|
||||
|
||||
VLOG(1) << "FLAGS_tolerant_times = " << FLAGS_tolerant_times;
|
||||
}
|
||||
|
||||
void MultiBinBufferedAllocator::FreeImpl(Allocation *allocation) {
|
||||
auto bin_index = FindDivisionPlanBinIndex(division_plan_, allocation->size());
|
||||
{
|
||||
platform::LockGuardPtr<std::mutex> guard(mtx_[bin_index]);
|
||||
allocations_[bin_index].emplace(allocation->size(),
|
||||
AllocationPtr(allocation));
|
||||
}
|
||||
}
|
||||
|
||||
void MultiBinBufferedAllocator::FreeCache(size_t size, size_t bin_index) {
|
||||
size_t accumulated_size = 0;
|
||||
// FIXME(zjl): free the largest first when there is no extra
|
||||
for (size_t i = allocations_.size() - 1; i != static_cast<size_t>(-1); --i) {
|
||||
platform::LockGuardPtr<std::mutex> lock(mtx_[i]);
|
||||
if (allocations_[i].empty()) continue;
|
||||
auto it = --allocations_[i].end();
|
||||
do {
|
||||
accumulated_size += it->second->size();
|
||||
underlying_allocator_->Free(it->second.release());
|
||||
allocations_[i].erase(it--);
|
||||
if (accumulated_size >= size) {
|
||||
return;
|
||||
}
|
||||
} while (!allocations_[i].empty());
|
||||
}
|
||||
}
|
||||
|
||||
Allocation *MultiBinBufferedAllocator::AllocateImpl(size_t size, Attr attr) {
|
||||
auto bin_index = FindDivisionPlanBinIndex(division_plan_, size);
|
||||
auto upper_size = TolerantUpperSize(size);
|
||||
|
||||
for (; upper_size >= division_plan_[bin_index]; ++bin_index) {
|
||||
auto &allocation = allocations_[bin_index];
|
||||
platform::LockGuardPtr<std::mutex> lock(mtx_[bin_index]);
|
||||
auto it = allocation.lower_bound(size);
|
||||
if (it != allocation.end() && it->second->size() < upper_size) {
|
||||
auto ret = std::move(it->second);
|
||||
allocation.erase(it);
|
||||
return ret.release();
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
return underlying_allocator_->Allocate(size, attr).release();
|
||||
} catch (BadAlloc &) {
|
||||
VLOG(2) << "BadAlloc raises, try to free " << size << " caches";
|
||||
FreeCache(size, bin_index);
|
||||
return underlying_allocator_->Allocate(size, attr).release();
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,54 @@
|
||||
// Copyright (c) 2019 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 <map>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/fluid/memory/allocation/allocator.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
class MultiBinBufferedAllocator : public Allocator {
|
||||
public:
|
||||
explicit MultiBinBufferedAllocator(
|
||||
std::shared_ptr<Allocator> underlying_allocator);
|
||||
|
||||
MultiBinBufferedAllocator(std::shared_ptr<Allocator> underlying_allocator,
|
||||
const std::vector<size_t>& division_plan);
|
||||
|
||||
bool IsAllocThreadSafe() const override { return mtx_.front() != nullptr; }
|
||||
|
||||
void ClearCache() { FreeCache(static_cast<size_t>(-1), 0); }
|
||||
|
||||
protected:
|
||||
Allocation* AllocateImpl(size_t size, Attr attr) override;
|
||||
void FreeImpl(Allocation* allocation) override;
|
||||
|
||||
private:
|
||||
void FreeCache(size_t size, size_t bin_index);
|
||||
|
||||
std::shared_ptr<Allocator> underlying_allocator_;
|
||||
std::vector<std::multimap<size_t, AllocationPtr>> allocations_;
|
||||
std::vector<size_t> division_plan_;
|
||||
std::vector<std::unique_ptr<std::mutex>> mtx_;
|
||||
};
|
||||
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,148 @@
|
||||
// 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/multi_bin_buffered_allocator.h"
|
||||
#include <gtest/gtest.h>
|
||||
#include <vector>
|
||||
#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 {
|
||||
|
||||
inline std::shared_ptr<MultiBinBufferedAllocator> GetBufferedAllocator(
|
||||
Allocation *allocation, bool thread_safe) {
|
||||
std::shared_ptr<Allocator> allocator(new BestFitAllocator(allocation));
|
||||
if (thread_safe) {
|
||||
allocator.reset(new LockedAllocator(std::move(allocator)));
|
||||
}
|
||||
|
||||
return std::make_shared<MultiBinBufferedAllocator>(allocator);
|
||||
}
|
||||
|
||||
TEST(buffered_allocator, thread_safety) {
|
||||
std::unique_ptr<CPUAllocator> allocator(new CPUAllocator());
|
||||
auto chunk = allocator->Allocate(1 << 20, allocator->kDefault);
|
||||
{
|
||||
auto buf_allocator = GetBufferedAllocator(chunk.get(), true);
|
||||
ASSERT_EQ(buf_allocator->IsAllocThreadSafe(), true);
|
||||
}
|
||||
|
||||
{
|
||||
auto buf_allocator = GetBufferedAllocator(chunk.get(), false);
|
||||
ASSERT_EQ(buf_allocator->IsAllocThreadSafe(), false);
|
||||
}
|
||||
}
|
||||
|
||||
class StubAllocation : public Allocation {
|
||||
public:
|
||||
using Allocation::Allocation;
|
||||
};
|
||||
|
||||
class StubAllocator : public Allocator {
|
||||
public:
|
||||
void ResetCounter() {
|
||||
construct_count_ = 0;
|
||||
destruct_count_ = 0;
|
||||
}
|
||||
|
||||
size_t GetAllocCount() const { return construct_count_; }
|
||||
|
||||
size_t GetFreeCount() const { return destruct_count_; }
|
||||
|
||||
protected:
|
||||
void FreeImpl(Allocation *allocation) override {
|
||||
auto *alloc = dynamic_cast<StubAllocation *>(allocation);
|
||||
PADDLE_ENFORCE_NOT_NULL(alloc);
|
||||
if (alloc->ptr()) delete[] static_cast<uint8_t *>(alloc->ptr());
|
||||
++destruct_count_;
|
||||
delete allocation;
|
||||
}
|
||||
|
||||
Allocation *AllocateImpl(size_t size, Allocator::Attr attr) override {
|
||||
++construct_count_;
|
||||
if (size == 0) {
|
||||
return new StubAllocation(nullptr, 0, platform::CPUPlace());
|
||||
} else {
|
||||
return new StubAllocation(new uint8_t[size], size, platform::CPUPlace());
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
size_t construct_count_ = 0;
|
||||
size_t destruct_count_ = 0;
|
||||
};
|
||||
|
||||
constexpr size_t kZero = 0;
|
||||
constexpr size_t kOne = 1;
|
||||
constexpr size_t kTwo = 2;
|
||||
|
||||
TEST(buffered_allocator, lazy_free) {
|
||||
std::vector<int> original_alloc_size({1022, 1023, 1024, 1025, 1026});
|
||||
for (auto alloc_size : original_alloc_size) {
|
||||
auto stub_allocator = std::make_shared<StubAllocator>();
|
||||
auto *underlying_allocator = stub_allocator.get();
|
||||
auto allocator =
|
||||
std::make_shared<MultiBinBufferedAllocator>(stub_allocator);
|
||||
|
||||
{
|
||||
underlying_allocator->ResetCounter();
|
||||
auto x = allocator->Allocate(alloc_size, allocator->kDefault);
|
||||
ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
x = nullptr;
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
}
|
||||
|
||||
{
|
||||
underlying_allocator->ResetCounter();
|
||||
auto x = allocator->Allocate(900, allocator->kDefault);
|
||||
ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
auto y = allocator->Allocate(2048, allocator->kDefault);
|
||||
ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
x = nullptr;
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
y = nullptr;
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
|
||||
}
|
||||
|
||||
{
|
||||
underlying_allocator->ResetCounter();
|
||||
allocator->ClearCache();
|
||||
ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
|
||||
ASSERT_EQ(underlying_allocator->GetFreeCount(), kTwo);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(buffered_allocator, garbage_collection) {
|
||||
std::unique_ptr<CPUAllocator> cpu_allocator(new CPUAllocator());
|
||||
auto chunk = cpu_allocator->Allocate(2048, cpu_allocator->kDefault);
|
||||
auto allocator = GetBufferedAllocator(chunk.get(), false);
|
||||
auto x1 = allocator->Allocate(1600, allocator->kDefault);
|
||||
auto x2 = allocator->Allocate(400, allocator->kDefault);
|
||||
x1 = nullptr;
|
||||
x2 = nullptr;
|
||||
auto x3 = allocator->Allocate(1600, allocator->kDefault);
|
||||
ASSERT_NE(x3, nullptr);
|
||||
ASSERT_NE(x3->ptr(), nullptr);
|
||||
}
|
||||
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue