[ROCM] fix cupti, rccl on rocm (#54807)

ronny1996 · web-flow · commit 57da105ccb18 · 2023-06-28T11:22:54.000+08:00
* [ROCM] fix cupti, hipcub

* update

* update
diff --git a/cmake/cupti.cmake b/cmake/cupti.cmake
@@ -2,9 +2,15 @@ if(NOT WITH_GPU AND NOT WITH_ROCM)
   return()
 endif()
 
-set(CUPTI_ROOT
-    "/usr"
-    CACHE PATH "CUPTI ROOT")
+if(WITH_ROCM)
+  set(CUPTI_ROOT
+      "${ROCM_PATH}/CUPTI"
+      CACHE PATH "CUPTI ROOT")
+else()
+  set(CUPTI_ROOT
+      "/usr"
+      CACHE PATH "CUPTI ROOT")
+endif()
 find_path(
   CUPTI_INCLUDE_DIR cupti.h
   PATHS ${CUPTI_ROOT}
diff --git a/cmake/hip.cmake b/cmake/hip.cmake
@@ -106,7 +106,11 @@ list(APPEND HIP_CXX_FLAGS -Wno-duplicate-decl-specifier)
 list(APPEND HIP_CXX_FLAGS -Wno-implicit-int-float-conversion)
 list(APPEND HIP_CXX_FLAGS -Wno-pass-failed)
 list(APPEND HIP_CXX_FLAGS -DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_HIP)
-list(APPEND HIP_CXX_FLAGS -std=c++14)
+if(WITH_CINN)
+  list(APPEND HIP_CXX_FLAGS -std=c++14)
+else()
+  list(APPEND HIP_CXX_FLAGS -std=c++17)
+endif()
 
 if(CMAKE_BUILD_TYPE MATCHES Debug)
   list(APPEND HIP_CXX_FLAGS -g2)
diff --git a/paddle/phi/backends/dynload/rccl.cc b/paddle/phi/backends/dynload/rccl.cc
@@ -28,9 +28,17 @@ RCCL_RAND_ROUTINE_EACH(DEFINE_WRAP);
 RCCL_RAND_ROUTINE_EACH_AFTER_2212(DEFINE_WRAP)
 #endif
 
+#if NCCL_VERSION_CODE >= 2304
+RCCL_RAND_ROUTINE_EACH_AFTER_2304(DEFINE_WRAP)
+#endif
+
 #if NCCL_VERSION_CODE >= 2703
 RCCL_RAND_ROUTINE_EACH_AFTER_2703(DEFINE_WRAP)
 #endif
 
+#if NCCL_VERSION_CODE >= 21100
+RCCL_RAND_ROUTINE_EACH_AFTER_21100(DEFINE_WRAP)
+#endif
+
 }  // namespace dynload
 }  // namespace phi
diff --git a/paddle/phi/backends/dynload/rccl.h b/paddle/phi/backends/dynload/rccl.h
@@ -64,12 +64,23 @@ RCCL_RAND_ROUTINE_EACH(DECLARE_DYNAMIC_LOAD_RCCL_WRAP)
 RCCL_RAND_ROUTINE_EACH_AFTER_2212(DECLARE_DYNAMIC_LOAD_RCCL_WRAP)
 #endif
 
+#if NCCL_VERSION_CODE >= 2304
+#define RCCL_RAND_ROUTINE_EACH_AFTER_2304(__macro) __macro(ncclGetVersion);
+RCCL_RAND_ROUTINE_EACH_AFTER_2304(DECLARE_DYNAMIC_LOAD_RCCL_WRAP)
+#endif
+
 #if NCCL_VERSION_CODE >= 2703
 #define RCCL_RAND_ROUTINE_EACH_AFTER_2703(__macro) \
   __macro(ncclSend);                               \
   __macro(ncclRecv);
 RCCL_RAND_ROUTINE_EACH_AFTER_2703(DECLARE_DYNAMIC_LOAD_RCCL_WRAP)
 #endif
 
