From c42c7d99f159a4f4f07c0e0b87417e056b302f01 Mon Sep 17 00:00:00 2001
From: Hui Liu <96135754+hliuca@users.noreply.github.com>
Date: Thu, 11 Jan 2024 14:10:30 -0800
Subject: [PATCH] fix memoryType

---
 cmake/Utils.cmake           |  4 ++--
 rocgputreeshap              |  2 +-
 src/data/array_interface.cu | 32 ++++++++++++++++++++++++++++++--
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/cmake/Utils.cmake b/cmake/Utils.cmake
index 19ccdac8a..f295d1446 100644
--- a/cmake/Utils.cmake
+++ b/cmake/Utils.cmake
@@ -214,12 +214,12 @@ function(xgboost_link_rccl target)
   endif()
 
   if(BUILD_STATIC_LIB)
-    target_include_directories(${target} PUBLIC ${RCCL_INCLUDE_DIR})
+    target_include_directories(${target} PUBLIC ${RCCL_INCLUDE_DIR}/rccl)
     target_compile_definitions(${target} PUBLIC ${xgboost_rccl_flags})
     target_link_directories(${target} PUBLIC ${HIP_LIB_INSTALL_DIR})
     target_link_libraries(${target} PUBLIC ${RCCL_LIBRARY})
   else()
-    target_include_directories(${target} PRIVATE ${RCCL_INCLUDE_DIR})
+    target_include_directories(${target} PRIVATE ${RCCL_INCLUDE_DIR}/rccl)
     target_compile_definitions(${target} PRIVATE ${xgboost_rccl_flags})
     target_link_directories(${target} PUBLIC ${HIP_LIB_INSTALL_DIR})
     if(NOT USE_DLOPEN_RCCL)
diff --git a/rocgputreeshap b/rocgputreeshap
index 6ceffde02..2fea6734e 160000
--- a/rocgputreeshap
+++ b/rocgputreeshap
@@ -1 +1 @@
-Subproject commit 6ceffde024f8752954550ebcca98caa24b5d158d
+Subproject commit 2fea6734e83cf147c1bbe580ac4713cd50abcad5
diff --git a/src/data/array_interface.cu b/src/data/array_interface.cu
index b29987ff4..569196407 100644
--- a/src/data/array_interface.cu
+++ b/src/data/array_interface.cu
@@ -20,7 +20,6 @@ void ArrayInterfaceHandler::SyncCudaStream(std::int64_t stream) {
        *   case where 0 might be given should either use None, 1, or 2 instead for
        *   clarity.
        */
-    /* ignored for HIP */
 #if !defined(XGBOOST_USE_HIP)
       LOG(FATAL) << "Invalid stream ID in array interface: " << stream;
 #endif
@@ -42,7 +41,7 @@ bool ArrayInterfaceHandler::IsCudaPtr(void const* ptr) {
     return false;
   }
 
-#if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
+#if defined(XGBOOST_USE_CUDA)
   cudaPointerAttributes attr;
   auto err = cudaPointerGetAttributes(&attr, ptr);
   // reset error
@@ -64,6 +63,35 @@ bool ArrayInterfaceHandler::IsCudaPtr(void const* ptr) {
     // other errors, `cudaErrorNoDevice`, `cudaErrorInsufficientDriver` etc.
     return false;
   }
+#elif defined(XGBOOST_USE_HIP)
+  hipPointerAttribute_t attr;
+  auto err = hipPointerGetAttributes(&attr, ptr);
+  // reset error
+  CHECK_EQ(err, hipGetLastError());
+  if (err == hipErrorInvalidValue) {
+    return false;
+  } else if (err == hipSuccess) {
+#if HIP_VERSION_MAJOR < 6
+    switch (attr.memoryType) {
+      case hipMemoryTypeUnified:
+      case hipMemoryTypeHost:
+        return false;
+      default:
+        return true;
+    }
+#else
+    switch (attr.type) {
+      case hipMemoryTypeUnified:
+      case hipMemoryTypeHost:
+        return false;
+      default:
+        return true;
+    }
+#endif
+    return true;
+  } else {
+    return false;
+  }
 #endif
 }
 }  // namespace xgboost