From 469defa78a746d37823e6fbfaee3bbb4ebdd87ab Mon Sep 17 00:00:00 2001
From: David Yat Sin <David.YatSin@amd.com>
Date: Thu, 25 May 2023 22:47:25 +0000
Subject: [PATCH] Add agent query for nearest CPU agent

Add agent info query to return nearest CPU agent. This can be used to
determine which CPU agent is in the same NUMA region as the GPU agent.

Change-Id: I5400b4347ffbf4d2a836df31c4de443a38b0ecd1
---
 runtime/hsa-runtime/core/inc/amd_gpu_agent.h  |  2 ++
 .../core/runtime/amd_cpu_agent.cpp            |  3 ++
 .../core/runtime/amd_gpu_agent.cpp            | 30 +++++++++++--------
 runtime/hsa-runtime/inc/hsa_ext_amd.h         |  7 ++++-
 4 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/runtime/hsa-runtime/core/inc/amd_gpu_agent.h b/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
index ab48ca989c..d2a502749a 100644
--- a/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
+++ b/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
@@ -289,6 +289,8 @@ class GpuAgent : public GpuAgentInt {
     return current_coherency_type_;
   }
 
+  core::Agent* GetNearestCpuAgent(void) const;
+
   // Getter & setters.
 
   // @brief Returns Hive ID
diff --git a/runtime/hsa-runtime/core/runtime/amd_cpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_cpu_agent.cpp
index bdb6070f80..6367588722 100644
--- a/runtime/hsa-runtime/core/runtime/amd_cpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_cpu_agent.cpp
@@ -391,6 +391,9 @@ hsa_status_t CpuAgent::GetInfo(hsa_agent_info_t attribute, void* value) const {
     case HSA_AMD_AGENT_INFO_DRIVER_UID:
       *((uint32_t*)value) = 0;
       break;
+    case HSA_AMD_AGENT_INFO_NEAREST_CPU:
+      ((hsa_agent_t*)value)->handle = 0;
+      break;
     default:
       return HSA_STATUS_ERROR_INVALID_ARGUMENT;
       break;
diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index e2025c6daa..5fc8f07856 100644
--- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -1367,6 +1367,9 @@ hsa_status_t GpuAgent::GetInfo(hsa_agent_info_t attribute, void* value) const {
     case HSA_AMD_AGENT_INFO_DRIVER_UID:
       *((uint32_t*)value) = KfdGpuID();
       break;
+    case HSA_AMD_AGENT_INFO_NEAREST_CPU:
+      *((hsa_agent_t*)value) = GetNearestCpuAgent()->public_handle();
+      break;
     default:
       return HSA_STATUS_ERROR_INVALID_ARGUMENT;
       break;
@@ -1985,18 +1988,7 @@ void GpuAgent::Trim() {
 }
 
 void GpuAgent::InitNumaAllocator() {
-  Agent* nearCpu = nullptr;
-  uint32_t dist = -1u;
-  for (auto cpu : core::Runtime::runtime_singleton_->cpu_agents()) {
-    const core::Runtime::LinkInfo link_info =
-        core::Runtime::runtime_singleton_->GetLinkInfo(node_id(), cpu->node_id());
-    if (link_info.info.numa_distance < dist) {
-      dist = link_info.info.numa_distance;
-      nearCpu = cpu;
-    }
-  }
-
-  for (auto pool : nearCpu->regions()) {
+  for (auto pool : GetNearestCpuAgent()->regions()) {
     if (pool->kernarg()) {
       system_allocator_ = [pool](size_t size, size_t alignment,
                                  MemoryRegion::AllocateFlags alloc_flags) -> void* {
@@ -2016,5 +2008,19 @@ void GpuAgent::InitNumaAllocator() {
   assert(false && "Nearest NUMA node did not have a kernarg pool.");
 }
 
+core::Agent* GpuAgent::GetNearestCpuAgent() const {
+  core::Agent* nearCpu = nullptr;
+  uint32_t dist = -1u;
+  for (auto cpu : core::Runtime::runtime_singleton_->cpu_agents()) {
+    const core::Runtime::LinkInfo link_info =
+        core::Runtime::runtime_singleton_->GetLinkInfo(node_id(), cpu->node_id());
+    if (link_info.info.numa_distance < dist) {
+      dist = link_info.info.numa_distance;
+      nearCpu = cpu;
+    }
+  }
+  return nearCpu;
+}
+
 }  // namespace amd
 }  // namespace rocr
diff --git a/runtime/hsa-runtime/inc/hsa_ext_amd.h b/runtime/hsa-runtime/inc/hsa_ext_amd.h
index e493ac7ff8..ae1c90d1dc 100644
--- a/runtime/hsa-runtime/inc/hsa_ext_amd.h
+++ b/runtime/hsa-runtime/inc/hsa_ext_amd.h
@@ -399,7 +399,12 @@ typedef enum hsa_amd_agent_info_s {
    * Queries for driver unique identifier.
    * The type of this attribute is uint32_t.
    */
-  HSA_AMD_AGENT_INFO_DRIVER_UID = 0xA112
+  HSA_AMD_AGENT_INFO_DRIVER_UID = 0xA112,
+  /**
+   * Returns the hsa_agent_t of the nearest CPU agent
+   * The type of this attribute is hsa_agent_t.
+   */
+  HSA_AMD_AGENT_INFO_NEAREST_CPU = 0xA113
 } hsa_amd_agent_info_t;
 
 /**