Implement async scratch reclaim

For devices where the CP FW supports asynchronous scratch reclaim, ROCr
is able to claw-back scratch memory that was assigned to an AQL queue.
With that ability, ROCr does not have to rely on using USO
(use-scratch-once) when assigning large amounts of memory to a queue.
If we reach a situation where we are running low on device memory, ROCr
will attempt to claw-back the scratch memory.

Change-Id: Iddf8ec84e37ab8b9fdc58bafbe2b61fe2acb6eb7

runtime/hsa-runtime/core/common/hsa_table_interface.cpp

@@ -1292,6 +1292,10 @@ hsa_status_t HSA_API hsa_amd_vmem_get_alloc_properties_from_handle(
   return amdExtTable->hsa_amd_vmem_get_alloc_properties_from_handle_fn(alloc_handle, pool, type);
 }
 
+hsa_status_t HSA_API hsa_amd_agent_set_async_scratch_limit(hsa_agent_t agent, size_t threshold) {
+  return amdExtTable->hsa_amd_agent_set_async_scratch_limit_fn(agent, threshold);
+}
+
 // Tools only table interfaces.
 namespace rocr {
 

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

@@ -208,6 +208,13 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   /// @brief Enable use of GWS from this queue.
   hsa_status_t EnableGWS(int gws_slot_count);
 
+  /// @brief Update internal scratch limits based on agent limits. If current allocated scratch are
+  /// larger than new limits, perform async-reclaim.
+  void CheckScratchLimits();
+
+  /// @brief Async reclaim main scratch memory
+  void AsyncReclaimMainScratch();
+
  protected:
   bool _IsA(Queue::rtti_t id) const override { return id == &rtti_id_; }
 
@@ -236,6 +243,8 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   void FillComputeTmpRingSize();
   void FillComputeTmpRingSize_Gfx11();
 
+  void FreeMainScratchSpace();
+
   /// @brief Halt the queue without destroying it or fencing memory.
   void Suspend();
 
@@ -314,6 +323,9 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   // Mutex for queue_event_ manipulation
   static KernelMutex queue_lock_;
 
+  // Async scratch single limit - may be modified after init
+  size_t async_scratch_single_limit_;
+
   static int rtti_id_;
 
   // Forbid copying and moving of this object

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

@@ -166,6 +166,14 @@ class GpuAgentInt : public core::Agent {
   //
   // @retval Bus width in MHz.
   virtual uint32_t memory_max_frequency() const = 0;
+
+  // @brief Whether agent supports asynchronous scratch reclaim. Depends on CP FW
+  virtual bool AsyncScratchReclaimEnabled() const = 0;
+
+  // @brief Update the agent's scratch use-once threshold.
+  // Only valid when async scratch reclaim is supported
+  // @retval HSA_STATUS_SUCCESS if successful
+  virtual hsa_status_t SetAsyncScratchThresholds(size_t use_once_limit) = 0;
 };
 
 class GpuAgent : public GpuAgentInt {
@@ -352,6 +360,23 @@ class GpuAgent : public GpuAgentInt {
 
   void ReserveScratch();
 
+  // @brief If agent supports it, release scratch memory for all AQL queues on this agent.
+  void AsyncReclaimScratchQueues();
+
+  // @brief Returns true if scratch reclaim is enabled
+  __forceinline bool AsyncScratchReclaimEnabled() const override {
+    // TODO: Need to update min CP FW ucode version once it is released
+    return (core::Runtime::runtime_singleton_->flag().enable_scratch_async_reclaim() &&
+            isa()->GetMajorVersion() == 9 && isa()->GetMinorVersion() == 4 &&
+            properties_.EngineId.ui32.uCode > 999);
+  };
+
+  hsa_status_t SetAsyncScratchThresholds(size_t use_once_limit) override;
+
+  __forceinline size_t ScratchSingleLimitAsyncThreshold() const {
+    return scratch_limit_async_threshold_;
+  }
+
   void Trim() override;
 
   const std::function<void*(size_t size, size_t align, core::MemoryRegion::AllocateFlags flags)>&
@@ -534,6 +559,9 @@ class GpuAgent : public GpuAgentInt {
   // @brief Setup NUMA aware system memory allocator.
   void InitNumaAllocator();
 
+  // @brief Initialize scratch handler thresholds
+  void InitAsyncScratchThresholds();
+
   // @brief Register signal for notification when scratch may become available.
   // @p signal is notified by OR'ing with @p value.
   bool AddScratchNotifier(hsa_signal_t signal, hsa_signal_value_t value) {
@@ -579,6 +607,9 @@ class GpuAgent : public GpuAgentInt {
     KernelMutex lock_;
   } gws_queue_;
 
+  // @brief list of AQL queues owned by this agent. Indexed by queue pointer
+  std::vector<core::Queue*> aql_queues_;
+
   // Sets and Tracks pending SDMA status check or request counts
   void SetCopyRequestRefCount(bool set);
   void SetCopyStatusCheckRefCount(bool set);
@@ -588,6 +619,9 @@ class GpuAgent : public GpuAgentInt {
   // Tracks what SDMA blits have been used since initialization.
   uint32_t sdma_blit_used_mask_;
 
+  // Scratch limit thresholds when async scratch is enabled.
+  size_t scratch_limit_async_threshold_;
+
   ScratchCache scratch_cache_;
 
   // System memory allocator in the nearest NUMA node.

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

@@ -346,6 +346,9 @@ hsa_status_t hsa_amd_vmem_get_alloc_properties_from_handle(hsa_amd_vmem_alloc_ha
                                                            hsa_amd_memory_pool_t* pool,
                                                            hsa_amd_memory_type_t* type);
 
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_agent_set_async_scratch_limit(hsa_agent_t agent, size_t threshold);
+
 }  // namespace amd
 }  // namespace rocr
 

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

@@ -216,6 +216,11 @@ AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, Scr
   else
     queue_scratch_.mem_alignment_size = 1024;
 
+  queue_scratch_.use_once_limit = core::Runtime::runtime_singleton_->flag().scratch_single_limit();
+  queue_scratch_.async_reclaim = agent_->AsyncScratchReclaimEnabled();
+  if (queue_scratch_.async_reclaim)
+    queue_scratch_.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
+
   MAKE_NAMED_SCOPE_GUARD(EventGuard, [&]() {
     ScopedAcquire<KernelMutex> _lock(&queue_lock_);
     queue_count_--;
@@ -283,6 +288,8 @@ 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;
+
   // On the first queue creation, reserve some scratch memory on this agent.
   agent_->ReserveScratch();
 
@@ -778,6 +785,49 @@ hsa_status_t AqlQueue::SetPriority(HSA_QUEUE_PRIORITY priority) {
   return (err == HSAKMT_STATUS_SUCCESS ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR_OUT_OF_RESOURCES);
 }
 
+void AqlQueue::CheckScratchLimits() {
+  auto& scratch = queue_scratch_;
+  if (!scratch.async_reclaim) return;
+
+  scratch.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
+
+  if (scratch.main_size > scratch.use_once_limit) AsyncReclaimMainScratch();
+
+  return;
+}
+
+void AqlQueue::FreeMainScratchSpace() {
+  auto& scratch = queue_scratch_;
+  agent_->ReleaseQueueMainScratch(scratch);
+  scratch.main_size = 0;
+  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;
+
+  // Notify CP that we are trying to reclaim scratch. CP will assume scratch is reclaimed on next
+  // dispatch
+  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);
+
+  // Wait for scratch to be idle.
+  while (true) {
+    uint64_t last = amd_queue_.scratch_last_used_index;
+
+    if (std::min(last, last_used) < amd_queue_.read_dispatch_id) {
+      FreeMainScratchSpace();
+      return;
+    }
+  }
+}
+
 void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
                                          hsa_signal_value_t& waitVal, bool& changeWait) {
   // Insufficient scratch - recoverable, don't process dynamic scratch if errors are present.
@@ -790,6 +840,11 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
    * uint64_t all_slots_size;      // Size needed to fill all slots on this device
    * uint64_t dispatch_size;       // Size needed to fill wanted slots for this dispatch
    *
+   * //Default values:
+   * size_t use_once_limit = 128 MB      // When async reclaim not supported
+   *                       = 1GB per-XCC // When async reclaim is supported
+   *
+   *
    *******************************************************************************************/
 
   auto calc_dispatch_waves_per_group = [&](core::AqlPacket& pkt) {
@@ -845,7 +900,9 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
     return AlignUp(cu_count, engines) * agent_->properties().MaxSlotsScratchCU;
   };
 
-  scratch.use_once_limit = core::Runtime::runtime_singleton_->flag().scratch_single_limit();
+  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");
+
   scratch.cooperative = (amd_queue_.hsa_queue.type == HSA_QUEUE_TYPE_COOPERATIVE);
 
   uint64_t pkt_slot_idx = amd_queue_.read_dispatch_id & (amd_queue_.hsa_queue.size - 1);
@@ -855,7 +912,6 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
   pkt.AssertIsDispatchAndNeedsScratch();
 
   uint32_t device_slots = calc_device_slots();
-
   uint32_t groups = calc_dispatch_groups(pkt);
   uint32_t waves_per_group = calc_dispatch_waves_per_group(pkt);
 
@@ -864,12 +920,14 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
 
   const uint64_t lanes_per_wave = (error_code & 0x400) ? 32 : 64;
 
-  uint64_t device_size = pkt.dispatch.private_segment_size * lanes_per_wave * device_slots;
-  uint64_t dispatch_size = pkt.dispatch.private_segment_size * lanes_per_wave * dispatch_slots;
+  const uint64_t size_per_thread =
+      AlignUp(pkt.dispatch.private_segment_size, scratch.mem_alignment_size / lanes_per_wave);
+  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;
 
   agent_->ReleaseQueueMainScratch(scratch);
   scratch.main_size = device_size;
-  scratch.main_size_per_thread = pkt.dispatch.private_segment_size;
+  scratch.main_size_per_thread = size_per_thread;
   scratch.main_lanes_per_wave = lanes_per_wave;
   scratch.main_waves_per_group = waves_per_group;
 

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

@@ -83,6 +83,8 @@
 #define DEFAULT_SCRATCH_BYTES_PER_THREAD 2048
 #define MAX_WAVE_SCRATCH 8387584  // See COMPUTE_TMPRING_SIZE.WAVESIZE
 #define MAX_NUM_DOORBELLS 0x400
+#define MAX_SCRATCH_APERTURE_PER_XCC 4294967296
+#define DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC (1 << 30)  // 1 GB
 
 namespace rocr {
 namespace core {
@@ -109,6 +111,7 @@ GpuAgent::GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xna
       pending_copy_req_ref_(0),
       pending_copy_stat_check_ref_(0),
       sdma_blit_used_mask_(0),
+      scratch_limit_async_threshold_(0),
       scratch_cache_(
           [this](void* base, size_t size, bool large) { ReleaseScratch(base, size, large); }) {
   const bool is_apu_node = (properties_.NumCPUCores > 0);
@@ -200,6 +203,9 @@ GpuAgent::GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xna
 
   // Populate cache list.
   InitCacheList();
+
+  // Initialize thresholds for async-scratch handling
+  InitAsyncScratchThresholds();
 }
 
 GpuAgent::~GpuAgent() {
@@ -496,10 +502,10 @@ void GpuAgent::InitScratchPool() {
   size_t max_scratch_len = queue_scratch_len_ * max_queues_;
 
 #if defined(HSA_LARGE_MODEL) && defined(__linux__)
-  const size_t max_scratch_device = properties_.NumXcc * 4294967296;
+  const size_t max_scratch_device = properties_.NumXcc * MAX_SCRATCH_APERTURE_PER_XCC;
   // For 64-bit linux use max queues unless otherwise specified
   if ((max_scratch_len == 0) || (max_scratch_len > max_scratch_device)) {
-    max_scratch_len = max_scratch_device;  // 4GB per XCC apeture max
+    max_scratch_len = max_scratch_device;  // 4GB per XCC aperture max
   }
 #endif
 
@@ -518,6 +524,15 @@ void GpuAgent::InitScratchPool() {
   }
 }
 
+void GpuAgent::InitAsyncScratchThresholds() {
+  scratch_limit_async_threshold_ =
+      core::Runtime::runtime_singleton_->flag().scratch_single_limit_async();
+
+  if (!scratch_limit_async_threshold_)
+    scratch_limit_async_threshold_ =
+        DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC * properties().NumXcc;
+}
+
 void GpuAgent::ReserveScratch()
 {
   size_t reserved_sz = core::Runtime::runtime_singleton_->flag().scratch_single_limit();
@@ -1555,6 +1570,10 @@ hsa_status_t GpuAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
     return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
   }
 
+  // Asynchronous reclaim flag bit is set by CP FW on queue-connect, we will update this when
+  // we get the first scratch request.
+  scratch.async_reclaim = false;
+
   scratch.main_lanes_per_wave = 64;
   scratch.main_size_per_thread = AlignUp(private_segment_size, 1024 / scratch.main_lanes_per_wave);
   if (scratch.main_size_per_thread > 262128) {
@@ -1586,6 +1605,7 @@ hsa_status_t GpuAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
   auto aql_queue =
       new AqlQueue(this, size, node_id(), scratch, event_callback, data, is_kv_device_);
   *queue = aql_queue;
+  aql_queues_.push_back(aql_queue);
 
   if (doorbell_queue_map_) {
     // Calculate index of the queue doorbell within the doorbell aperture.
@@ -1805,6 +1825,28 @@ void GpuAgent::ReleaseScratch(void* base, size_t size, bool large) {
   ClearScratchNotifiers();
 }
 
+// Go through all the AQL queues and try to release scratch memory
+void GpuAgent::AsyncReclaimScratchQueues() {
+  for (auto iter : aql_queues_) {
+    auto aqlQueue = static_cast<AqlQueue*>(iter);
+    aqlQueue->AsyncReclaimMainScratch();
+  }
+}
+
+hsa_status_t GpuAgent::SetAsyncScratchThresholds(size_t use_once_limit) {
+  if (use_once_limit > properties_.NumXcc * MAX_SCRATCH_APERTURE_PER_XCC)
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+
+  scratch_limit_async_threshold_ = use_once_limit;
+
+  for (auto iter : aql_queues_) {
+    auto aqlQueue = static_cast<AqlQueue*>(iter);
+    aqlQueue->CheckScratchLimits();
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
 void GpuAgent::TranslateTime(core::Signal* signal, hsa_amd_profiling_dispatch_time_t& time) {
   uint64_t start, end;
   signal->GetRawTs(false, start, end);
@@ -2126,6 +2168,7 @@ lazy_ptr<core::Blit>& GpuAgent::GetBlitObject(const core::Agent& dst_agent,
 
 void GpuAgent::Trim() {
   Agent::Trim();
+  AsyncReclaimScratchQueues();
   ScopedAcquire<KernelMutex> lock(&scratch_lock_);
   scratch_cache_.trim(false);
 }

runtime/hsa-runtime/core/runtime/hsa_api_trace.cpp

@@ -80,7 +80,7 @@ void HsaApiTable::Init() {
   // they can add preprocessor macros on the new functions
 
   constexpr size_t expected_core_api_table_size = 1016;
-  constexpr size_t expected_amd_ext_table_size = 552;
+  constexpr size_t expected_amd_ext_table_size = 560;
   constexpr size_t expected_image_ext_table_size = 120;
   constexpr size_t expected_finalizer_ext_table_size = 64;
 
@@ -437,6 +437,7 @@ void HsaApiTable::UpdateAmdExts() {
   amd_ext_api.hsa_amd_vmem_retain_alloc_handle_fn = AMD::hsa_amd_vmem_retain_alloc_handle;
   amd_ext_api.hsa_amd_vmem_get_alloc_properties_from_handle_fn =
       AMD::hsa_amd_vmem_get_alloc_properties_from_handle;
+  amd_ext_api.hsa_amd_agent_set_async_scratch_limit_fn = AMD::hsa_amd_agent_set_async_scratch_limit;
 }
 
 void LoadInitialHsaApiTable() {

runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp

@@ -1367,5 +1367,23 @@ hsa_status_t hsa_amd_vmem_get_alloc_properties_from_handle(hsa_amd_vmem_alloc_ha
   CATCH;
 }
 
+hsa_status_t HSA_API hsa_amd_agent_set_async_scratch_limit(hsa_agent_t _agent, size_t threshold) {
+  TRY;
+  IS_OPEN();
+
+  core::Agent* agent = core::Agent::Convert(_agent);
+  if (agent == NULL || !agent->IsValid() || agent->device_type() != core::Agent::kAmdGpuDevice)
+    return HSA_STATUS_ERROR_INVALID_AGENT;
+
+  AMD::GpuAgentInt* gpu_agent = static_cast<AMD::GpuAgentInt*>(agent);
+
+  if (!core::Runtime::runtime_singleton_->flag().enable_scratch_async_reclaim() ||
+      !gpu_agent->AsyncScratchReclaimEnabled())
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+
+  return gpu_agent->SetAsyncScratchThresholds(threshold);
+  CATCH;
+}
+
 }   //  namespace amd
 }   //  namespace rocr

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

@@ -64,7 +64,8 @@ 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.
+  // 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;
 
   explicit Flag() { Refresh(); }
@@ -120,11 +121,26 @@ class Flag {
     // memory
     if (os::IsEnvVarSet("HSA_SCRATCH_SINGLE_LIMIT")) {
       var = os::GetEnvVar("HSA_SCRATCH_SINGLE_LIMIT");
-      scratch_single_limit_ = atoi(var.c_str());
+      char* end;
+      scratch_single_limit_ = strtoul(var.c_str(), &end, 10);
     } else {
       scratch_single_limit_ = DEFAULT_SCRATCH_SINGLE_LIMIT;
     }
 
+    // On GPUs that support asynchronous scratch reclaim
+    // Scratch memory sizes > HSA_SCRATCH_SINGLE_LIMIT_ASYNC will trigger a use-once scheme
+    if (os::IsEnvVarSet("HSA_SCRATCH_SINGLE_LIMIT_ASYNC")) {
+      var = os::GetEnvVar("HSA_SCRATCH_SINGLE_LIMIT_ASYNC");
+      char* end;
+      scratch_single_limit_async_ = strtoul(var.c_str(), &end, 10);
+    } else {
+      scratch_single_limit_async_ = 0;  // DEFAULT_SCRATCH_SINGLE_LIMIT_ASYNC_PER_XCC;
+    }
+
+    // On GPUs that support asynchronous scratch reclaim this can be used to disable this feature.
+    var = os::GetEnvVar("HSA_ENABLE_SCRATCH_ASYNC_RECLAIM");
+    enable_scratch_async_reclaim_ = (var == "0") ? false : true;
+
     tools_lib_names_ = os::GetEnvVar("HSA_TOOLS_LIB");
 
     var = os::GetEnvVar("HSA_TOOLS_REPORT_LOAD_FAILURE");
@@ -252,6 +268,10 @@ class Flag {
 
   size_t scratch_single_limit() const { return scratch_single_limit_; }
 
+  bool enable_scratch_async_reclaim() const { return enable_scratch_async_reclaim_; }
+
+  size_t scratch_single_limit_async() const { return scratch_single_limit_async_; }
+
   std::string tools_lib_names() const { return tools_lib_names_; }
 
   bool disable_image() const { return disable_image_; }
@@ -330,6 +350,8 @@ class Flag {
 
   size_t scratch_mem_size_;
   size_t scratch_single_limit_;
+  size_t scratch_single_limit_async_;
+  bool enable_scratch_async_reclaim_;
 
   std::string tools_lib_names_;
   std::string svm_profile_;

runtime/hsa-runtime/hsacore.so.def

@@ -245,6 +245,7 @@ global:
 	hsa_amd_vmem_import_shareable_handle;
 	hsa_amd_vmem_retain_alloc_handle;
 	hsa_amd_vmem_get_alloc_properties_from_handle;
+	hsa_amd_agent_set_async_scratch_limit;
 
 local:
     *;

runtime/hsa-runtime/inc/hsa_api_trace.h

@@ -69,7 +69,7 @@
 // Step Ids of the Api tables exported by Hsa Core Runtime
 #define HSA_API_TABLE_STEP_VERSION                0x00
 #define HSA_CORE_API_TABLE_STEP_VERSION           0x00
-#define HSA_AMD_EXT_API_TABLE_STEP_VERSION        0x00
+#define HSA_AMD_EXT_API_TABLE_STEP_VERSION        0x01
 #define HSA_FINALIZER_API_TABLE_STEP_VERSION      0x00
 #define HSA_IMAGE_API_TABLE_STEP_VERSION          0x00
 #define HSA_AQLPROFILE_API_TABLE_STEP_VERSION     0x00
@@ -247,6 +247,7 @@ struct AmdExtTable {
   decltype(hsa_amd_vmem_retain_alloc_handle)* hsa_amd_vmem_retain_alloc_handle_fn;
   decltype(hsa_amd_vmem_get_alloc_properties_from_handle)*
       hsa_amd_vmem_get_alloc_properties_from_handle_fn;
+  decltype(hsa_amd_agent_set_async_scratch_limit)* hsa_amd_agent_set_async_scratch_limit_fn;
 };
 
 // Table to export HSA Core Runtime Apis

runtime/hsa-runtime/inc/hsa_ext_amd.h

@@ -3026,6 +3026,29 @@ hsa_status_t hsa_amd_vmem_retain_alloc_handle(hsa_amd_vmem_alloc_handle_t* memor
 hsa_status_t hsa_amd_vmem_get_alloc_properties_from_handle(
     hsa_amd_vmem_alloc_handle_t memory_handle, hsa_amd_memory_pool_t* pool,
     hsa_amd_memory_type_t* type);
+
+/**
+ * @brief Set the asynchronous scratch limit threshold on all the queues for this agent.
+ * Dispatches that are enqueued on HW queues on this agent that are smaller than threshold will not
+ * result in a scratch use-once method. This API is only supported on devices that support
+ * asynchronous scratch reclaim.
+ *
+ * @param[in] agent A valid agent.
+ *
+ * @param[in] threshold Threshold size in bytes
+ *
+ * @retval ::HSA_STATUS_SUCCESS The function has been executed successfully.
+ *
+ * @retval ::HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
+ * initialized.
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_AGENT The agent is invalid.
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT This agent does not support asynchronous scratch
+ * reclaim
+ */
+hsa_status_t HSA_API hsa_amd_agent_set_async_scratch_limit(hsa_agent_t agent, size_t threshold);
+
 #ifdef __cplusplus
 }  // end extern "C" block
 #endif