Correct scratch in use computation.

Old logic did not consider memory held in the scratch cache to be
free when deciding whether or not to reclaim.

Change-Id: I7f7c7549c72d743edbf7c53489fe9a453dc4177a

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

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<KernelMutex> 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<KernelMutex> lock(&scratch_lock_);
     // Check scratch cache
+    scratch.large = large;
     if (scratch_cache_.alloc(scratch)) return;
 
     // Attempt new allocation.