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.


[ROCm/ROCR-Runtime commit: aa2f98e6f9]

projects/rocr-runtime/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_;
 

projects/rocr-runtime/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_;
 

projects/rocr-runtime/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_);
   }
 

projects/rocr-runtime/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_; }
 

projects/rocr-runtime/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);
 

projects/rocr-runtime/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);

projects/rocr-runtime/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;

projects/rocr-runtime/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

projects/rocr-runtime/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];