From e03e46730c2b2826d734edd6117e7e377214db57 Mon Sep 17 00:00:00 2001
From: ronnywang <524019753@qq.com>
Date: Tue, 9 Mar 2021 11:02:19 +0800
Subject: [PATCH] [ROCM] fix gather_op, sigmoid_cross_entropy_with_logits_op,
 test=develop (#31467)

---
 .../operators/sigmoid_cross_entropy_with_logits_op.cu      | 4 ++++
 paddle/fluid/platform/cuda_helper.h                        | 7 +++++++
 2 files changed, 11 insertions(+)

diff --git a/paddle/fluid/operators/sigmoid_cross_entropy_with_logits_op.cu b/paddle/fluid/operators/sigmoid_cross_entropy_with_logits_op.cu
index b9300f1b23b..8611249a29f 100644
--- a/paddle/fluid/operators/sigmoid_cross_entropy_with_logits_op.cu
+++ b/paddle/fluid/operators/sigmoid_cross_entropy_with_logits_op.cu
@@ -29,7 +29,11 @@ namespace operators {
 
 using Tensor = framework::Tensor;
 
+#ifdef __HIPCC__
+static constexpr int kNumCUDAThreads = 256;
+#else
 static constexpr int kNumCUDAThreads = 512;
+#endif
 static constexpr int kNumMaxinumNumBlocks = 4096;
 
 static inline int NumBlocks(const int N) {
diff --git a/paddle/fluid/platform/cuda_helper.h b/paddle/fluid/platform/cuda_helper.h
index 30c38236c52..fa4ef3f8c12 100644
--- a/paddle/fluid/platform/cuda_helper.h
+++ b/paddle/fluid/platform/cuda_helper.h
@@ -76,10 +76,17 @@ namespace platform {
  *
 */
 
+#ifdef __HIPCC__
+#define CUDA_KERNEL_LOOP_TYPE(i, num, index_type)                     \
+  int64_t __index__ = hipBlockIdx_x * hipBlockDim_x + hipThreadIdx_x; \
+  for (index_type i = __index__; __index__ < (num);                   \
+       __index__ += hipBlockDim_x * hipGridDim_x, i = __index__)
+#else
 #define CUDA_KERNEL_LOOP_TYPE(i, num, index_type)            \
   int64_t __index__ = blockIdx.x * blockDim.x + threadIdx.x; \
   for (index_type i = __index__; __index__ < (num);          \
        __index__ += blockDim.x * gridDim.x, i = __index__)
+#endif
 
 #define CUDA_KERNEL_LOOP(i, num) CUDA_KERNEL_LOOP_TYPE(i, num, int)
 
-- 
GitLab