From f7eaff474b29b369618e014683d2d4532a67faa4 Mon Sep 17 00:00:00 2001
From: Maneesh Gupta <maneesh.gupta@amd.com>
Date: Mon, 6 May 2019 15:54:31 +0530
Subject: [PATCH] Implement hipExtGetLinkTypeAndHopCount for ROCm devices

Change-Id: Ie5bb4f640ac6d189c7fceeab22627a7494fd10bd


[ROCm/hip commit: 035ef04e1f44db28af80b9b079ab0607774f0a84]
---
 .../include/hip/hcc_detail/hip_runtime_api.h  | 14 ++++++
 projects/hip/src/hip_device.cpp               | 49 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/projects/hip/include/hip/hcc_detail/hip_runtime_api.h b/projects/hip/include/hip/hcc_detail/hip_runtime_api.h
index e8196af702..b6eb21e0ef 100644
--- a/projects/hip/include/hip/hcc_detail/hip_runtime_api.h
+++ b/projects/hip/include/hip/hcc_detail/hip_runtime_api.h
@@ -542,6 +542,20 @@ hipError_t hipSetDeviceFlags(unsigned flags);
  */
 hipError_t hipChooseDevice(int* device, const hipDeviceProp_t* prop);
 
+/**
+ * @brief Returns the link type and hop count between two devices
+ *
+ * @param [in] device1 Ordinal for device1
+ * @param [in] device2 Ordinal for device2
+ * @param [out] linktype Returns the link type (See hsa_amd_link_info_type_t) between the two devices
+ * @param [out] hopcount Returns the hop count between the two devices
+ *
+ * Queries and returns the HSA link type and the hop count between the two specified devices.
+ *
+ * @returns #hipSuccess, #hipInvalidDevice, #hipErrorRuntimeOther
+ */
+hipError_t hipExtGetLinkTypeAndHopCount(int device1, int device2, uint32_t* linktype, uint32_t* hopcount);
+
 // end doxygen Device
 /**
  * @}
diff --git a/projects/hip/src/hip_device.cpp b/projects/hip/src/hip_device.cpp
index 246ce8cf6f..6570b73316 100644
--- a/projects/hip/src/hip_device.cpp
+++ b/projects/hip/src/hip_device.cpp
@@ -531,3 +531,52 @@ hipError_t hipChooseDevice(int* device, const hipDeviceProp_t* prop) {
     }
     return ihipLogStatus(e);
 }
+
+#define HSA_ERROR_CHECK(hsa_error_code) \
+    if ((hsa_error_code != HSA_STATUS_SUCCESS) && (hsa_error_code != HSA_STATUS_INFO_BREAK)) { \
+        return ihipLogStatus(hipErrorRuntimeOther); \
+    }
+
+hipError_t hipExtGetLinkTypeAndHopCount(int device1, int device2, uint32_t* linktype, uint32_t* hopcount) {
+    HIP_INIT_API(hipExtGetLinkTypeAndHopCount, device1, device2, linktype, hopcount);
+
+    if ((device1 < 0) || (device1 >= g_deviceCnt) || (device2 < 0) || (device2 >= g_deviceCnt)) {
+        return ihipLogStatus(hipErrorInvalidDevice);
+    } else {
+        auto device1Handle = ihipGetDevice(device1);
+        auto device2Handle = ihipGetDevice(device2);
+
+        const auto& find_pool = [](hsa_amd_memory_pool_t pool, void* data) {
+            bool allowed;
+            hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED, &allowed);
+            if (allowed) {
+                hsa_amd_segment_t segment;
+                hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_SEGMENT, &segment);
+                if (HSA_AMD_SEGMENT_GLOBAL != segment) return HSA_STATUS_SUCCESS;
+
+                uint32_t flags;
+                hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &flags);
+                if (flags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED) {
+                    *((hsa_amd_memory_pool_t*)data) = pool;
+                    return HSA_STATUS_INFO_BREAK;
+                }
+            }
+            return HSA_STATUS_SUCCESS;
+        };
+
+        hsa_status_t err;
+        hsa_amd_memory_pool_t pool;
+        err = hsa_amd_agent_iterate_memory_pools(device2Handle->_hsaAgent, find_pool, (void*)&pool);
+        HSA_ERROR_CHECK(err);
+
+        hsa_amd_memory_pool_link_info_t link_info;
+        err = hsa_amd_agent_memory_pool_get_info(device1Handle->_hsaAgent, pool, HSA_AMD_AGENT_MEMORY_POOL_INFO_LINK_INFO, &link_info);
+        HSA_ERROR_CHECK(err);
+        *linktype = link_info.link_type;
+
+        err = hsa_amd_agent_memory_pool_get_info(device1Handle->_hsaAgent, pool, HSA_AMD_AGENT_MEMORY_POOL_INFO_NUM_LINK_HOPS, hopcount);
+        HSA_ERROR_CHECK(err);
+
+        return ihipLogStatus(hipSuccess);
+    }
+}