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_;