dtk 2504 update

CMakeLists.txt

@@ -157,7 +157,7 @@ if(USE_HIP)
   project(dgl C CXX HIP)
   include(cmake/modules/ROCM.cmake)
   # target_compile_features(dgl PRIVATE cxx_std_17)
-  set(CMAKE_HIP_ARCHITECTURES gfx906;gfx928;gfx926)
+  set(CMAKE_HIP_ARCHITECTURES gfx906;gfx928;gfx926;gfx936)
   # see https://github.com/NVIDIA/thrust/issues/1401
   if(NOT DEFINED ENV{ROCM_PATH})
     set(ROCM_PATH "/opt/dtk" CACHE PATH "Path to which RoCm has been installed")

dgl_sparse/CMakeLists.txt

@@ -21,10 +21,12 @@ list(GET TORCH_VERSION_LIST 0 TORCH_VERSION_MAJOR)
 list(GET TORCH_VERSION_LIST 1 TORCH_VERSION_MINOR)
 
 set(SPARSE_LINKER_LIBS "")
-
+list(APPEND CMAKE_PREFIX_PATH $ENV{ROCM_PATH})
+set(HIP_PATH $ENV{ROCM_PATH}/hip)
 if(USE_CUDA OR USE_HIP)
+  project(dgl_sparse C CXX HIP)
   add_definitions(-DDGL_USE_CUDA)
-  enable_language(HIP)
+  find_package(HIP REQUIRED PATHS ${HIP_PATH} NO_DEFAULT_PATH)
 endif()
 
 # For windows, define NOMINMAX to avoid conflict with std::min/max

graphbolt/src/cuda/neighbor_sampler.hip

@@ -45,11 +45,12 @@ struct radix_key_codec_base<__hip_bfloat16> : radix_key_codec_floating<__hip_bfl
 };
 }
 }
-
+#if HIP_VERSION_MAJOR<6
 __host__ __device__ bool operator>(const __hip_bfloat16& a, const __hip_bfloat16& b)
 {
   return float(a)>float(b);
 }
+#endif
 
 namespace graphbolt {
 namespace ops {

src/array/cuda/bf16.cuh

@@ -43,7 +43,7 @@ min(__hip_bfloat16 a, __hip_bfloat16 b) {
 #endif
 }
 
-#ifdef __HIPCC__ 
+#if HIP_VERSION_MAJOR < 6 
 // Arithmetic BF16 operations for architecture >= 8.0 are already defined in
 // hip/__hip_bfloat16.h
 // #if defined(__DTK_ARCH__) && (__DTK_ARCH__ < 800)
@@ -143,6 +143,8 @@ __device__ __forceinline__ bool operator<=(
 }
 // #endif  // defined(DTKRT_VERSION) && (DTKRT_VERSION < 12020)
 // #endif  // defined(__DTK_ARCH__) && (__DTK_ARCH__ < 800)
+#endif
+#if __HIPCC__
 __device__
 inline
 __hip_bfloat16 __shfl_down(__hip_bfloat16 var, unsigned int lane_delta, int width = warpSize) {