From 6ae29610a531cea66e94f8bdcf96f2c5cbdb3bf9 Mon Sep 17 00:00:00 2001
From: Ruihang Lai <ruihangl@cs.cmu.edu>
Date: Fri, 9 Aug 2024 09:44:59 -0400
Subject: [PATCH] [ROCm] Support ROCm 6 (#17256)

This PR updates some ROCm modules in order to support ROCm 6.
---
 cmake/modules/ROCM.cmake            | 1 +
 cmake/utils/FindRCCL.cmake          | 2 +-
 src/runtime/rocm/rocm_device_api.cc | 2 +-
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/cmake/modules/ROCM.cmake b/cmake/modules/ROCM.cmake
index 37fcd716464e..02c4c739934a 100644
--- a/cmake/modules/ROCM.cmake
+++ b/cmake/modules/ROCM.cmake
@@ -23,6 +23,7 @@ if(ROCM_FOUND)
   # avoid global retrigger of cmake
   include_directories(SYSTEM ${ROCM_INCLUDE_DIRS})
   add_definitions(-D__HIP_PLATFORM_HCC__=1)
+  add_definitions(-D__HIP_PLATFORM_AMD__=1)
 endif(ROCM_FOUND)
 
 
diff --git a/cmake/utils/FindRCCL.cmake b/cmake/utils/FindRCCL.cmake
index 93d8c8744630..95cb555178d0 100644
--- a/cmake/utils/FindRCCL.cmake
+++ b/cmake/utils/FindRCCL.cmake
@@ -32,7 +32,7 @@ macro(find_rccl use_rccl)
     find_path(RCCL_INCLUDE_DIR NAMES rccl.h)
     find_library(RCCL_LIBRARY NAMES rccl)
   else()
-    find_path(RCCL_INCLUDE_DIR NAMES rccl.h HINTS ${use_rccl} ${use_rccl}/include)
+    find_path(RCCL_INCLUDE_DIR NAMES rccl.h HINTS ${use_rccl} ${use_rccl}/include ${use_rccl}/include/rccl)
     find_library(RCCL_LIBRARY NAMES rccl HINTS ${use_rccl} ${use_rccl}/lib)
   endif()
   include(FindPackageHandleStandardArgs)
diff --git a/src/runtime/rocm/rocm_device_api.cc b/src/runtime/rocm/rocm_device_api.cc
index e2a5048ca030..c37e9fada5b2 100644
--- a/src/runtime/rocm/rocm_device_api.cc
+++ b/src/runtime/rocm/rocm_device_api.cc
@@ -113,7 +113,7 @@ class ROCMDeviceAPI final : public DeviceAPI {
       case kGcnArch: {
         hipDeviceProp_t prop;
         ROCM_CALL(hipGetDeviceProperties(&prop, device.device_id));
-        *rv = prop.gcnArch;
+        *rv = prop.gcnArchName;
         return;
       }
       case kApiVersion: {