[ROCM] fix reduce_sum nan in ROCM platform, test=develop (#31780)

46dd1d4a · Qi Li · GitHub · f72d197e · 46dd1d4a
隐藏空白更改
内联并排

Showing with 4 addition and 0 deletion

paddle/fluid/operators/reduce_ops/cub_reduce.h paddle/fluid/operators/reduce_ops/cub_reduce.h +4 -0

未找到文件。
--- a/paddle/fluid/operators/reduce_ops/cub_reduce.h
+++ b/paddle/fluid/operators/reduce_ops/cub_reduce.h
@@ -161,7 +161,11 @@ static inline std::vector<int> GetStrides(const std::vector<int>& dims,
  return strides;
 }

+#ifdef __HIPCC__
+constexpr int kMaxBlockDim = 256;
+#else
 constexpr int kMaxBlockDim = 512;
+#endif

 static inline int GetDesiredBlockDim(int block_dim) {
  return block_dim >= kMaxBlockDim