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: 8ebf5f9c48]

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<size_t, node> reserved_;
 };
 
 }  // namespace AMD

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<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;
+  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.

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