rocr: Update for new async scratch reclaim

Updating ROCr code to match new handshake protocol with CP FW for
asynchronous scratch reclaim.
Increase previous limits when scratch reclaim feature is available.

runtime/hsa-runtime/core/inc/amd_aql_queue.h

@@ -264,6 +264,9 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   /// @brief Halt the queue without destroying it or fencing memory.
   void Suspend();
 
+  /// @brief Resume the queue.
+  void Resume();
+
   /// @brief Handle insufficient scratch
   void HandleInsufficientScratch(hsa_signal_value_t& error_code, hsa_signal_value_t& waitVal,
                                  bool& changeWait);
@@ -327,6 +330,10 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   // CU mask lock
   KernelMutex mask_lock_;
 
+  // Mutex to prevent AsyncReclaimScratch and HandleInsufficientScratch from
+  // happening at the same time.
+  KernelMutex scratch_lock_;
+
   // Current CU mask
   std::vector<uint32_t> cu_mask_;
 

runtime/hsa-runtime/core/inc/amd_gpu_agent.h

@@ -435,11 +435,12 @@ class GpuAgent : public GpuAgentInt {
 
   // @brief Returns true if scratch reclaim is enabled
   __forceinline bool AsyncScratchReclaimEnabled() const override {
+    const uint32_t GFX94X_MIN_CP_FW_VERSION_REQUIRED = 177;
     // TODO: Need to update min CP FW ucode version once it is released
     return (core::Runtime::runtime_singleton_->flag().enable_scratch_async_reclaim() &&
             supported_isas()[0]->GetMajorVersion() == 9 &&
-            supported_isas()[0]->GetMinorVersion() == 4 &&
-            properties_.EngineId.ui32.uCode > 999);
+            supported_isas()[0]->GetMinorVersion() >= 4 &&
+            properties_.EngineId.ui32.uCode >= GFX94X_MIN_CP_FW_VERSION_REQUIRED);
   };
 
   hsa_status_t SetAsyncScratchThresholds(size_t use_once_limit) override;
