Merge pull request #10729 from edgargabriel/topic/gfx-arch-cleanup

UCT/ROCM: cleanup GPU detection code

Author: edgargabriel

src/uct/rocm/base/rocm_base.c

@@ -511,46 +511,28 @@ ucs_status_t uct_rocm_base_get_link_type(hsa_amd_link_info_type_t *link_type)
 uct_rocm_amd_gpu_product_t uct_rocm_base_get_gpu_product(void)
 {
     uct_rocm_amd_gpu_product_t gpu_product = UCT_ROCM_AMD_GPU_MI200;
-    char product_name[64];
     char gfx_name[64];
     hsa_status_t status;
 
-    /* fetching data from GPU 0, assuming all GPUs on a node are
-       identical */
+    /* Query the gfx architecture name */
     status = hsa_agent_get_info(uct_rocm_base_agents.gpu_agents[0],
-                                (hsa_agent_info_t)
-                                        HSA_AMD_AGENT_INFO_PRODUCT_NAME,
-                                (void*)product_name);
+                                (hsa_agent_info_t)HSA_AGENT_INFO_NAME,
+                                gfx_name);
     if (status != HSA_STATUS_SUCCESS) {
-        ucs_debug("Error in hsa_agent_info %d", status);
+        ucs_debug("hsa_agent_get_info failed: %d", status);
         return gpu_product;
     }
 
-    if (NULL != strstr(product_name, "MI300A")) {
-        gpu_product = UCT_ROCM_AMD_GPU_MI300A;
-        ucs_trace("found MI300A GPU");
-    } else if (NULL != strstr(product_name, "MI300X")) {
-        gpu_product = UCT_ROCM_AMD_GPU_MI300X;
-        ucs_trace("found MI300X GPU");
+    if (NULL != strstr(gfx_name, "gfx94")) {
+        /* This is an MI300/325 GPU */
+        gpu_product = UCT_ROCM_AMD_GPU_MI300;
+        ucs_trace("found gfx94* GPU");
+    } else if (NULL != strstr(gfx_name, "gfx95")) {
+        /* This is an MI35x GPU */
+        gpu_product = UCT_ROCM_AMD_GPU_MI350;
+        ucs_trace("found gfx950 GPU");
     } else {
-        /* In case product_name is not set correctly, query the gfx
-           architecture name */
-        status = hsa_agent_get_info(uct_rocm_base_agents.gpu_agents[0],
-                                    (hsa_agent_info_t)HSA_AGENT_INFO_NAME,
-                                    (void*)gfx_name);
-        if (status != HSA_STATUS_SUCCESS) {
-            ucs_debug("error in hsa_agent_info %d", status);
-            return gpu_product;
-        }
-
-        if (NULL != strstr(gfx_name, "gfx94")) {
-            /* This is an MI300 GPU, but cannot say whether its the A or X
-               variant. Assuming A variant for now*/
-            gpu_product = UCT_ROCM_AMD_GPU_MI300A;
-            ucs_trace("found gfx94* GPU, assuming MI300A");
-        } else {
-            ucs_trace("assuming MI100/MI200 GPU");
-        }
+        ucs_trace("assuming MI100/MI200 GPU");
     }
 
     return gpu_product;

src/uct/rocm/base/rocm_base.h

@@ -16,8 +16,8 @@
 typedef enum uct_rocm_amd_gpu_product {
     UCT_ROCM_AMD_GPU_UNDEFINED = -1,
     UCT_ROCM_AMD_GPU_MI200,
-    UCT_ROCM_AMD_GPU_MI300A,
-    UCT_ROCM_AMD_GPU_MI300X
+    UCT_ROCM_AMD_GPU_MI300,
+    UCT_ROCM_AMD_GPU_MI350,
 } uct_rocm_amd_gpu_product_t;
 
 hsa_status_t uct_rocm_base_init(void);

src/uct/rocm/copy/rocm_copy_iface.c

@@ -212,7 +212,7 @@ uct_rocm_copy_get_default_bandwidth(const uct_perf_attr_t *perf_attr)
 }
 
 static UCS_F_ALWAYS_INLINE double
-uct_rocm_copy_get_mi300a_bandwidth(const uct_perf_attr_t *perf_attr)
+uct_rocm_copy_get_mi300_bandwidth(const uct_perf_attr_t *perf_attr)
 {
     switch (perf_attr->operation) {
     case UCT_EP_OP_GET_SHORT:
@@ -250,9 +250,8 @@ uct_rocm_copy_estimate_perf(uct_iface_h tl_iface, uct_perf_attr_t *perf_attr)
         if (!(perf_attr->field_mask & UCT_PERF_ATTR_FIELD_OPERATION)) {
             bandwidth.shared = 0;
         } else {
-            if (gpu_product == UCT_ROCM_AMD_GPU_MI300A ||
-                gpu_product == UCT_ROCM_AMD_GPU_MI300X) {
-                bandwidth.shared = uct_rocm_copy_get_mi300a_bandwidth(perf_attr);
+            if (gpu_product == UCT_ROCM_AMD_GPU_MI300) {
+                bandwidth.shared = uct_rocm_copy_get_mi300_bandwidth(perf_attr);
             } else {
                 bandwidth.shared = uct_rocm_copy_get_default_bandwidth(perf_attr);
             }