Skip to content

Commit b59a9bf

Browse files
reyoungreyoung
committed
Clean buffered_allocator
test=develop
1 parent 26fb34c commit b59a9bf

File tree

8 files changed

+105
-195
lines changed

8 files changed

+105
-195
lines changed

‎paddle/fluid/memory/allocation/buffered_allocator.cc

Lines changed: 29 additions & 151 deletions
Original file line numberDiff line numberDiff line change
@@ -22,41 +22,6 @@ namespace memory {
2222
namespace allocation {
2323

2424
BufferedAllocator::BufferedAllocator(std::unique_ptr<Allocator>&& allocator) {
25-
std::vector<size_t> division_plan(8 * sizeof(size_t));
26-
for (size_t i = 0; i < 8 * sizeof(size_t); ++i) {
27-
division_plan[i] = (static_cast<size_t>(1) << i);
28-
}
29-
InitAndEnforceCheck(std::move(allocator), division_plan);
30-
}
31-
32-
BufferedAllocator::BufferedAllocator(std::unique_ptr<Allocator>&& allocator,
33-
const std::vector<size_t>& division_plan) {
34-
InitAndEnforceCheck(std::move(allocator), division_plan);
35-
}
36-
37-
BufferedAllocator::~BufferedAllocator() { FlushImpl(); }
38-
39-
void BufferedAllocator::FlushImpl() {
40-
for (auto& v : allocations_) {
41-
for (auto& pair : v) {
42-
underlying_allocator_->FreeUniquePtr(std::move(pair.second));
43-
}
44-
v.clear();
45-
}
46-
}
47-
48-
void BufferedAllocator::Flush() {
49-
if (mtx_) {
50-
std::lock_guard<std::mutex> lock(*mtx_);
51-
FlushImpl();
52-
} else {
53-
FlushImpl();
54-
}
55-
}
56-
57-
void BufferedAllocator::InitAndEnforceCheck(
58-
std::unique_ptr<Allocator>&& allocator,
59-
const std::vector<size_t>& division_plan) {
6025
underlying_allocator_.reset(
6126
dynamic_cast<UnmanagedAllocator*>(allocator.release()));
6227
PADDLE_ENFORCE_NOT_NULL(
@@ -65,141 +30,54 @@ void BufferedAllocator::InitAndEnforceCheck(
6530
if (underlying_allocator_->IsAllocThreadSafe()) {
6631
mtx_.reset(new std::mutex());
6732
}
68-
constexpr size_t kMax = std::numeric_limits<size_t>::max();
69-
if (division_plan.empty()) {
70-
division_plan_.assign({0, kMax});
71-
} else {
72-
auto from = division_plan.front() == 0 ? division_plan.begin() + 1
73-
: division_plan.begin();
74-
auto to = division_plan.back() == kMax ? division_plan.end() - 1
75-
: division_plan.end();
76-
division_plan_.reserve(to - from + 2);
77-
division_plan_.push_back(0);
78-
division_plan_.insert(division_plan_.end(), from, to);
79-
division_plan_.push_back(kMax);
80-
for (size_t i = 1; i < division_plan_.size(); ++i) {
81-
PADDLE_ENFORCE_LT(division_plan_[i - 1], division_plan_[i],
82-
"Division plan must be strictly sorted");
83-
}
84-
}
85-
allocations_.resize(division_plan_.size() - 1);
86-
}
87-
88-
void BufferedAllocator::InsertAllocationImpl(
89-
std::unique_ptr<Allocation>&& allocation) {
90-
auto size = allocation->size();
91-
auto idx = GetListIndex(size);
92-
allocations_[idx].emplace(size, std::move(allocation));
93-
}
94-
95-
void BufferedAllocator::InsertAllocation(
96-
std::unique_ptr<Allocation>&& allocation) {
97-
if (mtx_) {
98-
std::lock_guard<std::mutex> lock(*mtx_);
99-
InsertAllocationImpl(std::move(allocation));
100-
} else {
101-
InsertAllocationImpl(std::move(allocation));
102-
}
10333
}
10434

105-
bool BufferedAllocator::Match(size_t actual_size, size_t requested_size) {
106-
return (actual_size >> 1) < requested_size;
107-
}
108-
109-
size_t BufferedAllocator::GetListIndex(size_t size) {
110-
auto it =
111-
std::upper_bound(division_plan_.begin(), division_plan_.end(), size);
112-
return static_cast<size_t>(it - division_plan_.begin()) - 1;
113-
}
35+
BufferedAllocator::~BufferedAllocator() { FreeCache(-1UL); }
11436

115-
std::unique_ptr<Allocation> BufferedAllocator::RemoveAllocationImpl(
116-
size_t size) {
117-
auto idx = GetListIndex(size);
118-
auto& allocation_map = allocations_[idx];
119-
auto it = allocation_map.lower_bound(size);
120-
// Only remove allocation whose size is not more than twice of requested size
121-
if (it != allocation_map.end()) {
122-
if (Match(it->second->size(), size)) {
123-
auto ret = std::move(it->second);
124-
allocation_map.erase(it);
125-
return ret;
126-
} else {
127-
return nullptr;
128-
}
129-
} else {
130-
while (++idx < allocations_.size() && Match(division_plan_[idx], size)) {
131-
auto& allocation_map = allocations_[idx];
132-
if (!allocation_map.empty()) {
133-
auto it = allocation_map.begin();
134-
if (Match(it->second->size(), size)) {
135-
auto ret = std::move(it->second);
136-
allocation_map.erase(it);
137-
return ret;
138-
} else {
139-
return nullptr;
140-
}
141-
}
37+
std::unique_ptr<Allocation> BufferedAllocator::Allocate(size_t size,
38+
Allocator::Attr attr) {
39+
std::unique_ptr<Allocation> result;
40+
{
41+
platform::LockGuardPtr<std::mutex> guard(mtx_);
42+
auto it = allocations_.lower_bound(size);
43+
if (it != allocations_.end() && it->first < size * 2) {
44+
result = std::move(it->second);
45+
allocations_.erase(it);
14246
}
143-
return nullptr;
14447
}
145-
}
14648

147-
std::unique_ptr<Allocation> BufferedAllocator::RemoveAllocation(size_t size) {
148-
if (mtx_) {
149-
std::lock_guard<std::mutex> lock(*mtx_);
150-
return RemoveAllocationImpl(size);
151-
} else {
152-
return RemoveAllocationImpl(size);
49+
if (result) {
50+
return result;
15351
}
154-
}
15552

156-
std::unique_ptr<Allocation> BufferedAllocator::Allocate(size_t size,
157-
Allocator::Attr attr) {
158-
auto ret = RemoveAllocation(size);
159-
if (!ret) {
160-
try {
161-
return underlying_allocator_->Allocate(size, attr);
162-
} catch (BadAlloc&) {
163-
// if allocation failed, try to free some memorys from buffers
164-
FreeAllocations(size);
165-
return underlying_allocator_->Allocate(size, attr);
166-
}
53+
try {
54+
return underlying_allocator_->Allocate(size, attr);
55+
} catch (BadAlloc&) {
56+
FreeCache(size);
57+
return underlying_allocator_->Allocate(size, attr);
16758
}
168-
return ret;
16959
}
17060

171-
void BufferedAllocator::FreeAllocationsImpl(size_t size) {
61+
void BufferedAllocator::FreeCache(size_t size) {
62+
platform::LockGuardPtr<std::mutex> guard(mtx_);
17263
if (UNLIKELY(size == 0)) return;
17364
size_t cur = 0;
174-
for (auto& alloc_map : allocations_) {
175-
// use reverse iterator to free large allocations first
176-
while (!alloc_map.empty()) {
177-
auto it = --(alloc_map.end());
178-
cur += it->second->size();
179-
underlying_allocator_->FreeUniquePtr(std::move(it->second));
180-
alloc_map.erase(it);
181-
if (cur >= size) return;
182-
}
183-
}
184-
}
185-
186-
void BufferedAllocator::FreeAllocations(size_t size) {
187-
if (mtx_) {
188-
std::lock_guard<std::mutex> lock(*mtx_);
189-
FreeAllocationsImpl(size);
190-
} else {
191-
FreeAllocationsImpl(size);
65+
while (!allocations_.empty()) { // free the largest
66+
auto it = --allocations_.end();
67+
cur += it->second->size();
68+
underlying_allocator_->FreeUniquePtr(std::move(it->second));
69+
allocations_.erase(it);
70+
if (cur >= size) return;
19271
}
19372
}
19473

19574
void BufferedAllocator::FreeUniquePtr(std::unique_ptr<Allocation> allocation) {
196-
InsertAllocation(std::move(allocation));
75+
platform::LockGuardPtr<std::mutex> guard(mtx_);
76+
allocations_.emplace(allocation->size(), std::move(allocation));
19777
}
19878

199-
bool BufferedAllocator::IsAllocThreadSafe() const { return mtx_ != nullptr; }
200-
201-
const std::vector<size_t>& BufferedAllocator::GetDivisionPlan() const {
202-
return division_plan_;
79+
bool BufferedAllocator::IsAllocThreadSafe() const {
80+
return this->underlying_allocator_->IsAllocThreadSafe();
20381
}
20482

20583
} // namespace allocation

‎paddle/fluid/memory/allocation/buffered_allocator.h

Lines changed: 5 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include <memory>
2020
#include <vector>
2121
#include "paddle/fluid/memory/allocation/allocator.h"
22+
#include "paddle/fluid/platform/lock_guard_ptr.h"
2223

2324
namespace paddle {
2425
namespace memory {
@@ -32,9 +33,6 @@ class BufferedAllocator : public UnmanagedAllocator {
3233
public:
3334
explicit BufferedAllocator(std::unique_ptr<Allocator>&& allocator);
3435

35-
BufferedAllocator(std::unique_ptr<Allocator>&& allocator,
36-
const std::vector<size_t>& division_plan);
37-
3836
~BufferedAllocator();
3937

4038
std::unique_ptr<Allocation> Allocate(
@@ -44,31 +42,14 @@ class BufferedAllocator : public UnmanagedAllocator {
4442

4543
bool IsAllocThreadSafe() const override;
4644

47-
const std::vector<size_t>& GetDivisionPlan() const;
48-
49-
void Flush();
45+
// only used in unittest
46+
inline void ClearCache() { FreeCache(-1UL); }
5047

5148
private:
52-
void InitAndEnforceCheck(std::unique_ptr<Allocator>&& allocator,
53-
const std::vector<size_t>& division_plan);
54-
55-
void InsertAllocation(std::unique_ptr<Allocation>&& allocation);
56-
void InsertAllocationImpl(std::unique_ptr<Allocation>&& allocation);
57-
58-
static bool Match(size_t actual_size, size_t requested_size);
59-
std::unique_ptr<Allocation> RemoveAllocation(size_t size);
60-
std::unique_ptr<Allocation> RemoveAllocationImpl(size_t size);
61-
62-
void FreeAllocations(size_t size);
63-
void FreeAllocationsImpl(size_t size);
64-
65-
void FlushImpl();
66-
67-
size_t GetListIndex(size_t size);
49+
void FreeCache(size_t size);
6850

6951
std::unique_ptr<UnmanagedAllocator> underlying_allocator_;
70-
std::vector<std::multimap<size_t, std::unique_ptr<Allocation>>> allocations_;
71-
std::vector<size_t> division_plan_;
52+
std::multimap<size_t, std::unique_ptr<Allocation>> allocations_;
7253
std::unique_ptr<std::mutex> mtx_;
7354
};
7455

‎paddle/fluid/memory/allocation/buffered_allocator_test.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ TEST(buffered_allocator, lazy_free) {
124124

125125
{
126126
underlying_allocator->ResetCounter();
127-
allocator->Flush();
127+
allocator->ClearCache();
128128
ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
129129
ASSERT_EQ(underlying_allocator->GetFreeCount(), kTwo);
130130
}

‎paddle/fluid/memory/malloc.cc

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,10 @@ DEFINE_bool(init_allocated_mem, false,
3030
"during unit testing.");
3131
DECLARE_double(fraction_of_gpu_memory_to_use);
3232

33-
DEFINE_bool(use_legacy_allocator, true,
34-
"Whether to use the legacy allocator. If the new allocators have"
35-
"been well tested, we should remove these flag.");
33+
DEFINE_string(
34+
allocator_strategy, "legacy",
35+
"The allocation strategy. Legacy means the original allocator of Fluid."
36+
"New means the experimental allocators of Fluid. in [legacy, new]");
3637

3738
namespace paddle {
3839
namespace memory {
@@ -274,15 +275,11 @@ size_t Usage::operator()(const platform::CUDAPinnedPlace& cuda_pinned) const {
274275
#endif
275276
}
276277

277-
size_t memory_usage(const platform::Place& p) {
278-
return boost::apply_visitor(Usage(), p);
279-
}
280-
281278
class LegacyAllocation : public Allocation {
282279
public:
283280
using Allocation::Allocation;
284281

285-
~LegacyAllocation() {
282+
~LegacyAllocation() final {
286283
boost::apply_visitor(FreeVisitor(this->ptr()), this->place());
287284
}
288285
};
@@ -291,7 +288,7 @@ class LegacyAllocation : public Allocation {
291288

292289
std::shared_ptr<Allocation> AllocShared(const platform::Place& place,
293290
size_t size, Allocator::Attr attr) {
294-
if (FLAGS_use_legacy_allocator) {
291+
if (FLAGS_allocator_strategy == "legacy") {
295292
void* p = boost::apply_visitor(legacy::AllocVisitor(size), place);
296293
return std::shared_ptr<Allocation>(
297294
new legacy::LegacyAllocation(p, size, place));
@@ -303,7 +300,7 @@ std::shared_ptr<Allocation> AllocShared(const platform::Place& place,
303300

304301
std::unique_ptr<Allocation> Alloc(const platform::Place& place, size_t size,
305302
Allocator::Attr attr) {
306-
if (FLAGS_use_legacy_allocator) {
303+
if (FLAGS_allocator_strategy == "legacy") {
307304
void* p = boost::apply_visitor(legacy::AllocVisitor(size), place);
308305
return std::unique_ptr<Allocation>(
309306
new legacy::LegacyAllocation(p, size, place));

‎paddle/fluid/operators/reader/create_recordio_file_reader_op.cc

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
// limitations under the License.
1414

1515
#include "paddle/fluid/operators/reader/reader_op_registry.h"
16+
#include "paddle/fluid/platform/lock_guard_ptr.h"
1617
#include "paddle/fluid/recordio/scanner.h"
1718

1819
namespace paddle {
@@ -33,11 +34,7 @@ class RecordIOFileReader : public framework::FileReader {
3334

3435
protected:
3536
void ReadNextImpl(std::vector<framework::LoDTensor>* out) override {
36-
std::unique_ptr<std::lock_guard<std::mutex>> guard;
37-
if (ThreadSafe) {
38-
guard.reset(new std::lock_guard<std::mutex>(*mutex_));
39-
}
40-
37+
platform::LockGuardPtr<std::mutex> guard(mutex_);
4138
bool ok = framework::ReadFromRecordIO(&scanner_, dev_ctx_, out);
4239
if (!ok) {
4340
out->clear();

0 commit comments

Comments
 (0)