From 6c390e32ccbc11d8e22002014b8e3bd1253fcadb Mon Sep 17 00:00:00 2001 From: Honglei Huang Date: Tue, 15 Jul 2025 10:49:11 +0800 Subject: [PATCH] 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 --- runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp | 8 ++++++++ .../hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp | 4 ++++ runtime/hsa-runtime/core/inc/amd_kfd_driver.h | 1 + runtime/hsa-runtime/core/inc/amd_xdna_driver.h | 1 + runtime/hsa-runtime/core/inc/driver.h | 6 ++++++ runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp | 10 +++++----- 6 files changed, 25 insertions(+), 5 deletions(-) diff --git a/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp b/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp index b9751f68b3..baf30aaac9 100644 --- a/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp +++ b/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; diff --git a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp index 76c4842d5f..f4555e4005 100644 --- a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp +++ b/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 diff --git a/runtime/hsa-runtime/core/inc/amd_kfd_driver.h b/runtime/hsa-runtime/core/inc/amd_kfd_driver.h index fc0ee82cba..e9feb3a39c 100644 --- a/runtime/hsa-runtime/core/inc/amd_kfd_driver.h +++ b/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; diff --git a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h b/runtime/hsa-runtime/core/inc/amd_xdna_driver.h index 0aee5bdfcb..7ff5cf8e6a 100644 --- a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h +++ b/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; diff --git a/runtime/hsa-runtime/core/inc/driver.h b/runtime/hsa-runtime/core/inc/driver.h index 8df7e73686..f03ec94a0d 100644 --- a/runtime/hsa-runtime/core/inc/driver.h +++ b/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_; diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp index 9a0c621672..3be9215299 100644 --- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp +++ b/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 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();