From d476ff16eb1212e5781066904839b3470ad9f7cf Mon Sep 17 00:00:00 2001 From: David Yat Sin Date: Thu, 6 Apr 2023 02:41:40 +0000 Subject: [PATCH] Adding scratch memory reservation Some applications will keep trying to allocate device memory until the allocation fails. This causes all device memory to be used up and we are then unable to allocate scratch memory for dispatches. Reserve enough memory for 1 small scratch allocation. Change-Id: I968400d41540ba1aca8f28581f229693eec02225 [ROCm/ROCR-Runtime commit: 8ebf5f9c4877c12cebaa8ee6b07bdcdba402a6e1] --- .../hsa-runtime/core/inc/scratch_cache.h | 60 ++++++++++++++++--- .../core/runtime/amd_gpu_agent.cpp | 38 +++++++++--- .../runtime/hsa-runtime/core/util/flag.h | 18 ++++++ 3 files changed, 101 insertions(+), 15 deletions(-) diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/scratch_cache.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/scratch_cache.h index a3a4e20250..3b81b20690 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/scratch_cache.h +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/scratch_cache.h @@ -109,7 +109,7 @@ class ScratchCache { ScratchCache& operator=(const ScratchCache& rhs) = delete; ScratchCache& operator=(ScratchCache&& rhs) = delete; - ScratchCache(deallocator_t deallocator) : dealloc(deallocator), available_bytes(0) {} + ScratchCache(deallocator_t deallocator) : dealloc(deallocator), available_bytes_(0) {} ~ScratchCache() { assert(map.empty() && "ScratchCache not empty at shutdown."); } @@ -124,7 +124,7 @@ class ScratchCache { it->second.alloc(); info.queue_base = it->second.base; info.scratch_node = it; - available_bytes -= it->first; + available_bytes_ -= it->first; return true; } it++; @@ -138,7 +138,7 @@ class ScratchCache { it->second.alloc(); info.queue_base = it->second.base; info.scratch_node = it; - available_bytes -= it->first; + available_bytes_ -= it->first; return true; } it++; @@ -147,6 +147,14 @@ class ScratchCache { } void free(ScratchInfo& info) { + if (info.scratch_node == map.end()) { + // This is reserved scratch memory. Do not de-allocate, just mark it as free. + assert(!reserved_.second.isFree() && "free called when reserved node already free."); + reserved_.second.free(); + available_bytes_ += reserved_.first; + return; + } + assert(!info.scratch_node->second.isFree() && "free called on free scratch node."); auto it = info.scratch_node; if (it->second.trimPending()) { @@ -155,7 +163,7 @@ class ScratchCache { return; } it->second.free(); - available_bytes += it->first; + available_bytes_ += it->first; } bool trim(bool trim_nodes_in_use) { @@ -163,7 +171,7 @@ class ScratchCache { auto it = map.begin(); while (it != map.end()) { if (it->second.isFree()) { - available_bytes -= it->first; + available_bytes_ -= it->first; dealloc(it->second.base, it->first, it->second.large); auto temp = it; it++; @@ -186,14 +194,50 @@ class ScratchCache { info.scratch_node = it; } - size_t free_bytes() const { - return available_bytes; + size_t free_bytes() const { return available_bytes_; } + + void reserve(size_t bytes, void* base) { + assert(!reserved_.first && "Already reserved memory."); + + node n; + n.base = base; + n.large = 0; + + available_bytes_ += bytes; + + reserved_ = std::make_pair(bytes, n); + } + + bool use_reserved(ScratchInfo& info) { + if (!reserved_.second.isFree() || info.size > reserved_.first) { + debug_print("reserved node is already in use or too small (requested:%ld reserved:%ld)\n", + info.size, reserved_.first); + return false; + } + reserved_.second.large = info.large; + reserved_.second.alloc(); + info.queue_base = reserved_.second.base; + // Special case to indicate that this node is reserved memory + info.scratch_node = map.end(); + available_bytes_ -= reserved_.first; + return true; + } + + void free_reserve() { + available_bytes_ -= reserved_.first; + if (reserved_.first) dealloc(reserved_.second.base, reserved_.first, reserved_.second.large); + + reserved_.first = 0; + reserved_.second.base = NULL; + reserved_.second.large = 0; } private: map_t map; deallocator_t dealloc; - size_t available_bytes; + size_t available_bytes_; + + std::pair reserved_; }; } // namespace AMD diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp index 7a3bd2eb6a..6d8efe0618 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp @@ -212,6 +212,8 @@ GpuAgent::~GpuAgent() { } scratch_cache_.trim(true); + scratch_cache_.free_reserve(); + if (scratch_pool_.base() != NULL) { hsaKmtFreeMemory(scratch_pool_.base(), scratch_pool_.size()); } @@ -478,6 +480,18 @@ void GpuAgent::InitScratchPool() { } else { new (&scratch_pool_) SmallHeap(); } + + size_t reserved_sz = core::Runtime::runtime_singleton_->flag().scratch_single_limit(); + if (reserved_sz) { + HSAuint64 alt_va; + void* reserved_base = scratch_pool_.alloc(reserved_sz); + assert(reserved_base && "Could not allocate reserved memory"); + + if (hsaKmtMapMemoryToGPU(reserved_base, reserved_sz, &alt_va) == HSAKMT_STATUS_SUCCESS) + scratch_cache_.reserve(reserved_sz, reserved_base); + else + throw AMD::hsa_exception(HSA_STATUS_ERROR_OUT_OF_RESOURCES, "Reserve scratch memory failed."); + } } void GpuAgent::InitCacheList() { @@ -1345,11 +1359,12 @@ void GpuAgent::AcquireQueueScratch(ScratchInfo& scratch) { */ ScopedAcquire lock(&scratch_lock_); size_t small_limit = scratch_pool_.size() >> 3; - // Lift limit for 2.10 release RCCL workaround. - size_t single_limit = 146800640; //small_limit >> 2; + const size_t single_scratch_limit = + core::Runtime::runtime_singleton_->flag().scratch_single_limit(); bool use_reclaim = true; - bool large = (scratch.size > single_limit) || - (scratch_pool_.size() - scratch_pool_.remaining() - scratch_cache_.free_bytes() + scratch.size > small_limit); + bool large = (scratch.size > single_scratch_limit) || + ((scratch_pool_.size() - scratch_pool_.remaining() - scratch_cache_.free_bytes() + + scratch.size) > small_limit); if ((isa_->GetMajorVersion() < 8) || core::Runtime::runtime_singleton_->flag().no_scratch_reclaim()) { large = false; @@ -1382,7 +1397,7 @@ void GpuAgent::AcquireQueueScratch(ScratchInfo& scratch) { if (scratch_cache_.alloc(scratch)) return; // Attempt new allocation. - for (int i = 0; i < 2; i++) { + for (int i = 0; i < 3; i++) { if (large) scratch.queue_base = scratch_pool_.alloc_high(scratch.size); else @@ -1406,8 +1421,17 @@ void GpuAgent::AcquireQueueScratch(ScratchInfo& scratch) { scratch.queue_base = nullptr; // Release cached scratch and retry. - // First iteration trims unused blocks, second trims all. - scratch_cache_.trim(i == 1); + // First iteration trims unused blocks, second trims all. 3rd uses reserved memory + switch (i) { + case 0: + scratch_cache_.trim(false); + break; + case 1: + scratch_cache_.trim(true); + break; + case 2: + if (scratch_cache_.use_reserved(scratch)) return; + } } // Retry if large may yield needed space. diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h index 28e9815bef..5c5e8f3101 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h @@ -63,6 +63,9 @@ class Flag { static_assert(XNACK_DISABLE == 0, "XNACK_REQUEST enum values improperly changed."); static_assert(XNACK_ENABLE == 1, "XNACK_REQUEST enum values improperly changed."); + // Lift limit for 2.10 release RCCL workaround. + const size_t DEFAULT_SCRATCH_SINGLE_LIMIT = 146800640; // small_limit >> 2; + explicit Flag() { Refresh(); } virtual ~Flag() {} @@ -95,9 +98,21 @@ class Flag { var = os::GetEnvVar("HSA_MAX_QUEUES"); max_queues_ = static_cast(atoi(var.c_str())); + // Maximum amount of scratch mem that can be used per process per gpu var = os::GetEnvVar("HSA_SCRATCH_MEM"); scratch_mem_size_ = atoi(var.c_str()); + // Scratch memory sizes > HSA_SCRATCH_SINGLE_LIMIT will trigger a use-once scheme + // We also reserve HSA_SCRATCH_SINGLE_LIMIT per process per gpu to guarrantee we + // have sufficient memory to for scratch in case user tried to allocate all device + // memory + if (os::IsEnvVarSet("HSA_SCRATCH_SINGLE_LIMIT")) { + var = os::GetEnvVar("HSA_SCRATCH_SINGLE_LIMIT"); + scratch_single_limit_ = atoi(var.c_str()); + } else { + scratch_single_limit_ = DEFAULT_SCRATCH_SINGLE_LIMIT; + } + tools_lib_names_ = os::GetEnvVar("HSA_TOOLS_LIB"); var = os::GetEnvVar("HSA_TOOLS_REPORT_LOAD_FAILURE"); @@ -213,6 +228,8 @@ class Flag { size_t scratch_mem_size() const { return scratch_mem_size_; } + size_t scratch_single_limit() const { return scratch_single_limit_; } + std::string tools_lib_names() const { return tools_lib_names_; } bool disable_image() const { return disable_image_; } @@ -285,6 +302,7 @@ class Flag { uint32_t max_queues_; size_t scratch_mem_size_; + size_t scratch_single_limit_; std::string tools_lib_names_; std::string svm_profile_;