@@ -619,7 +620,7 @@ class GpuAgent : public GpuAgentInt {
 
   // @brief Mappings from doorbell index to queue, for trap handler.
   // Correlates with output of s_sendmsg(MSG_GET_DOORBELL) for queue identification.
-  amd_queue_t** doorbell_queue_map_;
+  amd_queue_v2_t** doorbell_queue_map_;
 
   // @brief The GPU memory bus width in bit.
   uint32_t memory_bus_width_;
@@ -716,7 +717,7 @@ class GpuAgent : public GpuAgentInt {
   uint32_t sdma_blit_used_mask_;
 
   // Scratch limit thresholds when async scratch is enabled.
-  size_t scratch_limit_async_threshold_;
+  uint64_t scratch_limit_async_threshold_;
 
   ScratchCache scratch_cache_;
 

runtime/hsa-runtime/core/inc/intercept_queue.h

@@ -64,7 +64,7 @@ class QueueWrapper : public Queue {
   std::unique_ptr<Queue> wrapped;
 
   explicit QueueWrapper(std::unique_ptr<Queue> queue) : Queue(), wrapped(std::move(queue)) {
-    memcpy(&amd_queue_, &wrapped->amd_queue_, sizeof(amd_queue_t));
+    memcpy(&amd_queue_, &wrapped->amd_queue_, sizeof(amd_queue_));
     wrapped->set_public_handle(wrapped.get(), public_handle_);
   }
 

runtime/hsa-runtime/core/inc/queue.h

@@ -151,10 +151,10 @@ struct AqlPacket {
 
 class Queue;
 
-/// @brief Helper structure to simplify conversion of amd_queue_t and
+/// @brief Helper structure to simplify conversion of amd_queue_v2_t and
 /// core::Queue object.
 struct SharedQueue {
-  amd_queue_t amd_queue;
+  amd_queue_v2_t amd_queue;
   Queue* core_queue;
 };
 
@@ -378,7 +378,7 @@ class Queue : public Checked<0xFA3906A679F9DB49>, private LocalQueue {
   static void DefaultErrorHandler(hsa_status_t status, hsa_queue_t* source, void* data);
 
   // Handle of AMD Queue struct
-  amd_queue_t& amd_queue_;
+  amd_queue_v2_t& amd_queue_;
 
   hsa_queue_t* public_handle() const { return public_handle_; }
 

runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp

@@ -73,6 +73,9 @@
 namespace rocr {
 namespace AMD {
 
+
+#define SCRATCH_ALT_RATIO 4
+
 AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, ScratchInfo& scratch,
                    core::HsaEventCallback callback, void* err_data, bool is_kv)
     : Queue(agent->node_id(), agent->isMES() ? (MemoryRegion::AllocateGTTAccess | MemoryRegion::AllocateNonPaged) : 0),
@@ -226,7 +229,7 @@ AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, Scr
   if (queue_scratch_.async_reclaim) {
     queue_scratch_.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
     queue_scratch_.use_alt_limit = core::Runtime::runtime_singleton_->flag().enable_scratch_alt()
-        ? (queue_scratch_.use_once_limit / 4)
+        ? (queue_scratch_.use_once_limit / SCRATCH_ALT_RATIO)
         : 0;
   }
 
@@ -301,8 +304,12 @@ AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, Scr
   queue_id_ = queue_rsrc.QueueId;
   MAKE_NAMED_SCOPE_GUARD(QueueGuard, [&]() { hsaKmtDestroyQueue(queue_id_); });
 
-  amd_queue_.scratch_last_used_index = UINT64_MAX;
-  amd_queue_.alt_scratch_last_used_index = UINT64_MAX;
+  amd_queue_.scratch_max_use_index = UINT64_MAX;
+  amd_queue_.alt_scratch_max_use_index = UINT64_MAX;
+
+  // Set flag to notify CP FW that SW supports the new amd_queue_v2
+  if (agent_->AsyncScratchReclaimEnabled())
+    amd_queue_.caps |= AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM;
 
   // On the first queue creation, reserve some scratch memory on this agent.
   agent_->ReserveScratch();
@@ -819,6 +826,14 @@ void AqlQueue::Suspend() {
   assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtUpdateQueue failed.");
 }
 
+void AqlQueue::Resume() {
+  if (suspended_) {
+    suspended_ = false;
+    auto err = hsaKmtUpdateQueue(queue_id_, 100, priority_, ring_buf_, ring_buf_alloc_bytes_, NULL);
+    assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtUpdateQueue failed.");
+  }
+}
+
 hsa_status_t AqlQueue::Inactivate() {
   bool active = active_.exchange(false, std::memory_order_relaxed);
   if (active) {
@@ -845,12 +860,14 @@ void AqlQueue::CheckScratchLimits() {
 
   scratch.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
   scratch.use_alt_limit = core::Runtime::runtime_singleton_->flag().enable_scratch_alt()
-      ? (scratch.use_once_limit / 4)
+      ? (scratch.use_once_limit / SCRATCH_ALT_RATIO)
       : 0;
 
-  if (scratch.main_size > scratch.use_once_limit) AsyncReclaimMainScratch();
+  if (scratch.main_size > scratch.use_once_limit)
+    AsyncReclaimMainScratch();
 
-  if (scratch.alt_size > scratch.use_alt_limit) AsyncReclaimAltScratch();
+  if (scratch.alt_size > scratch.use_alt_limit)
+    AsyncReclaimAltScratch();
 
   return;
 }
@@ -862,34 +879,60 @@ void AqlQueue::FreeMainScratchSpace() {
   scratch.main_size_per_thread = 0;
   scratch.main_queue_process_offset = 0;
   InitScratchSRD();
-
-  HSA::hsa_signal_store_relaxed(amd_queue_.queue_inactive_signal, 0);
 }
 
 void AqlQueue::AsyncReclaimMainScratch() {
-  auto& scratch = queue_scratch_;
-  if (!scratch.async_reclaim || !scratch.main_size) return;
+  auto getMaxMainScratchUseIndex = [&]() {
+    uint64_t max = 0;
+    for (int i = 0; i < agent_->properties().NumXcc; i++) {
+      if (amd_queue_.scratch_last_used_index[i].main > max)
+        max = amd_queue_.scratch_last_used_index[i].main;
+    }
+    return max;
+  };
+
+  auto& scratch = queue_scratch_;
+  if (!scratch.async_reclaim || !scratch.main_size) {
+    return;
+  }
+
+  assert((amd_queue_.caps & AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM) &&
+          "This version of CP FW should support async scratch, but flag is not set");
 
-  // Notify CP that we are trying to reclaim scratch. CP will assume scratch is reclaimed on next
-  // dispatch
   tool::notify_event_scratch_async_reclaim_start(public_handle(),
                                                  HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_NONE);
 
-  amd_queue_.scratch_wave64_lane_byte_size = 0;
-  uint64_t last_used =
-      atomic::Exchange(&amd_queue_.scratch_last_used_index, UINT64_MAX, std::memory_order_relaxed);
+  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
 
-  // Wait for scratch to be idle.
-  while (true) {
-    uint64_t last = amd_queue_.scratch_last_used_index;
+  // Unmap the queue. CP will check amd_queue_ fields on re-map
+  Suspend();
 
-    if (std::min(last, last_used) < amd_queue_.read_dispatch_id) {
-      FreeMainScratchSpace();
-      tool::notify_event_scratch_async_reclaim_end(public_handle(),
-                                                   HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_NONE);
-      return;
-    }
+  /*
+   * amd_queue_.scratch_last_used_index[*].main is updated by CP FW every time a
+   * dispatch packet is launched and it needs scratch memory.
+   * If amd_queue_.scratch_last_used_index[*].main > amd_queue_.read_dispatch_id
+   * then this XCC is currently running a dispatch that uses scratch.
+   * Setting max_scratch_use_index to max(amd_queue_.scratch_last_used_index[*].main)
+   * prevents CP from trying to use main-scratch after
+   * amd_queue_.scratch_max_use_index. If CP sees a dispatch that needs scratch,
+   * it will raise a new signal. CP may use alt-scratch in the meantime.
+   */
+  amd_queue_.scratch_max_use_index = getMaxMainScratchUseIndex();
+
+
+  Resume();
+
+  // If current dispatch is using scratch, wait for it to finish
+  while (amd_queue_.scratch_max_use_index > amd_queue_.read_dispatch_id) {
+    //TODO: if mwaitx supported, //mwaitx(amd_queue_.read_dispatch_id);
+    os::YieldThread();
   }
+
+  FreeMainScratchSpace();
+  tool::notify_event_scratch_async_reclaim_end(public_handle(),
+                                                HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_NONE);
+
+  return;
 }
 
 void AqlQueue::FreeAltScratchSpace() {
@@ -899,34 +942,52 @@ void AqlQueue::FreeAltScratchSpace() {
   scratch.alt_size_per_thread = 0;
   scratch.alt_queue_process_offset = 0;
   InitScratchSRD();
-
-  HSA::hsa_signal_store_relaxed(amd_queue_.queue_inactive_signal, 0);
 }
 
 void AqlQueue::AsyncReclaimAltScratch() {
-  auto& scratch = queue_scratch_;
-  if (!scratch.async_reclaim || !scratch.alt_size) return;
+  /*
+   * See AsyncReclaimMainScratch() for scratch reclaim handshake protocol with
+   * CP FW.
+   */
+  auto getMaxAltScratchUseIndex = [&]() {
+    uint64_t max = 0;
+    for (int i = 0; i < agent_->properties().NumXcc; i++) {
+      if (amd_queue_.scratch_last_used_index[i].alt > max)
+        max = amd_queue_.scratch_last_used_index[i].alt;
+    }
+    return max;
+  };
+
+  auto& scratch = queue_scratch_;
+  if (!scratch.async_reclaim || !scratch.alt_size) {
+    return;
+  }
+
+  assert((amd_queue_.caps & AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM) &&
+          "This version of CP FW should support async scratch, but flag is not set");
 
-  // Notify CP that we are trying to reclaim scratch. CP will assume scratch is reclaimed on next
-  // dispatch
   tool::notify_event_scratch_async_reclaim_start(public_handle(),
                                                  HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT);
 
-  amd_queue_.alt_scratch_wave64_lane_byte_size = 0;
-  uint64_t last_used = atomic::Exchange(
-      &amd_queue_.alt_scratch_last_used_index, UINT64_MAX,
-      std::memory_order_relaxed);  // TODO: Confirm this is the correct memory order
-  // Wait for scratch to be idle.
-  while (true) {
-    uint64_t last = amd_queue_.alt_scratch_last_used_index;
+  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
 
-    if (std::min(last, last_used) < amd_queue_.read_dispatch_id) {
-      FreeAltScratchSpace();
-      tool::notify_event_scratch_async_reclaim_end(public_handle(),
-                                                   HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT);
-      return;
-    }
+  // Unmap the queue. CP will check amd_queue_ fields on re-map
+  Suspend();
+
+  amd_queue_.alt_scratch_max_use_index = getMaxAltScratchUseIndex();
+
+  Resume();
+
+  // If current dispatch is using alt scratch, wait for it to finish
+  while (amd_queue_.alt_scratch_max_use_index > amd_queue_.read_dispatch_id) {
+    //DYSDEBUG TODO: if mwaitx supported, //mwaitx(amd_queue_.read_dispatch_id);
+    os::YieldThread();
   }
+
+  FreeAltScratchSpace();
+  tool::notify_event_scratch_async_reclaim_end(public_handle(),
+                                                HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT);
+  return;
 }
 
 void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
@@ -943,9 +1004,11 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
    *
    * //Default values:
    * size_t use_once_limit = 128 MB      // When async reclaim not supported
-   *                       = 1GB per-XCC // When async reclaim is supported
+   *                                     // DEFAULT_SCRATCH_SINGLE_LIMIT
+   *                       = 3GB per-XCC // When async reclaim is supported
+   *                                     // DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC
    *
-   * size_t use_alt_limit  = 256 MB per-XCC
+   * size_t use_alt_limit  = 768 MB per-XCC // use_once_limit/SCRATCH_ALT_RATIO
    *
    * if (async-scratch-reclaim-supported
    *     && dispatch_slots < max_scratch_slots
@@ -1043,8 +1106,9 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
     return AlignUp(cu_count, engines) * agent_->properties().MaxSlotsScratchCU;
   };
 
-  assert((!scratch.async_reclaim || (amd_queue_.caps & AMD_QUEUE_CAPS_ASYNC_RECLAIM)) &&
-         "Asynchronous scratch reclaim capability not set, but this FW version should support it");
+  assert(core::Runtime::runtime_singleton_->flag().enable_scratch_async_reclaim() &&
+         (!scratch.async_reclaim || (amd_queue_.caps & AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM)) &&
+          "Asynchronous scratch reclaim capability not set, but this FW version should support it");
 
   scratch.cooperative = (amd_queue_.hsa_queue.type == HSA_QUEUE_TYPE_COOPERATIVE);
 
@@ -1070,6 +1134,8 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
   const uint64_t device_size = size_per_thread * lanes_per_wave * device_slots;
   const uint64_t dispatch_size = size_per_thread * lanes_per_wave * dispatch_slots;
 
+  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
+
   // scratch.use_alt_limit will be 0 if alt scratch is not supported or disabled
   if (dispatch_size < scratch.use_alt_limit && dispatch_slots < device_slots) {
     // Try to use ALT scratch
@@ -1086,8 +1152,13 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
       scratch.alt_dispatch_limit_y = pkt->dispatch.grid_size_y;
       scratch.alt_dispatch_limit_z = pkt->dispatch.grid_size_z;
 
-      // Update queue SRD
       InitScratchSRD();
+      /*
+       * Indicate to CP FW that any dispatch may use alt scratch memory.
+       * If ROCr wants to reclain scratch memory, it will set
+       * amd_queue_.alt_scratch_max_use_index to a lower value
+       */
+      amd_queue_.alt_scratch_max_use_index = UINT64_MAX;
       // Restart the queue.
       HSA::hsa_signal_store_screlease(amd_queue_.queue_inactive_signal, 0);
       tool::notify_event_scratch_alloc_end(public_handle(), HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT,
@@ -1136,16 +1207,28 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
   } else if (scratch.alt_size && scratch.main_size > scratch.alt_size) {
     // Not using use-scratch-once, and dispatches that would fit in alt-scratch would also fit in
     // main scratch. No need for alt-scratch.
-    AsyncReclaimAltScratch();
+    tool::notify_event_scratch_async_reclaim_start(public_handle(),
+                                                 HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT);
+    FreeAltScratchSpace();
+    tool::notify_event_scratch_async_reclaim_end(public_handle(),
+                                                 HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_ALT);
   }
 
   // Reset scratch memory related entities for the queue
   InitScratchSRD();
+  /*
+   * Indicate to CP FW that any dispatch may use alt scratch memory.
+   * If ROCr wants to reclain scratch memory, it will set
+   * amd_queue_.alt_scratch_max_use_index to a lower value
+   */
+  amd_queue_.scratch_max_use_index = UINT64_MAX;
+
   // Restart the queue.
   HSA::hsa_signal_store_screlease(amd_queue_.queue_inactive_signal, 0);
 
   auto alloc_flag = (scratch.large) ? HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_USE_ONCE
                                     : HSA_AMD_EVENT_SCRATCH_ALLOC_FLAG_NONE;
+
   tool::notify_event_scratch_alloc_end(public_handle(), alloc_flag, dispatch_id, scratch.main_size,
                                        dispatch_slots);
 

runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp

@@ -86,7 +86,6 @@
 #define DEFAULT_SCRATCH_BYTES_PER_THREAD 2048
 #define MAX_WAVE_SCRATCH 8387584  // See COMPUTE_TMPRING_SIZE.WAVESIZE
 #define MAX_NUM_DOORBELLS 0x400
-#define DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC (1 << 30)  // 1 GB
 
 namespace rocr {
 
@@ -556,12 +555,17 @@ void GpuAgent::InitScratchPool() {
 }
 
 void GpuAgent::InitAsyncScratchThresholds() {
+  if (!AsyncScratchReclaimEnabled()) return;
+
   scratch_limit_async_threshold_ =
       core::Runtime::runtime_singleton_->flag().scratch_single_limit_async();
 
-  if (!scratch_limit_async_threshold_)
+  if (!scratch_limit_async_threshold_) {
+    // User did not set env var HSA_SCRATCH_SINGLE_LIMIT_ASYNC
     scratch_limit_async_threshold_ =
-        DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC * properties().NumXcc;
+      core::Runtime::runtime_singleton_->flag().DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC *
+      (uint64_t)(properties().NumXcc);
+  }
 }
 
 void GpuAgent::ReserveScratch()
@@ -1769,8 +1773,9 @@ void GpuAgent::AcquireQueueMainScratch(ScratchInfo& scratch) {
   bool use_reclaim = true;
 
   large = (scratch.main_size > scratch.use_once_limit) ||
-      ((scratch_pool_.size() - scratch_pool_.remaining() - scratch_cache_.free_bytes() +
-        scratch.main_size) > small_limit);
+          (!AsyncScratchReclaimEnabled() &&
+            ((scratch_pool_.size() - scratch_pool_.remaining() - scratch_cache_.free_bytes() +
+             scratch.main_size) > small_limit));
 
   if ((isa_->GetMajorVersion() < 8) ||
       core::Runtime::runtime_singleton_->flag().no_scratch_reclaim()) {
@@ -1905,8 +1910,8 @@ void GpuAgent::AcquireQueueMainScratch(ScratchInfo& scratch) {
       : uintptr_t(scratch.main_queue_base) - uintptr_t(scratch_pool_.base());
 }
 
+/* Should be called with scratch_lock_ */
 void GpuAgent::ReleaseQueueMainScratch(ScratchInfo& scratch) {
-  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
   if (scratch.main_queue_base == nullptr) return;
 
   scratch_cache_.freeMain(scratch);
@@ -1980,8 +1985,8 @@ void GpuAgent::AcquireQueueAltScratch(ScratchInfo& scratch) {
   scratch.alt_queue_process_offset = uintptr_t(scratch.alt_queue_base);
 }
 
+/* Should be called with scratch_lock_ */
 void GpuAgent::ReleaseQueueAltScratch(ScratchInfo& scratch) {
-  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
   if (scratch.alt_queue_base == nullptr) return;
 
   scratch_cache_.freeAlt(scratch);
@@ -2023,7 +2028,6 @@ hsa_status_t GpuAgent::SetAsyncScratchThresholds(size_t use_once_limit) {
     auto aqlQueue = static_cast<AqlQueue*>(iter);
     aqlQueue->CheckScratchLimits();
   }
-
   return HSA_STATUS_SUCCESS;
 }
 
@@ -2232,10 +2236,10 @@ void GpuAgent::BindTrapHandler() {
     AssembleShader("TrapHandler", AssembleTarget::ISA, trap_code_buf_, trap_code_buf_size_);
 
     // Make an empty map from doorbell index to queue.
-    // The trap handler uses this to retrieve a wave's amd_queue_t*.
-    auto doorbell_queue_map_size = MAX_NUM_DOORBELLS * sizeof(amd_queue_t*);
+    // The trap handler uses this to retrieve a wave's amd_queue_v2_t*.
+    auto doorbell_queue_map_size = MAX_NUM_DOORBELLS * sizeof(amd_queue_v2_t*);
 
-    doorbell_queue_map_ = (amd_queue_t**)system_allocator()(doorbell_queue_map_size, 0x1000, 0);
+    doorbell_queue_map_ = (amd_queue_v2_t**)system_allocator()(doorbell_queue_map_size, 0x1000, 0);
     assert(doorbell_queue_map_ != NULL && "Doorbell queue map allocation failed");
 
     memset(doorbell_queue_map_, 0, doorbell_queue_map_size);

runtime/hsa-runtime/core/util/flag.h

@@ -66,8 +66,9 @@ class Flag {
 
   // Lift limit for 2.10 release RCCL workaround. This limit is not used when asynchronous scratch
   // reclaim is supported
-  const size_t DEFAULT_SCRATCH_SINGLE_LIMIT = 146800640;  // small_limit >> 2;
-  const size_t DEFAULT_PCS_MAX_DEVICE_BUFFER_SIZE = 256 * 1024 * 1024;
+  const size_t DEFAULT_SCRATCH_SINGLE_LIMIT = (140 * (1UL<<20));  // small_limit >> 2;
+  const size_t DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC = (3 * (1UL<<30));  // 3 GB
+  const size_t DEFAULT_PCS_MAX_DEVICE_BUFFER_SIZE = (256 * (1UL<<20)); //256 MB
 
   explicit Flag() { Refresh(); }
 
@@ -135,6 +136,9 @@ class Flag {
 
     // On GPUs that support asynchronous scratch reclaim
     // Scratch memory sizes > HSA_SCRATCH_SINGLE_LIMIT_ASYNC will trigger a use-once scheme
+    // Note: This only sets the initial value for the threshold. If
+    // hsa_amd_agent_set_async_scratch_limit is called after initialization, the threshold
+    // will be updated.
     if (os::IsEnvVarSet("HSA_SCRATCH_SINGLE_LIMIT_ASYNC")) {
       var = os::GetEnvVar("HSA_SCRATCH_SINGLE_LIMIT_ASYNC");
       char* end;

runtime/hsa-runtime/inc/amd_hsa_queue.h

@@ -64,14 +64,18 @@ enum amd_queue_properties_t {
 // AMD Queue Capabilities.
 typedef uint32_t amd_queue_capabilities32_t;
 enum amd_queue_capabilities_t {
-  /* Whether this CP queue supports dual-scratch and async-reclaim */
-  AMD_HSA_BITS_CREATE_ENUM_ENTRIES(AMD_QUEUE_CAPS_ASYNC_RECLAIM, 0, 1),
+  /* This version of CP FW supports dual-scratch and async-reclaim */
+  AMD_HSA_BITS_CREATE_ENUM_ENTRIES(AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM, 0, 1),
+
+  /*
+   * This version of ROCr supports async-reclaim and CP FW may access the
+   * V2 fields.
+   */
+  AMD_HSA_BITS_CREATE_ENUM_ENTRIES(AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM, 1, 1),
 };
 
-// Members tagged with "async-reclaim" are ignored by CP FW's that do not support
-// AMD_QUEUE_CAPS_ASYNC_RECLAIM. CP FW's that support async-reclaim also support
-// dual-scratch (alternate scratch).
-
+/* This is the original amd_queue_t structure. The definition is only kept
+ * for reference purposes. This structure should not be used. */
 typedef struct AMD_QUEUE_ALIGN amd_queue_s {
   hsa_queue_t hsa_queue;
   uint32_t caps;
@@ -92,18 +96,59 @@ typedef struct AMD_QUEUE_ALIGN amd_queue_s {
   uint64_t scratch_backing_memory_byte_size;
   uint32_t scratch_wave64_lane_byte_size;
   amd_queue_properties32_t queue_properties;
-  volatile uint64_t scratch_last_used_index;     /* async-reclaim */
+  uint32_t reserved3[2];
   hsa_signal_t queue_inactive_signal;
-  uint32_t reserved4[2];
-  volatile uint64_t alt_scratch_last_used_index; /* async-reclaim */
-  uint64_t alt_scratch_backing_memory_location;  /* async-reclaim */
-  uint64_t alt_scratch_backing_memory_byte_size; /* async-reclaim */
-  uint32_t alt_scratch_dispatch_limit_x;         /* async-reclaim */
-  uint32_t alt_scratch_dispatch_limit_y;         /* async-reclaim */
-  uint32_t alt_scratch_dispatch_limit_z;         /* async-reclaim */
-  uint32_t alt_scratch_wave64_lane_byte_size;    /* async-reclaim */
-  uint32_t alt_compute_tmpring_size;             /* async-reclaim */
-  uint32_t reserved5;
+  uint32_t reserved4[14];
 } amd_queue_t;
 
+/*
+ * AMD_QUEUE Version 2
+ * amd_queue_v2_t is backwards compatible with amd_queue_t structure and can
+ * be used with previous versions of CP FW. The added fields tagged as V2 are
+ * ignored when running previous versions of CP FW.
+ * CP FW will not try to access elements beyond the original 64-bytes
+ * (sizeof(amd_queue_t)) unless the AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM bit is set.
+ */
+
+#define MAX_NUM_XCC 128
+typedef struct scratch_last_used_index_xcc_s {
+  volatile uint64_t main;
+  volatile uint64_t alt;
+} scratch_last_used_index_xcc_t;
+
+typedef struct AMD_QUEUE_ALIGN amd_queue_v2_s {
+  hsa_queue_t hsa_queue;
+  uint32_t caps;
+  uint32_t reserved1[3];
+  volatile uint64_t write_dispatch_id;
+  uint32_t group_segment_aperture_base_hi;
+  uint32_t private_segment_aperture_base_hi;
+  uint32_t max_cu_id;
+  uint32_t max_wave_id;
+  volatile uint64_t max_legacy_doorbell_dispatch_id_plus_1;
+  volatile uint32_t legacy_doorbell_lock;
+  uint32_t reserved2[9];
+  volatile uint64_t read_dispatch_id;
+  uint32_t read_dispatch_id_field_base_byte_offset;
+  uint32_t compute_tmpring_size;
+  uint32_t scratch_resource_descriptor[4];
+  uint64_t scratch_backing_memory_location;
+  uint32_t reserved3[2];
+  uint32_t scratch_wave64_lane_byte_size;
+  amd_queue_properties32_t queue_properties;
+  volatile uint64_t scratch_max_use_index;       /* V2 */
+  hsa_signal_t queue_inactive_signal;
+  volatile uint64_t alt_scratch_max_use_index;  /* V2 */
+  uint32_t alt_scratch_resource_descriptor[4];   /* V2 */
+  uint64_t alt_scratch_backing_memory_location;  /* V2 */
+  uint32_t alt_scratch_dispatch_limit_x;         /* V2 */
+  uint32_t alt_scratch_dispatch_limit_y;         /* V2 */
+  uint32_t alt_scratch_dispatch_limit_z;         /* V2 */
+  uint32_t alt_scratch_wave64_lane_byte_size;    /* V2 */
+  uint32_t alt_compute_tmpring_size;             /* V2 */
+  uint32_t reserved5;
+
+  scratch_last_used_index_xcc_t scratch_last_used_index[MAX_NUM_XCC];
+} amd_queue_v2_t;
+
 #endif // AMD_HSA_QUEUE_H

runtime/hsa-runtime/inc/amd_hsa_signal.h

@@ -71,7 +71,7 @@ typedef struct AMD_SIGNAL_ALIGN amd_signal_s {
   uint64_t start_ts;
   uint64_t end_ts;
   union {
-    amd_queue_t* queue_ptr;
+    amd_queue_v2_t* queue_ptr;
     uint64_t reserved2;
   };
   uint32_t reserved3[2];