提交 0c25da39 编写于 作者: Y Yu Yang

Refine auto_increment_allocator

上级 ab87a882
...@@ -31,7 +31,7 @@ namespace allocation { ...@@ -31,7 +31,7 @@ namespace allocation {
// invoke its `allocate` method. // invoke its `allocate` method.
// //
// NOTE(yy): The AutoIncrementAllocator will prefer to allocate memory from // NOTE(yy): The AutoIncrementAllocator will prefer to allocate memory from
// the latest sucessful allocator. // the latest successful allocator.
// //
// NOTE(yy): We may need to release an underlying allocator if it allocate // NOTE(yy): We may need to release an underlying allocator if it allocate
// nothing. However, it is generally not useful, since it will make performance // nothing. However, it is generally not useful, since it will make performance
...@@ -76,27 +76,26 @@ class AutoIncrementAllocator : public ManagedAllocator { ...@@ -76,27 +76,26 @@ class AutoIncrementAllocator : public ManagedAllocator {
} }
} catch (...) { } catch (...) {
// if there is another type of allocation, just rethrow it. // if there is another type of allocation, just rethrow it.
std::rethrow_exception(std::current_exception()); throw;
} }
} }
// No suitable allocator
// This happens when the first allocator is exhausted and // This happens when the first allocator is exhausted and
// there are more than 1 allocation requests // there are more than 1 allocation requests
// In this situation, the first allocation request would success // In this situation, the first allocation request would success
// and the second allocation request would fail if we do not use // and the second allocation request would fail if we do not use
// the newly created allocator by the first allocation request. // the newly created allocator by the first allocation request.
for (size_t new_allocator_num = allocator_num_.load(); for (cur = allocator_num; cur < allocator_num_; ++cur) {
allocator_num < new_allocator_num; ++allocator_num) {
try { try {
auto ret = callback(*underlying_allocators_[allocator_num]); auto ret = callback(*underlying_allocators_[cur]);
prev_success_allocator_ = allocator_num; prev_success_allocator_ = cur;
return std::move(ret); return std::move(ret);
} catch (BadAlloc&) { } catch (BadAlloc&) {
} catch (...) { } catch (...) {
std::rethrow_exception(std::current_exception()); throw;
} }
} }
// No suitable allocator
ManagedAllocator* new_allocator; ManagedAllocator* new_allocator;
{ {
...@@ -108,7 +107,7 @@ class AutoIncrementAllocator : public ManagedAllocator { ...@@ -108,7 +107,7 @@ class AutoIncrementAllocator : public ManagedAllocator {
underlying_allocators_[old_size] = creator_(); underlying_allocators_[old_size] = creator_();
new_allocator = underlying_allocators_[old_size].get(); new_allocator = underlying_allocators_[old_size].get();
prev_success_allocator_ = old_size; prev_success_allocator_ = old_size;
allocator_num_.fetch_add(1); ++allocator_num_;
} }
PADDLE_ENFORCE( PADDLE_ENFORCE(
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册