Skip to content

Commit

Permalink
[CustomDevice] remove stream_safe_allocator when the allocater strage…
Browse files Browse the repository at this point in the history 
… is naive_best_fit (PaddlePaddle#58035)
  • Loading branch information
@ronny1996 @Frida-a
ronny1996 authored and Frida-a committed Oct 14, 2023
1 parent 5359962 commit b767335
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 59 deletions.
171 changes: 115 additions & 56 deletions paddle/fluid/memory/allocation/allocator_facade.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,10 +212,6 @@ class AllocatorFacadePrivate {
platform::CustomPlace(dev_type, dev_id));
}
}
if (FLAGS_use_stream_safe_cuda_allocator) {
WrapStreamSafeCustomDeviceAllocatorForDefault();
is_stream_safe_cuda_allocator_used_ = true;
}
#endif
break;
}
Expand Down Expand Up @@ -576,24 +572,30 @@ class AllocatorFacadePrivate {
#endif

#ifdef PADDLE_WITH_CUSTOM_DEVICE
bool HasCustomDevice(const platform::CustomPlace& place,
phi::stream::stream_t stream) {
bool HasCustomDeviceAllocator(const platform::CustomPlace& place,
phi::stream::stream_t stream) {
auto it = custom_device_allocators_.find(place);
if (it == custom_device_allocators_.end()) {
return false;
}
auto& allocator_map = it->second;
const std::map<phi::stream::stream_t, std::shared_ptr<Allocator>>&
allocator_map = it->second;
return allocator_map.find(stream) != allocator_map.end();
}

const std::shared_ptr<Allocator>& GetAllocator(
const platform::CustomPlace& place,
phi::stream::stream_t stream,
bool create_if_not_found = false) {
if (stream == GetDefaultStream(place)) {
VLOG(7) << "Get Allocator by passing in a default stream";
return GetAllocator(place, /* A non-zero num to choose allocator_ */ 1);
}

/* shared_lock_guard */ {
std::shared_lock<std::shared_timed_mutex> lock_guard(
custom_device_allocator_mutex_);
if (LIKELY(HasCustomDevice(place, stream))) {
if (LIKELY(HasCustomDeviceAllocator(place, stream))) {
return custom_device_allocators_[place][stream];
} else {
PADDLE_ENFORCE_NE(create_if_not_found,
Expand Down Expand Up @@ -627,17 +629,11 @@ class AllocatorFacadePrivate {
return iter->second;
}

void RecordStream(std::shared_ptr<phi::Allocation> allocation,
phi::stream::stream_t stream) {
std::shared_ptr<StreamSafeCustomDeviceAllocation>
stream_safe_custom_device_allocation =
std::dynamic_pointer_cast<StreamSafeCustomDeviceAllocation>(
allocation);
if (stream_safe_custom_device_allocation != nullptr) {
stream_safe_custom_device_allocation->RecordStream(stream);
} else {
VLOG(6) << "RecordStream for a non-StreamSafeCustomDeviceAllocation";
}
phi::stream::stream_t GetDefaultStream(
const platform::CustomPlace& place) const {
const std::shared_ptr<StreamSafeCustomDeviceAllocator>& allocator =
GetDefaultStreamSafeCustomDeviceAllocator(place);
return allocator->GetDefaultStream();
}

void SetDefaultStream(const platform::CustomPlace& place,
Expand All @@ -662,6 +658,34 @@ class AllocatorFacadePrivate {
<< ") in " << place;
}

void RecordStream(std::shared_ptr<phi::Allocation> allocation,
phi::stream::stream_t stream) {
std::shared_ptr<StreamSafeCustomDeviceAllocation>
stream_safe_custom_device_allocation =
std::dynamic_pointer_cast<StreamSafeCustomDeviceAllocation>(
allocation);
if (stream_safe_custom_device_allocation != nullptr) {
stream_safe_custom_device_allocation->RecordStream(stream);
} else {
VLOG(6) << "RecordStream for a non-StreamSafeCustomDeviceAllocation";
}
}

phi::stream::stream_t GetStream(
const std::shared_ptr<phi::Allocation>& allocation) const {
const std::shared_ptr<StreamSafeCustomDeviceAllocation>
stream_safe_custom_device_allocation =
std::dynamic_pointer_cast<StreamSafeCustomDeviceAllocation>(
allocation);
if (stream_safe_custom_device_allocation != nullptr) {
return stream_safe_custom_device_allocation->GetOwningStream();
}

VLOG(6) << "GetStream for a non-StreamSafeCustomDeviceAllocation";
return static_cast<phi::CustomContext*>(
platform::DeviceContextPool::Instance().Get(allocation->place()))
->stream();
}
#endif

private:
Expand Down Expand Up @@ -1108,10 +1132,41 @@ class AllocatorFacadePrivate {
allocators_[p] = std::make_shared<NaiveBestFitAllocator>(p);
}

void InitNaiveBestFitCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
std::shared_ptr<Allocator> CreateCustomDeviceAllocator(
platform::CustomPlace p) {
return std::make_shared<CustomAllocator>(p);
}

void InitStreamSafeCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
PADDLE_ENFORCE_EQ(
strategy_,
AllocatorStrategy::kAutoGrowth,
platform::errors::Unimplemented(
"Only support auto-growth strategey for "
"StreamSafeCustomDeviceAllocator, "
"the allocator strategy %d is unsupported for multi-stream",
static_cast<int>(strategy_)));
if (LIKELY(!HasCustomDeviceAllocator(p, stream))) {
VLOG(8) << "Init StreamSafeCustomDeviceAllocator for stream " << stream
<< " in place " << p;
InitAutoGrowthCustomDeviceAllocator(p, stream);
WrapStreamSafeCustomDeviceAllocator(p, stream);
}
}

void InitAutoGrowthCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
auto chunk_size = FLAGS_auto_growth_chunk_size_in_mb << 20;
VLOG(4) << "FLAGS_auto_growth_chunk_size_in_mb is "
<< FLAGS_auto_growth_chunk_size_in_mb;

auto custom_allocator =
std::make_shared<paddle::memory::allocation::CustomAllocator>(p);
auto alignment = phi::DeviceManager::GetMinChunkSize(p);
custom_device_allocators_[p][stream] =
std::make_shared<NaiveBestFitAllocator>(p);
std::make_shared<AutoGrowthBestFitAllocator>(
custom_allocator, alignment, chunk_size, allow_free_idle_chunk_);
}

void InitAutoGrowthCustomDeviceAllocator(platform::CustomPlace p,
Expand Down Expand Up @@ -1146,39 +1201,13 @@ class AllocatorFacadePrivate {
}
}

void InitAutoGrowthCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
auto chunk_size = FLAGS_auto_growth_chunk_size_in_mb << 20;
VLOG(4) << "FLAGS_auto_growth_chunk_size_in_mb is "
<< FLAGS_auto_growth_chunk_size_in_mb;

auto custom_allocator =
std::make_shared<paddle::memory::allocation::CustomAllocator>(p);
auto alignment = phi::DeviceManager::GetMinChunkSize(p);
custom_device_allocators_[p][stream] =
std::make_shared<AutoGrowthBestFitAllocator>(
custom_allocator, alignment, chunk_size, allow_free_idle_chunk_);
}

void WrapStreamSafeCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
std::shared_ptr<Allocator>& allocator =
custom_device_allocators_[p][stream];
allocator =
std::make_shared<StreamSafeCustomDeviceAllocator>(allocator, p, stream);
}

void InitStreamSafeCustomDeviceAllocator(platform::CustomPlace p,
phi::stream::stream_t stream) {
VLOG(8) << "Init CustomDevice allocator for stream " << stream
<< " in place " << p;
if (strategy_ == AllocatorStrategy::kAutoGrowth) {
InitAutoGrowthCustomDeviceAllocator(p, stream);
} else {
InitNaiveBestFitCustomDeviceAllocator(p, stream);
}
WrapStreamSafeCustomDeviceAllocator(p, stream);
}
#endif

void InitSystemAllocators() {
Expand Down Expand Up @@ -1419,12 +1448,20 @@ AllocationPtr AllocatorFacade::Alloc(const platform::Place& place,
const phi::Stream& stream) {
#ifdef PADDLE_WITH_CUSTOM_DEVICE
if (platform::is_custom_place(place)) {
if (!GetPrivate()->IsStreamSafeCUDAAllocatorUsed()) {
VLOG(6) << "Warning: StreamSafeCustomDeviceAllocator is not used!";
return Alloc(place, size);
}
platform::CustomPlace p(place);
phi::stream::stream_t s =
reinterpret_cast<phi::stream::stream_t>(stream.id());
return GetPrivate()
->GetAllocator(p, s, /* create_if_not_found = */ true)
->Allocate(size);
if (LIKELY(size > 0 && FLAGS_use_system_allocator == false)) {
phi::stream::stream_t s =
reinterpret_cast<phi::stream::stream_t>(stream.id());
return GetPrivate()
->GetAllocator(p, s, /* create_if_not_found = */ true)
->Allocate(size);
} else {
return GetPrivate()->GetAllocator(p, size)->Allocate(size);
}
}
#endif
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
Expand Down Expand Up @@ -1552,20 +1589,42 @@ void AllocatorFacade::RemoveMemoryPoolOfCUDAGraph(int64_t id) {
#endif

#ifdef PADDLE_WITH_CUSTOM_DEVICE
uint64_t AllocatorFacade::Release(const platform::CustomPlace& place,
phi::stream::stream_t stream) {
AllocatorFacadePrivate* m = GetPrivate();
if (!m->IsStreamSafeCUDAAllocatorUsed()) {
VLOG(6) << "Warning: StreamSafeCustomDeviceAllocator is not used!";
return Release(place);
}

return m->GetAllocator(place, stream)->Release(place);
}

void AllocatorFacade::RecordStream(std::shared_ptr<phi::Allocation> allocation,
phi::stream::stream_t stream) {
GetPrivate()->RecordStream(allocation, stream);
}

const std::shared_ptr<Allocator>& AllocatorFacade::GetAllocator(
const platform::Place& place, phi::stream::stream_t stream) {
AllocatorFacadePrivate* m = GetPrivate();

if (!m->IsStreamSafeCUDAAllocatorUsed()) {
VLOG(6) << "Warning: StreamSafeCustomDeviceAllocator is not used!";
return GetAllocator(place);
}

if (platform::is_custom_place(place) && FLAGS_use_system_allocator == false) {
return m->GetAllocator(place,
stream,
/*create_if_not_found=*/true);
}
return m->GetAllocator(place, /* A non-zero num to choose allocator_ */ 1);
}

void AllocatorFacade::RecordStream(std::shared_ptr<phi::Allocation> allocation,
phi::stream::stream_t stream) {
GetPrivate()->RecordStream(allocation, stream);
phi::stream::stream_t AllocatorFacade::GetStream(
const std::shared_ptr<phi::Allocation>& allocation) const {
return GetPrivate()->GetStream(allocation);
}

void AllocatorFacade::SetDefaultStream(const platform::CustomPlace& place,
Expand Down
9 changes: 6 additions & 3 deletions paddle/fluid/memory/allocation/allocator_facade.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,11 +97,14 @@ class AllocatorFacade {
#endif

#ifdef PADDLE_WITH_CUSTOM_DEVICE
uint64_t Release(const platform::CustomPlace& place,
phi::stream::stream_t stream);
void RecordStream(std::shared_ptr<Allocation> allocation,
phi::stream::stream_t stream);
const std::shared_ptr<Allocator>& GetAllocator(const platform::Place& place,
phi::stream::stream_t stream);
void RecordStream(std::shared_ptr<phi::Allocation> allocation,
phi::stream::stream_t stream);

phi::stream::stream_t GetStream(
const std::shared_ptr<Allocation>& allocation) const;
void SetDefaultStream(const platform::CustomPlace& place,
phi::stream::stream_t stream);
#endif
Expand Down

0 comments on commit b767335

Please sign in to comment.