+#if NCCL_VERSION_CODE >= 21100
+#define RCCL_RAND_ROUTINE_EACH_AFTER_21100(__macro) \
+  __macro(ncclRedOpCreatePreMulSum);                \
+  __macro(ncclRedOpDestroy);
+RCCL_RAND_ROUTINE_EACH_AFTER_21100(DECLARE_DYNAMIC_LOAD_RCCL_WRAP)
+#endif
 }  // namespace dynload
 }  // namespace phi
diff --git a/paddle/phi/kernels/funcs/blas/blas_impl.hip.h b/paddle/phi/kernels/funcs/blas/blas_impl.hip.h
@@ -999,12 +999,10 @@ inline void Blas<phi::GPUContext>::GEMM(bool transA,
                                         int ldc) const {
   // Note that cublas follows fortran order, so the order is different from
   // the cblas convention.
-  rocblas_operation cuTransA = (transA == CblasNoTrans)
-                                   ? rocblas_operation_none
-                                   : rocblas_operation_transpose;
-  rocblas_operation cuTransB = (transB == CblasNoTrans)
-                                   ? rocblas_operation_none
-                                   : rocblas_operation_transpose;
+  rocblas_operation cuTransA =
+      transA ? rocblas_operation_none : rocblas_operation_transpose;
+  rocblas_operation cuTransB =
+      transB ? rocblas_operation_none : rocblas_operation_transpose;
   PADDLE_ENFORCE_GE(
       context_.GetComputeCapability(),
       80,
diff --git a/paddle/phi/kernels/funcs/top_k_function_cuda.h b/paddle/phi/kernels/funcs/top_k_function_cuda.h
@@ -54,6 +54,15 @@ struct radix_key_codec_base<phi::dtype::float16>
 template <>
 struct radix_key_codec_base<phi::dtype::bfloat16>
     : radix_key_codec_integral<phi::dtype::bfloat16, uint16_t> {};
+
+#if ROCM_VERSION_MAJOR >= 5 && ROCM_VERSION_MINOR >= 4
+template <>
+struct float_bit_mask<phi::dtype::float16> : float_bit_mask<rocprim::half> {};
+
+template <>
+struct float_bit_mask<phi::dtype::bfloat16>
+    : float_bit_mask<rocprim::bfloat16> {};
+#endif
 }  // namespace detail
 }  // namespace rocprim
 namespace cub = hipcub;
diff --git a/paddle/phi/kernels/gpu/argsort_kernel.cu b/paddle/phi/kernels/gpu/argsort_kernel.cu
@@ -40,6 +40,19 @@ namespace detail {
 template <>
 struct radix_key_codec_base<phi::dtype::float16>
     : radix_key_codec_integral<phi::dtype::float16, uint16_t> {};
+
+template <>
+struct radix_key_codec_base<phi::dtype::bfloat16>
+    : radix_key_codec_integral<phi::dtype::bfloat16, uint16_t> {};
+
+#if ROCM_VERSION_MAJOR >= 5 && ROCM_VERSION_MINOR >= 4
+template <>
+struct float_bit_mask<phi::dtype::float16> : float_bit_mask<rocprim::half> {};
+
+template <>
+struct float_bit_mask<phi::dtype::bfloat16>
+    : float_bit_mask<rocprim::bfloat16> {};
+#endif
 }  // namespace detail
 }  // namespace rocprim
 #else
diff --git a/test/custom_kernel/CMakeLists.txt b/test/custom_kernel/CMakeLists.txt
@@ -7,8 +7,7 @@ string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
 set(CUSTOM_ENVS
     PADDLE_SOURCE_DIR=${PADDLE_SOURCE_DIR}
     PADDLE_BINARY_DIR=${PADDLE_BINARY_DIR}
-    CUSTOM_DEVICE_ROOT=${CMAKE_BINARY_DIR}/python/paddle/fluid/tests/custom_kernel
-)
+    CUSTOM_DEVICE_ROOT=${CMAKE_BINARY_DIR}/test)
 
 foreach(TEST_OP ${TEST_OPS})
   py_test(${TEST_OP} SRCS ${TEST_OP}.py ENVS ${CUSTOM_ENVS})
diff --git a/test/custom_kernel/custom_kernel_dot.cc b/test/custom_kernel/custom_kernel_dot.cc
diff --git a/test/custom_kernel/custom_kernel_dot_c.cc b/test/custom_kernel/custom_kernel_dot_c.cc
