rocr/driver: add AvailableMemory API to driver

This commit introduces a new AvailableMemory API to the KfdDriver and
 XdnaDriver classes.

- Implemented AvailableMemory in KfdDriver to return the available memory size
  using hsaKmtAvailableMemory.
- Added a stub implementation of AvailableMemory in XdnaDriver that returns an error.
- Updated the GpuAgent class to use the new AvailableMemory API instead of
  directly calling hsaKmtAvailableMemory.

Signed-off-by: Honglei Huang <Honglei1.Huang@amd.com>

runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp

@@ -616,6 +616,14 @@ hsa_status_t KfdDriver::GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t KfdDriver::AvailableMemory(uint32_t node_id, uint64_t* available_size) const {
+  assert(available_size);
+
+  if (HSAKMT_CALL(hsaKmtAvailableMemory(node_id, available_size)) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t KfdDriver::IsModelEnabled(bool* enable) const {
   // AIE does not support streaming performance monitor.
   HSAKMT_STATUS status = HSAKMT_STATUS_ERROR;

runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -898,5 +898,9 @@ hsa_status_t XdnaDriver::GetWallclockFrequency(uint32_t node_id, uint64_t* frequ
   return HSA_STATUS_ERROR;
 }
 
+hsa_status_t XdnaDriver::AvailableMemory(uint32_t node_id, uint64_t* available_size) const {
+  return HSA_STATUS_ERROR;
+}
+
 } // namespace AMD
 } // namespace rocr

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

@@ -129,6 +129,7 @@ public:
   hsa_status_t GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const override;
   hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const override;
   hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const override;
+  hsa_status_t AvailableMemory(uint32_t node_id, uint64_t* available_size) const override;
 
   hsa_status_t OpenSMI(uint32_t node_id, int* fd) const override;
 

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

@@ -243,6 +243,7 @@ public:
   hsa_status_t GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const override;
   hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const override;
   hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const override;
+  hsa_status_t AvailableMemory(uint32_t node_id, uint64_t* available_size) const override;
 
   hsa_status_t IsModelEnabled(bool* enable) const override;
 

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

@@ -308,6 +308,12 @@ public:
   /// @return HSA_STATUS_SUCCESS if scratch memory allocated successfully.
   virtual hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const = 0;
 
+  /// @brief Inquires memory available for allocation as a memory buffer
+  /// @param[in] node_id Node ID of the agent
+  /// @param[out] available_size Available memory size in bytes
+  /// @return HSA_STATUS_SUCCESS if the driver successfully returns the available memory size.
+  virtual hsa_status_t AvailableMemory(uint32_t node_id, uint64_t* available_size) const = 0;
+
   /// Unique identifier for supported kernel-mode drivers.
   const DriverType kernel_driver_type_;
 

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

@@ -548,8 +548,8 @@ void GpuAgent::ReserveScratch()
   }
 
   size_t available;
-  HSAKMT_STATUS err = HSAKMT_CALL(hsaKmtAvailableMemory(node_id(), &available));
-  assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtAvailableMemory failed");
+  hsa_status_t err = driver().AvailableMemory(node_id(), &available);
+  assert(err == HSA_STATUS_SUCCESS && "AvailableMemory failed");
   ScopedAcquire<KernelMutex> lock(&scratch_lock_);
   if (!scratch_cache_.reserved_bytes() && reserved_sz && available > 8 * reserved_sz) {
     HSAuint64 alt_va;
@@ -1606,11 +1606,11 @@ hsa_status_t GpuAgent::GetInfo(hsa_agent_info_t attribute, void* value) const {
       return GetInfo((hsa_agent_info_t)HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, value);
     case HSA_AMD_AGENT_INFO_MEMORY_AVAIL: {
       HSAuint64 availableBytes;
-      HSAKMT_STATUS status;
+      hsa_status_t status;
 
-      status = HSAKMT_CALL(hsaKmtAvailableMemory(node_id(), &availableBytes));
+      status = driver().AvailableMemory(node_id(), &availableBytes);
 
-      if (status != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+      if (status != HSA_STATUS_SUCCESS) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
 
       for (auto r : regions()) availableBytes += ((AMD::MemoryRegion*)r)->GetCacheSize();