From 4ac840269cdfff771f076d75bb178f3bd310e1e3 Mon Sep 17 00:00:00 2001
From: David Yat Sin <david.yatsin@amd.com>
Date: Tue, 26 Apr 2022 21:11:37 -0400
Subject: [PATCH] Add API for available GPU memory

Add support for AMD Agent to return amount of memory available

Change-Id: I5c32e2cebbaa2993b044250aefe434e4cc02d8c2
Signed-off-by: David Yat Sin <david.yatsin@amd.com>
---
 .../hsa-runtime/core/inc/amd_memory_region.h    |  2 ++
 .../hsa-runtime/core/runtime/amd_gpu_agent.cpp  | 17 +++++++++++++++++
 runtime/hsa-runtime/core/util/simple_heap.h     |  2 ++
 runtime/hsa-runtime/inc/hsa_ext_amd.h           |  8 +++++++-
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/runtime/hsa-runtime/core/inc/amd_memory_region.h b/runtime/hsa-runtime/core/inc/amd_memory_region.h
index 185889040f..e12f0d760d 100644
--- a/runtime/hsa-runtime/core/inc/amd_memory_region.h
+++ b/runtime/hsa-runtime/core/inc/amd_memory_region.h
@@ -138,6 +138,8 @@ class MemoryRegion : public core::MemoryRegion {
 
   void Trim() const;
 
+  HSAuint64 GetCacheSize() const { return fragment_allocator_.cache_size(); }
+
   __forceinline bool IsLocalMemory() const {
     return ((mem_props_.HeapType == HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE) ||
             (mem_props_.HeapType == HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC));
diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index 2bfdbbb13a..138cd008b6 100644
--- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -1032,6 +1032,23 @@ hsa_status_t GpuAgent::GetInfo(hsa_agent_info_t attribute, void* value) const {
         break;
       }
       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;
+
+      status = hsaKmtAvailableMemory(node_id(), &availableBytes);
+
+      if (status != HSAKMT_STATUS_SUCCESS)
+        return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+
+      for (auto r: regions())
+        availableBytes += ((AMD::MemoryRegion*)r)->GetCacheSize();
+
+      availableBytes += scratch_cache_.free_bytes();
+
+      *((uint64_t*)value) = availableBytes;
+      break;
+    }
     default:
       return HSA_STATUS_ERROR_INVALID_ARGUMENT;
       break;
diff --git a/runtime/hsa-runtime/core/util/simple_heap.h b/runtime/hsa-runtime/core/util/simple_heap.h
index ef4c8cb477..f284e2650c 100644
--- a/runtime/hsa-runtime/core/util/simple_heap.h
+++ b/runtime/hsa-runtime/core/util/simple_heap.h
@@ -282,6 +282,8 @@ template <typename Allocator> class SimpleHeap {
     cache_size_ = 0;
   }
 
+  size_t cache_size() const { return cache_size_; }
+
   size_t default_block_size() const { return block_allocator_.block_size(); }
 
   // Prevent reuse of the block containing ptr.  No further fragments will be allocated from the
diff --git a/runtime/hsa-runtime/inc/hsa_ext_amd.h b/runtime/hsa-runtime/inc/hsa_ext_amd.h
index 9adca9d788..1f9c264cf8 100644
--- a/runtime/hsa-runtime/inc/hsa_ext_amd.h
+++ b/runtime/hsa-runtime/inc/hsa_ext_amd.h
@@ -318,7 +318,13 @@ typedef enum hsa_amd_agent_info_s {
    * cooperative dispatch.
    * The type of this attribute is uint32_t.
    */
-  HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT = 0xA014
+  HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT = 0xA014,
+  /**
+   * Queries the amount of memory available in bytes accross all global pools
+   * owned by the agent.
+   * The type of this attribute is uint64_t.
+   */
+  HSA_AMD_AGENT_INFO_MEMORY_AVAIL = 0xA015
 } hsa_amd_agent_info_t;
 
 typedef struct hsa_amd_hdp_flush_s {