diff --git a/runtime/hsa-runtime/core/inc/scratch_cache.h b/runtime/hsa-runtime/core/inc/scratch_cache.h index e85e22660f..1e079babb7 100644 --- a/runtime/hsa-runtime/core/inc/scratch_cache.h +++ b/runtime/hsa-runtime/core/inc/scratch_cache.h @@ -107,7 +107,7 @@ class ScratchCache { ScratchCache& operator=(const ScratchCache& rhs) = delete; ScratchCache& operator=(ScratchCache&& rhs) = delete; - ScratchCache(deallocator_t deallocator) : dealloc(deallocator) {} + ScratchCache(deallocator_t deallocator) : dealloc(deallocator), available_bytes(0) {} ~ScratchCache() { assert(map.empty() && "ScratchCache not empty at shutdown."); } @@ -122,6 +122,7 @@ class ScratchCache { it->second.alloc(); info.queue_base = it->second.base; info.scratch_node = it; + available_bytes -= it->first; return true; } it++; @@ -136,6 +137,7 @@ class ScratchCache { info.queue_base = it->second.base; info.size = it->first; info.scratch_node = it; + available_bytes -= it->first; return true; } it++; @@ -152,6 +154,8 @@ class ScratchCache { return; } it->second.free(); + available_bytes += it->first; + assert(it->first == info.size && "Scratch cache size mismatch."); } bool trim(bool trim_nodes_in_use) { @@ -159,6 +163,7 @@ class ScratchCache { auto it = map.begin(); while (it != map.end()) { if (it->second.isFree()) { + available_bytes -= it->first; dealloc(it->second.base, it->first, it->second.large); auto temp = it; it++; @@ -181,9 +186,14 @@ class ScratchCache { info.scratch_node = it; } + size_t free_bytes() const { + return available_bytes; + } + private: map_t map; deallocator_t dealloc; + size_t available_bytes; }; } // namespace AMD diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp index 1d5205f527..eb1e010cb3 100644 --- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp +++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp @@ -1107,12 +1107,13 @@ void GpuAgent::AcquireQueueScratch(ScratchInfo& scratch) { Limit total bound small scratch allocations to 1/8th of scratch pool and 1/4 of that for a single allocation. */ + 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; bool use_reclaim = true; bool large = (scratch.size > single_limit) || - (scratch_pool_.size() - scratch_pool_.remaining() + scratch.size > small_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; @@ -1140,8 +1141,8 @@ void GpuAgent::AcquireQueueScratch(ScratchInfo& scratch) { // Lambda called in place. // Used to allow exit from nested loops. [&]() { - ScopedAcquire lock(&scratch_lock_); // Check scratch cache + scratch.large = large; if (scratch_cache_.alloc(scratch)) return; // Attempt new allocation.