From c2407f4778f257718e9d6e6c63d45488da62b882 Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Fri, 24 Mar 2023 07:06:00 +0000
Subject: [PATCH 1/8] change judgement for DropoutGradGPUKernelDriver

---
 paddle/phi/kernels/funcs/dropout_impl.cu.h | 60 +++++++++++-----------
 1 file changed, 31 insertions(+), 29 deletions(-)

diff --git a/paddle/phi/kernels/funcs/dropout_impl.cu.h b/paddle/phi/kernels/funcs/dropout_impl.cu.h
index de82561c1fefb..6d0ceba2d9096 100644
--- a/paddle/phi/kernels/funcs/dropout_impl.cu.h
+++ b/paddle/phi/kernels/funcs/dropout_impl.cu.h
@@ -310,7 +310,7 @@ void DropoutFwGPUKernelDriver(
     auto* mask_data = mask->data<uint8_t>();
     size_t size = phi::product(mask->dims());
 
-    if (dropout_prob == 1.0f) {
+    if (abs(dropout_prob / 1.0f - 1) < 5.96e-08) {
 #ifdef PADDLE_WITH_HIP
       PADDLE_ENFORCE_GPU_SUCCESS(
           hipMemsetAsync(y_data, 0, x_numel * sizeof(T), stream));
@@ -454,41 +454,43 @@ void DropoutGradGPUKernelDriver(const phi::GPUContext& dev_ctx,
     // y = factor * x
     ScaleByDropoutFactor<T, MT>(dev_ctx, grad_y, grad_x, factor);
   } else {
-    phi::DenseTensor broadcasted_mask;
-    if (is_dropout_nd) {
-      broadcasted_mask.Resize(grad_y.dims());
-      dev_ctx.template Alloc<uint8_t>(&broadcasted_mask);
-
-      std::vector<const phi::DenseTensor*> broadcast_ins = {&mask};
-      std::vector<phi::DenseTensor*> broadcast_outs = {&broadcasted_mask};
-      phi::funcs::BroadcastKernel<phi::ElementwiseType::kUnary,
-                                  uint8_t,
-                                  uint8_t>(dev_ctx,
-                                           broadcast_ins,
-                                           &broadcast_outs,
-                                           -1,
-                                           kps::IdentityFunctor<uint8_t>());
-    }
-
-    std::vector<const phi::DenseTensor*> ins = {
-        &grad_y, is_dropout_nd ? &broadcasted_mask : &mask};
-    std::vector<phi::DenseTensor*> outs = {grad_x};
-    if (upscale_in_train) {
-      if (dropout_prob == 1.0f) {
+    if (upscale_in_train && abs(dropout_prob / 1.0f - 1) < 5.96e-08) {
 #ifdef PADDLE_WITH_HIP
-        hipMemset(grad_x->data<T>(), 0, grad_x->numel() * sizeof(T));
+      hipMemset(grad_x->data<T>(), 0, grad_x->numel() * sizeof(T));
 #else
-        cudaMemset(grad_x->data<T>(), 0, grad_x->numel() * sizeof(T));
+      cudaMemset(grad_x->data<T>(), 0, grad_x->numel() * sizeof(T));
 #endif
+    } else {
+      MT factor = upscale_in_train
+                      ? static_cast<MT>(1.0f / (1.0f - dropout_prob))
+                      : static_cast<MT>(1.0f);
+      if (is_dropout_nd) {
+        phi::DenseTensor broadcasted_mask;
+
+        broadcasted_mask.Resize(grad_y.dims());
+        dev_ctx.template Alloc<uint8_t>(&broadcasted_mask);
+
+        std::vector<const phi::DenseTensor*> broadcast_ins = {&mask};
+        std::vector<phi::DenseTensor*> broadcast_outs = {&broadcasted_mask};
+        phi::funcs::BroadcastKernel<phi::ElementwiseType::kUnary,
+                                    uint8_t,
+                                    uint8_t>(dev_ctx,
+                                             broadcast_ins,
+                                             &broadcast_outs,
+                                             -1,
+                                             kps::IdentityFunctor<uint8_t>());
+
+        std::vector<const phi::DenseTensor*> ins = {&grad_y, &broadcasted_mask};
+        std::vector<phi::DenseTensor*> outs = {grad_x};
+        phi::funcs::ElementwiseKernel<T>(
+            dev_ctx, ins, &outs, CudaDropoutGradFunctor<T>(factor));
+
       } else {
-        MT factor = static_cast<MT>(1.0f / (1.0f - dropout_prob));
+        std::vector<const phi::DenseTensor*> ins = {&grad_y, &mask};
+        std::vector<phi::DenseTensor*> outs = {grad_x};
         phi::funcs::ElementwiseKernel<T>(
             dev_ctx, ins, &outs, CudaDropoutGradFunctor<T>(factor));
       }
-    } else {
-      MT factor = static_cast<MT>(1.0f);
-      phi::funcs::ElementwiseKernel<T>(
-          dev_ctx, ins, &outs, CudaDropoutGradFunctor<T>(factor));
     }
   }
 }

From 48ff6cb1d5265222a93c05267692b5c5eaab75b0 Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Tue, 28 Mar 2023 05:23:55 +0000
Subject: [PATCH 2/8] add UnrollerWithoutVecSize and after this Loaddata to be
 refined

---
 paddle/phi/kernels/funcs/broadcast_function.h | 134 ++++++++++--------
 paddle/phi/kernels/funcs/elementwise_base.h   |  27 +++-
 2 files changed, 97 insertions(+), 64 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index f96a1764c24a5..054d6cb06aa61 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -31,20 +31,28 @@ namespace funcs {
 
 enum BroadcastLoadType { kMixed = 1, kBroadcast = 2, kElementwise = 3 };
 
-template <typename InT, typename OutT, int Arity>
+template <typename InT, typename OutT, int Arity, typename Functor>
 struct LoaderTypeClassifier {
  public:
   int64_t numel{0};
-  int vec_size{1};
+  int vec_size{4};
   int broadcast_num{0};
   bool all_elementwise{true};
-  phi::Array<int, Arity> use_broadcast;
+  phi::Array<bool, Arity> use_broadcast;
   phi::Array<const _ptr_ InT *__restrict__, Arity> ins_data;
+  // phi::Array<const _ptr_ char *__restrict__, Arity> ins_data;
 
   LoaderTypeClassifier() {}
   LoaderTypeClassifier(const std::vector<const DenseTensor *> &ins,
                        std::vector<DenseTensor *> *outs) {
+    using Traits = phi::funcs::FunctionTraits<Functor>;
+    using ArgsT = typename Traits::ArgsTuple;
+    ArgsT arg;
+    // The Arg VecSize=1 is to match the Unroller template.
     uint64_t out_addr = reinterpret_cast<uint64_t>((*outs)[0]->data<OutT>());
+
+    UnrollerWithoutVecSize<VecSizeGetter, Arity>::step(ins, arg, &vec_size);
+
     for (auto i = 1; i < outs->size(); ++i) {
       PADDLE_ENFORCE_EQ(
           (*outs)[i]->dims(),
@@ -56,28 +64,27 @@ struct LoaderTypeClassifier {
       out_addr =
           (out_addr | reinterpret_cast<uint64_t>((*outs)[i]->data<OutT>()));
     }
-    int out_vec_size =
-        phi::GetVectorizedSize<OutT>(reinterpret_cast<OutT *>(out_addr));
 
-    uint64_t in_addr = static_cast<uint64_t>(0);
+    vec_size = std::min(
+        vec_size,
+        phi::GetVectorizedSize<OutT>(reinterpret_cast<OutT *>(out_addr)));
     numel = (*outs)[0]->numel();
+    // UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, &ins_data);
+
+#pragma unroll
     for (int i = 0; i < Arity; ++i) {
+      // FIXME: get type of ins_data find all parts used in_data
       auto in_data = ins[i]->data<InT>();
       ins_data[i] = (const _ptr_ InT *)(in_data);
-
       bool is_same_dim = ins[i]->numel() == numel;
       if (is_same_dim) {
         use_broadcast[i] = false;
-        in_addr = (in_addr | reinterpret_cast<uint64_t>(in_data));
       } else {
         use_broadcast[i] = true;
         broadcast_num++;
       }
       all_elementwise &= is_same_dim;
     }
-    int in_vec_size = std::min(
-        4, phi::GetVectorizedSize<InT>(reinterpret_cast<InT *>(in_addr)));
-    vec_size = std::min(out_vec_size, in_vec_size);
   }
 };
 
@@ -89,7 +96,7 @@ struct BroadcastDataLoader {
       T args[Arity][VecSize],
       const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
       const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<int, Arity> &use_broadcast,
+      const phi::Array<bool, Arity> &use_broadcast,
       const int block_offset,
       const int num,
       const uint32_t numel) {
@@ -114,7 +121,7 @@ struct BroadcastDataLoader<T, VecSize, Arity, true, kElementwise> {
       T args[Arity][VecSize],
       const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
       const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<int, Arity> &use_broadcast,
+      const phi::Array<bool, Arity> &use_broadcast,
       const int block_offset,
       const int num,
       const uint32_t numel) {
@@ -140,7 +147,7 @@ struct BroadcastDataLoader<T, VecSize, Arity, false, kElementwise> {
       T args[Arity][VecSize],
       const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
       const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<int, Arity> &use_broadcast,
+      const phi::Array<bool, Arity> &use_broadcast,
       const int block_offset,
       const int num,
       const uint32_t numel) {
@@ -168,7 +175,7 @@ struct BroadcastDataLoader<T, VecSize, Arity, IsBoundary, kBroadcast> {
       T args[Arity][VecSize],
       const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
       const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<int, Arity> &use_broadcast,
+      const phi::Array<bool, Arity> &use_broadcast,
       const int block_offset,
       const int num,
       const uint32_t numel) {
@@ -224,7 +231,7 @@ template <typename InT,
 __device__ void VectorizedBroadcastKernelImpl(
     const phi::Array<const _ptr_ InT *__restrict__, Arity> &ins,
     phi::Array<_ptr_ OutT *, NumOuts> outs,
-    const phi::Array<int, Arity> &use_broadcast,
+    const phi::Array<bool, Arity> &use_broadcast,
     const uint32_t numel,
     const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
     int num,
@@ -274,7 +281,7 @@ template <typename Functor,
 __global__ void VectorizedBroadcastKernel(
     phi::Array<const _ptr_ InT *__restrict__, Arity> ins,
     phi::Array<_ptr_ OutT *, NumOuts> outs,
-    phi::Array<int, Arity> use_broadcast,
+    phi::Array<bool, Arity> use_broadcast,
     uint32_t numel,
     phi::Array<kps::details::BroadcastConfig, Arity> configs,
     int main_offset,
@@ -363,7 +370,7 @@ __global__ void VectorizedBroadcastKernel(
 
 template <typename InT,
           typename OutT,
-          typename Func,
+          typename Functor,
           int Arity,
           int NumOuts,
           int VecSize>
@@ -371,9 +378,9 @@ void LaunchBroadcastKernel(
     const KPDevice &ctx,
     const std::vector<const DenseTensor *> &ins,
     std::vector<DenseTensor *> *outs,
-    Func func,
+    Functor func,
     const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-    const LoaderTypeClassifier<InT, OutT, Arity> &loader_classifier) {
+    const LoaderTypeClassifier<InT, OutT, Arity, Functor> &loader_classifier) {
   phi::Array<_ptr_ OutT *, NumOuts> outs_data;
   for (int i = 0; i < NumOuts; ++i) {
     outs_data[i] = (_ptr_ OutT *)(ctx.Alloc<OutT>((*outs)[i]));
@@ -388,7 +395,7 @@ void LaunchBroadcastKernel(
   int main_offset = (numel / (read_lens * threads)) * read_lens * threads;
   int tail_tid = numel % (read_lens * threads);
 
-  VectorizedBroadcastKernel<Func, InT, OutT, Arity, NumOuts, VecSize, false>
+  VectorizedBroadcastKernel<Functor, InT, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                        outs_data,
                                        loader_classifier.use_broadcast,
@@ -409,7 +416,7 @@ void LaunchBroadcastKernel(
   int tail_tid = numel % (VecSize * threads);
 
   if (loader_classifier.all_elementwise) {
-    VectorizedBroadcastKernel<Func,
+    VectorizedBroadcastKernel<Functor,
                               InT,
                               OutT,
                               Arity,
@@ -427,7 +434,13 @@ void LaunchBroadcastKernel(
                                          func);
   } else if (loader_classifier.broadcast_num > (Arity >> 1)) {
     constexpr BroadcastLoadType type_ = (Arity > 1) ? kBroadcast : kMixed;
-    VectorizedBroadcastKernel<Func, InT, OutT, Arity, NumOuts, VecSize, type_>
+    VectorizedBroadcastKernel<Functor,
+                              InT,
+                              OutT,
+                              Arity,
+                              NumOuts,
+                              VecSize,
+                              type_>
         <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                          outs_data,
                                          loader_classifier.use_broadcast,
@@ -438,7 +451,13 @@ void LaunchBroadcastKernel(
                                          VecSize,
                                          func);
   } else {
-    VectorizedBroadcastKernel<Func, InT, OutT, Arity, NumOuts, VecSize, kMixed>
+    VectorizedBroadcastKernel<Functor,
+                              InT,
+                              OutT,
+                              Arity,
+                              NumOuts,
+                              VecSize,
+                              kMixed>
         <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                          outs_data,
                                          loader_classifier.use_broadcast,
@@ -847,6 +866,7 @@ template <ElementwiseType ET,
           typename InT,
           typename OutT,
           typename Functor,
+          int kArity,
           int NumOuts = 1>
 void BroadcastKernelForDifferentVecSize(
     const KPDevice &ctx,
@@ -854,43 +874,14 @@ void BroadcastKernelForDifferentVecSize(
     std::vector<DenseTensor *> *outs,
     int axis,
     Functor func) {
-  using Traits = phi::funcs::FunctionTraits<Functor>;
-  const int kArity =
-      Traits::has_pointer_args ? static_cast<int>(ET) : Traits::arity;
-  PADDLE_ENFORCE_EQ(
-      ins.size(),
-      kArity,
-      phi::errors::InvalidArgument("The number of inputs is expected to be "
-                                   "equal to the "
-                                   "arity of functor. But received: the "
-                                   "number of inputs "
-                                   "is %d, the arity of functor is %d.",
-                                   ins.size(),
-                                   kArity));
-  PADDLE_ENFORCE_LE(
-      kArity,
-      3,
-      phi::errors::InvalidArgument("Currently only broadcast of ternary is "
-                                   "supported "
-                                   "and verified, but received %d.",
-                                   kArity));
-  PADDLE_ENFORCE_EQ(
-      outs->size(),
-      NumOuts,
-      phi::errors::InvalidArgument("Number of outputs shall equal to number "
-                                   "of functions, "
-                                   "but number of outputs is %d, of "
-                                   "functions is %d.",
-                                   outs->size(),
-                                   NumOuts));
-
 #ifndef PADDLE_WITH_XPU_KP
   constexpr bool kEnabledInt64IndexKernel = (NumOuts == 1 && kArity <= 3);
   bool use_int64_index_kernel =
       kEnabledInt64IndexKernel &&
       (*outs)[0]->numel() >= std::numeric_limits<int32_t>::max();
   if (use_int64_index_kernel) {
-    auto loader_classifier = LoaderTypeClassifier<InT, OutT, kArity>(ins, outs);
+    auto loader_classifier =
+        LoaderTypeClassifier<InT, OutT, kArity, Functor>(ins, outs);
     switch (loader_classifier.vec_size) {
       case VecSizeL: {
         LaunchBroadcastKernelWithInt64IndexHelper<InT,
@@ -949,7 +940,7 @@ void BroadcastKernelForDifferentVecSize(
       phi::errors::InvalidArgument(
           "XPU only support inputs is 2, but received %d", ins.size()));
 
-  auto loader_classifier = LoaderTypeClassifier<InT, OutT, kArity>();
+  auto loader_classifier = LoaderTypeClassifier<InT, OutT, kArity, Functor>();
   const auto dims_simplifier =
       BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
   if (VLOG_IS_ON(6)) {
@@ -968,7 +959,8 @@ void BroadcastKernelForDifferentVecSize(
   bool is_optimize = configs[0].cmp_type != type;
   int vec_size = is_optimize ? VecSizeL : VecSizeM;
 #else
-  auto loader_classifier = LoaderTypeClassifier<InT, OutT, kArity>(ins, outs);
+  auto loader_classifier =
+      LoaderTypeClassifier<InT, OutT, kArity, Functor>(ins, outs);
   if (!loader_classifier.all_elementwise) {
     const auto dims_simplifier =
         BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
@@ -1013,6 +1005,8 @@ void BroadcastKernelForDifferentVecSize(
   }
 }
 
+// FIXME: delete (ElementwiseType ET & typename InT)
+// default: axis = 1
 template <ElementwiseType ET,
           typename InT,
           typename OutT,
@@ -1025,6 +1019,30 @@ void BroadcastKernel(const KPDevice &ctx,
                      Functor func) {
   // When there are multiple inputs, the outputs's rank should be equal the
   // maximum rank of all inputs.
+  // FIXME: delete ET ？
+  using Traits = phi::funcs::FunctionTraits<Functor>;
+  const int kArity =
+      Traits::has_pointer_args ? static_cast<int>(ET) : Traits::arity;
+  PADDLE_ENFORCE_EQ(
+      ins.size(),
+      kArity,
+      phi::errors::InvalidArgument("The number of inputs is expected to be "
+                                   "equal to the "
+                                   "arity of functor. But received: the "
+                                   "number of inputs "
+                                   "is %d, the arity of functor is %d.",
+                                   ins.size(),
+                                   kArity));
+  PADDLE_ENFORCE_EQ(
+      outs->size(),
+      NumOuts,
+      phi::errors::InvalidArgument("Number of outputs shall equal to number "
+                                   "of functions, "
+                                   "but number of outputs is %d, of "
+                                   "functions is %d.",
+                                   outs->size(),
+                                   NumOuts));
+
   int max_rank = 0;
   int min_rank = phi::DDim::kMaxRank;
   for (auto *in : ins) {
@@ -1037,7 +1055,7 @@ void BroadcastKernel(const KPDevice &ctx,
     max_rank = std::max(max_rank, (*outs)[0]->dims().size());
   }
   axis = axis == -1 ? max_rank - min_rank : axis;
-  BroadcastKernelForDifferentVecSize<ET, InT, OutT, Functor, NumOuts>(
+  BroadcastKernelForDifferentVecSize<ET, InT, OutT, Functor, kArity, NumOuts>(
       ctx, ins, outs, axis, func);
 }
 
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
index 1d40d4d8c2957..e2a8d0b569819 100644
--- a/paddle/phi/kernels/funcs/elementwise_base.h
+++ b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -35,7 +35,7 @@ namespace kps = phi::kps;
 
 namespace phi {
 
-enum ElementwiseType { kUnary = 1, kBinary = 2, kTernary = 3, kAny = -1 };
+enum ElementwiseType { kUnary = 1, kBinary = 2, kTernary = 3 };
 /* Packing scalar type T(float, int etc.) into Array<T, NumOuts> type
    for supporting multiple-output feature in elementwise system.*/
 template <class T, int Num>
@@ -508,6 +508,22 @@ struct Unroller<Func, VecSize, End, End> {
   static HOSTDEVICE inline void step(Args &&...args) {}
 };
 
+// static unroller without VecSize for broadcast
+template <template <int Index> typename Func, int End, int Begin = 0>
+struct UnrollerWithoutVecSize {
+  template <typename... Args>
+  static HOSTDEVICE inline void step(Args &&...args) {
+    Func<Begin>::Apply(std::forward<Args>(args)...);
+    UnrollerWithoutVecSize<Func, End, Begin + 1>::step(args...);
+  }
+};
+
+template <template <int Index> typename Func, int End>
+struct UnrollerWithoutVecSize<Func, End, End> {
+  template <typename... Args>
+  static HOSTDEVICE inline void step(Args &&...args) {}
+};
+
 template <int Index, int VecSize>
 struct Loader {
   template <typename Array, typename ArgsT>
@@ -536,7 +552,7 @@ struct Loader {
   }
 };
 
-template <int Index, int VecSize>
+template <int Index>
 struct InputSetter {
   template <typename Array>
   static HOSTDEVICE void Apply(
@@ -545,7 +561,7 @@ struct InputSetter {
   }
 };
 
-template <int Index, int VecSize>
+template <int Index>
 struct VecSizeGetter {
   template <typename ArgsT>
   static HOSTDEVICE void Apply(const std::vector<const DenseTensor *> &ins,
@@ -569,8 +585,7 @@ int GetVectorizedSizeForTensors(const std::vector<const DenseTensor *> &ins,
   int vec_size = 4;
   uint64_t addr = static_cast<uint64_t>(0);
   ArgsT arg;
-  // The Arg VecSize=1 is to match the Unroller template.
-  Unroller<VecSizeGetter, 1, Arity>::step(ins, arg, &vec_size);
+  UnrollerWithoutVecSize<VecSizeGetter, Arity>::step(ins, arg, &vec_size);
   for (auto iter = outs.begin(); iter != outs.end(); ++iter) {
     addr = (addr | reinterpret_cast<uint64_t>((*iter)->data<OutT>()));
   }
@@ -802,7 +817,7 @@ void LaunchElementwiseCudaKernel(const KPDevice &ctx,
   phi::Array<const _ptr_ char *__restrict__, Arity> ins_data;
   phi::Array<_ptr_ OutT *, NumOuts> outs_data;
 
-  Unroller<InputSetter, VecSize, Arity>::step(ins, &ins_data);
+  UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, &ins_data);
   for (int i = 0; i < NumOuts; ++i) {
     outs_data[i] = (_ptr_ OutT *)(ctx.Alloc<OutT>((*outs)[i]));
   }

From 9971126efd2c37da96ef96aac6c37d6e13940901 Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Fri, 7 Apr 2023 10:34:35 +0000
Subject: [PATCH 3/8] pass unittest

---
 paddle/phi/kernels/funcs/broadcast_function.h | 281 ++++++++----------
 paddle/phi/kernels/funcs/elementwise_base.h   |  80 +----
 .../kernels/primitive/datamover_primitives.h  |  92 +++++-
 3 files changed, 218 insertions(+), 235 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 054d6cb06aa61..cbb6190bbe465 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -31,7 +31,30 @@ namespace funcs {
 
 enum BroadcastLoadType { kMixed = 1, kBroadcast = 2, kElementwise = 3 };
 
-template <typename InT, typename OutT, int Arity, typename Functor>
+template <int Index>
+struct UseBroadcast {
+  template <typename ArgsT, typename Array1, typename Array2>
+  static HOSTDEVICE void Apply(
+      const std::vector<const DenseTensor *> &ins_tensor,
+      const ArgsT &args,
+      int64_t numel,
+      Array1 *ins_data,
+      Array2 *use_broadcast,
+      int *broadcast_num,
+      bool *all_elementwise) {
+    (*ins_data)[Index] = (const _ptr_ char *)(ins_tensor[Index]->data());
+    bool is_same_dim = ins_tensor[Index]->numel() == numel;
+    if (is_same_dim) {
+      (*use_broadcast)[Index] = false;
+    } else {
+      (*use_broadcast)[Index] = true;
+      (*broadcast_num)++;
+    }
+    *all_elementwise &= is_same_dim;
+  }
+};
+
+template <typename OutT, int Arity, typename Functor>
 struct LoaderTypeClassifier {
  public:
   int64_t numel{0};
@@ -39,8 +62,7 @@ struct LoaderTypeClassifier {
   int broadcast_num{0};
   bool all_elementwise{true};
   phi::Array<bool, Arity> use_broadcast;
-  phi::Array<const _ptr_ InT *__restrict__, Arity> ins_data;
-  // phi::Array<const _ptr_ char *__restrict__, Arity> ins_data;
+  phi::Array<const _ptr_ char *__restrict__, Arity> ins_data;
 
   LoaderTypeClassifier() {}
   LoaderTypeClassifier(const std::vector<const DenseTensor *> &ins,
@@ -48,7 +70,6 @@ struct LoaderTypeClassifier {
     using Traits = phi::funcs::FunctionTraits<Functor>;
     using ArgsT = typename Traits::ArgsTuple;
     ArgsT arg;
-    // The Arg VecSize=1 is to match the Unroller template.
     uint64_t out_addr = reinterpret_cast<uint64_t>((*outs)[0]->data<OutT>());
 
     UnrollerWithoutVecSize<VecSizeGetter, Arity>::step(ins, arg, &vec_size);
@@ -69,156 +90,81 @@ struct LoaderTypeClassifier {
         vec_size,
         phi::GetVectorizedSize<OutT>(reinterpret_cast<OutT *>(out_addr)));
     numel = (*outs)[0]->numel();
-    // UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, &ins_data);
-
-#pragma unroll
-    for (int i = 0; i < Arity; ++i) {
-      // FIXME: get type of ins_data find all parts used in_data
-      auto in_data = ins[i]->data<InT>();
-      ins_data[i] = (const _ptr_ InT *)(in_data);
-      bool is_same_dim = ins[i]->numel() == numel;
-      if (is_same_dim) {
-        use_broadcast[i] = false;
-      } else {
-        use_broadcast[i] = true;
-        broadcast_num++;
-      }
-      all_elementwise &= is_same_dim;
-    }
+    UnrollerWithoutVecSize<UseBroadcast, Arity>::step(ins,
+                                                      arg,
+                                                      numel,
+                                                      &ins_data,
+                                                      &use_broadcast,
+                                                      &broadcast_num,
+                                                      &all_elementwise);
   }
 };
 
 #ifndef PADDLE_WITH_XPU_KP
+
 // Common broadcast/elementwise Loader.
-template <typename T, int VecSize, int Arity, bool IsBoundary, int LoadType>
+template <int Index, int VecSize, bool IsBoundary, int LoadType>
 struct BroadcastDataLoader {
-  __device__ __forceinline__ void operator()(
-      T args[Arity][VecSize],
-      const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
-      const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<bool, Arity> &use_broadcast,
-      const int block_offset,
-      const int num,
-      const uint32_t numel) {
-#pragma unroll
-    for (int i = 0; i < Arity; ++i) {
-      kps::Init<T, VecSize>(args[i], static_cast<T>(1.0f));
-      if (use_broadcast[i]) {
-        kps::ReadDataBc<T, VecSize, 1, IsBoundary>(
-            args[i], ins[i], block_offset, configs[i], numel, VecSize);
-      } else {
-        kps::ReadData<T, VecSize, 1, IsBoundary>(
-            args[i], ins[i] + block_offset, num, VecSize);
-      }
+  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
+  static __device__ __forceinline__ void Apply(const Array1 &ins,
+                                               ArgsT *args,
+                                               const Array2 &configs,
+                                               const Array3 &use_broadcast,
+                                               const int block_offset,
+                                               const int num,
+                                               const uint32_t numel) {
+    using Type = std::tuple_element_t<Index, ArgsT>;
+    kps::Init<Type, ArgsT, Index, VecSize>(args, static_cast<Type>(1.0f));
+
+    if (use_broadcast[Index]) {
+      kps::ReadDataBc<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
+          args,
+          reinterpret_cast<const _ptr_ Type *>(ins[Index]),
+          block_offset,
+          configs[Index],
+          numel,
+          VecSize);
+    } else {
+      kps::ReadData<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
+          args,
+          reinterpret_cast<const _ptr_ Type *>(ins[Index]) + block_offset,
+          num,
+          VecSize);
     }
   }
 };
+// FIXME: add BroadcastDataLoaders Partial specialization here
 
-// Scalar elementwise Loader with consideration of IsBoundary.
-template <typename T, int VecSize, int Arity>
-struct BroadcastDataLoader<T, VecSize, Arity, true, kElementwise> {
-  __device__ __forceinline__ void operator()(
-      T args[Arity][VecSize],
-      const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
-      const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<bool, Arity> &use_broadcast,
-      const int block_offset,
-      const int num,
-      const uint32_t numel) {
-    int thread_offset = threadIdx.x * VecSize + block_offset;
-#pragma unroll
-    for (int i = 0; i < Arity; ++i) {
-#pragma unroll
-      for (int idx = 0; idx < VecSize; ++idx) {
-        args[i][idx] = static_cast<T>(1);
-        int index = thread_offset + idx;
-        if (index < numel) {
-          args[i][idx] = ins[i][index];
-        }
-      }
-    }
-  }
-};
+#endif
 
-// Vectorized elementwise Loader without consideration of IsBoundary.
-template <typename T, int VecSize, int Arity>
-struct BroadcastDataLoader<T, VecSize, Arity, false, kElementwise> {
-  __device__ __forceinline__ void operator()(
-      T args[Arity][VecSize],
-      const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
-      const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<bool, Arity> &use_broadcast,
-      const int block_offset,
-      const int num,
-      const uint32_t numel) {
-    using VecType = phi::kps::details::VectorType<T, VecSize>;
-    VecType vec_temp[Arity];
-
-    int thread_offset = threadIdx.x + blockIdx.x * blockDim.x;
-#pragma unroll
-    for (int i = 0; i < Arity; ++i) {
-      const VecType *__restrict__ vec_input =
-          reinterpret_cast<const VecType *__restrict__>(ins[i]);
-      vec_temp[i] = vec_input[thread_offset];
-#pragma unroll
-      for (int idx = 0; idx < VecSize; ++idx) {
-        args[i][idx] = vec_temp[i].val[idx];
-      }
-    }
+// static broadcast unroller
+template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
+          typename Func,
+          bool IsBoundary,
+          int LoadType,
+          int VecSize,
+          int End,
+          int Begin = 0>
+struct BcUnroller {
+  template <typename... Args>
+  static HOSTDEVICE inline void step(Args &&...args) {
+    Func<Begin, VecSize, IsBoundary, LoadType>::Apply(
+        std::forward<Args>(args)...);
+    BcUnroller<Func, IsBoundary, LoadType, VecSize, End, Begin + 1>::step(
+        args...);
   }
 };
 
-// Common broadcast data loader.
-template <typename T, int VecSize, int Arity, bool IsBoundary>
-struct BroadcastDataLoader<T, VecSize, Arity, IsBoundary, kBroadcast> {
-  __device__ __forceinline__ void operator()(
-      T args[Arity][VecSize],
-      const phi::Array<const _ptr_ T *__restrict__, Arity> &ins,
-      const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-      const phi::Array<bool, Arity> &use_broadcast,
-      const int block_offset,
-      const int num,
-      const uint32_t numel) {
-    uint32_t index_bc[Arity][VecSize];
-#pragma unroll
-    for (int j = 0; j < Arity; ++j) {
-#pragma unroll
-      for (int k = 0; k < VecSize; ++k) {
-        index_bc[j][k] = 0;
-        args[j][k] = static_cast<T>(1);
-      }
-    }
-
-    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
-#pragma unroll
-    for (int k = 0; k < VecSize; ++k) {
-      uint32_t idx = thread_offset + k;
-      if (IsBoundary) {
-        if (idx == numel) break;
-      }
-
-#pragma unroll
-      for (int i = 0; i < phi::DDim::kMaxRank; ++i) {
-        if (i == configs[0].rank) break;
-        auto fast_divmoder = configs[0].divmoders[i].Divmod(idx);
-        idx = fast_divmoder.val[0];
-#pragma unroll
-        for (int j = 0; j < Arity; ++j) {
-          index_bc[j][k] += fast_divmoder.val[1] * configs[j].strides[i];
-        }
-      }
-    }
-
-#pragma unroll
-    for (int j = 0; j < Arity; ++j) {
-#pragma unroll
-      for (int k = 0; k < VecSize; ++k) {
-        args[j][k] = ins[j][index_bc[j][k]];
-      }
-    }
-  }
+template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
+          typename Func,
+          bool IsBoundary,
+          int LoadType,
+          int VecSize,
+          int End>
+struct BcUnroller<Func, IsBoundary, LoadType, VecSize, End, End> {
+  template <typename... Args>
+  static HOSTDEVICE inline void step(Args &&...args) {}
 };
-#endif
 
 template <typename InT,
           typename OutT,
@@ -229,7 +175,7 @@ template <typename InT,
           bool IsBoundary,
           int LoadType>
 __device__ void VectorizedBroadcastKernelImpl(
-    const phi::Array<const _ptr_ InT *__restrict__, Arity> &ins,
+    const phi::Array<const _ptr_ char *__restrict__, Arity> &ins,
     phi::Array<_ptr_ OutT *, NumOuts> outs,
     const phi::Array<bool, Arity> &use_broadcast,
     const uint32_t numel,
@@ -238,9 +184,14 @@ __device__ void VectorizedBroadcastKernelImpl(
     int block_offset,
     int read_lens,
     Functor func) {
-  __simd__ InT args[Arity][VecSize];
+  using Traits = phi::funcs::FunctionTraits<Functor>;
+  using ArgsT = typename Traits::ArgsTuple;
+  __simd__ ArgsT args[VecSize];
   __simd__ ConditionalT<OutT, NumOuts> result[VecSize];
+
+// FIXME: rm xpu InT
 #ifdef PADDLE_WITH_XPU_KP
+  __simd__ InT args[Arity][VecSize];
 #pragma unroll
   for (int i = 0; i < Arity; ++i) {
     kps::Init<InT, VecSize>(args[i], static_cast<InT>(1.0f), read_lens);
@@ -253,19 +204,28 @@ __device__ void VectorizedBroadcastKernelImpl(
     }
   }
 #else
-  BroadcastDataLoader<InT, VecSize, Arity, IsBoundary, LoadType>()(
-      args, ins, configs, use_broadcast, block_offset, num, numel);
+
+  BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
+      ins, args, configs, use_broadcast, block_offset, num, numel);
 #endif
 
-  constexpr bool kCallElementwiseAny =
-      phi::funcs::FunctionTraits<Functor>::has_pointer_args;
-  phi::funcs::ElementwisePrimitiveCaller<InT,
-                                         ConditionalT<OutT, NumOuts>,
-                                         VecSize,
-                                         Functor,
-                                         Arity,
-                                         kCallElementwiseAny>()(
-      func, args, result, read_lens);
+  // constexpr bool kCallElementwiseAny =
+  //     phi::funcs::FunctionTraits<Functor>::has_pointer_args;
+
+  SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
+                                     VecSize,
+                                     Functor,
+                                     ArgsT,
+                                     Arity>()(func, args, result, read_lens);
+
+  // FIXME: clean this
+  // phi::funcs::ElementwisePrimitiveCaller<InT,
+  //                                        ConditionalT<OutT, NumOuts>,
+  //                                        VecSize,
+  //                                        Functor,
+  //                                        Arity,
+  //                                        kCallElementwiseAny>()(
+  //     func, args, result, read_lens);
   phi::funcs::
       ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, NumOuts>()(
           outs, result, block_offset, num, read_lens);
@@ -279,7 +239,7 @@ template <typename Functor,
           int VecSize,
           int LoadType>
 __global__ void VectorizedBroadcastKernel(
-    phi::Array<const _ptr_ InT *__restrict__, Arity> ins,
+    phi::Array<const _ptr_ char *__restrict__, Arity> ins,
     phi::Array<_ptr_ OutT *, NumOuts> outs,
     phi::Array<bool, Arity> use_broadcast,
     uint32_t numel,
@@ -380,7 +340,7 @@ void LaunchBroadcastKernel(
     std::vector<DenseTensor *> *outs,
     Functor func,
     const phi::Array<kps::details::BroadcastConfig, Arity> &configs,
-    const LoaderTypeClassifier<InT, OutT, Arity, Functor> &loader_classifier) {
+    const LoaderTypeClassifier<OutT, Arity, Functor> &loader_classifier) {
   phi::Array<_ptr_ OutT *, NumOuts> outs_data;
   for (int i = 0; i < NumOuts; ++i) {
     outs_data[i] = (_ptr_ OutT *)(ctx.Alloc<OutT>((*outs)[i]));
@@ -862,6 +822,7 @@ struct LaunchBroadcastKernelWithInt64IndexHelper<InT,
 };
 #endif
 
+// FIXME: delete ElementwiseType InT
 template <ElementwiseType ET,
           typename InT,
           typename OutT,
@@ -881,7 +842,7 @@ void BroadcastKernelForDifferentVecSize(
       (*outs)[0]->numel() >= std::numeric_limits<int32_t>::max();
   if (use_int64_index_kernel) {
     auto loader_classifier =
-        LoaderTypeClassifier<InT, OutT, kArity, Functor>(ins, outs);
+        LoaderTypeClassifier<OutT, kArity, Functor>(ins, outs);
     switch (loader_classifier.vec_size) {
       case VecSizeL: {
         LaunchBroadcastKernelWithInt64IndexHelper<InT,
@@ -940,7 +901,7 @@ void BroadcastKernelForDifferentVecSize(
       phi::errors::InvalidArgument(
           "XPU only support inputs is 2, but received %d", ins.size()));
 
-  auto loader_classifier = LoaderTypeClassifier<InT, OutT, kArity, Functor>();
+  auto loader_classifier = LoaderTypeClassifier<OutT, kArity, Functor>();
   const auto dims_simplifier =
       BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
   if (VLOG_IS_ON(6)) {
@@ -960,7 +921,7 @@ void BroadcastKernelForDifferentVecSize(
   int vec_size = is_optimize ? VecSizeL : VecSizeM;
 #else
   auto loader_classifier =
-      LoaderTypeClassifier<InT, OutT, kArity, Functor>(ins, outs);
+      LoaderTypeClassifier<OutT, kArity, Functor>(ins, outs);
   if (!loader_classifier.all_elementwise) {
     const auto dims_simplifier =
         BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
@@ -1006,7 +967,7 @@ void BroadcastKernelForDifferentVecSize(
 }
 
 // FIXME: delete (ElementwiseType ET & typename InT)
-// default: axis = 1
+// default: axis = -1
 template <ElementwiseType ET,
           typename InT,
           typename OutT,
@@ -1021,8 +982,7 @@ void BroadcastKernel(const KPDevice &ctx,
   // maximum rank of all inputs.
   // FIXME: delete ET ？
   using Traits = phi::funcs::FunctionTraits<Functor>;
-  const int kArity =
-      Traits::has_pointer_args ? static_cast<int>(ET) : Traits::arity;
+  const int kArity = Traits::arity;
   PADDLE_ENFORCE_EQ(
       ins.size(),
       kArity,
@@ -1069,6 +1029,7 @@ void ElementwiseCompute(const GPUContext &dev_ctx,
   std::vector<const DenseTensor *> ins = {&x, &y};
   std::vector<DenseTensor *> outs = {z};
   dev_ctx.template Alloc<OutType>(z);
+  // FIXME: delete ElementwiseType
   BroadcastKernel<ElementwiseType::kBinary, T, OutType, Functor, 1>(
       dev_ctx, ins, &outs, axis, func);
 }
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
index e2a8d0b569819..14f0cecf42d75 100644
--- a/paddle/phi/kernels/funcs/elementwise_base.h
+++ b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -35,6 +35,7 @@ namespace kps = phi::kps;
 
 namespace phi {
 
+// FIXME: delete this enum
 enum ElementwiseType { kUnary = 1, kBinary = 2, kTernary = 3 };
 /* Packing scalar type T(float, int etc.) into Array<T, NumOuts> type
    for supporting multiple-output feature in elementwise system.*/
@@ -527,12 +528,12 @@ struct UnrollerWithoutVecSize<Func, End, End> {
 template <int Index, int VecSize>
 struct Loader {
   template <typename Array, typename ArgsT>
-  static __device__ void Apply(const Array &in,
-                               ArgsT *args,
-                               kps::IndexType offset,
-                               int num,
-                               int read_lens,
-                               bool is_boundary) {
+  static __device__ __forceinline__ void Apply(const Array &in,
+                                               ArgsT *args,
+                                               kps::IndexType offset,
+                                               int num,
+                                               int read_lens,
+                                               bool is_boundary) {
     using Type = std::tuple_element_t<Index, ArgsT>;
     kps::Init<Type, ArgsT, Index, VecSize>(
         args, static_cast<Type>(1.0f), read_lens);
@@ -595,73 +596,6 @@ int GetVectorizedSizeForTensors(const std::vector<const DenseTensor *> &ins,
   return vec_size;
 }
 
-template <typename InT,
-          typename OutT,
-          int VecSize,
-          typename Functor,
-          int Arity,
-          bool CallElementwiseAny = false>
-struct ElementwisePrimitiveCaller {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens);
-};
-
-template <typename InT, typename OutT, int VecSize, typename Functor, int Arity>
-struct ElementwisePrimitiveCaller<InT, OutT, VecSize, Functor, Arity, true> {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens) {
-    kps::ElementwiseAny<InT, OutT, VecSize, 1, Arity, Functor>(
-        result, args, func);
-  }
-};
-
-template <typename InT, typename OutT, int VecSize, typename Functor>
-struct ElementwisePrimitiveCaller<InT, OutT, VecSize, Functor, 0, false> {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens) {
-    kps::ElementwiseConstant<InT, OutT, VecSize, 1, Functor>(result, func);
-  }
-};
-
-template <typename InT, typename OutT, int VecSize, typename Functor>
-struct ElementwisePrimitiveCaller<InT, OutT, VecSize, Functor, 1, false> {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens) {
-    kps::ElementwiseUnary<InT, OutT, VecSize, 1, Functor>(
-        result, args[0], func);
-  }
-};
-
-template <typename InT, typename OutT, int VecSize, typename Functor>
-struct ElementwisePrimitiveCaller<InT, OutT, VecSize, Functor, 2, false> {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens) {
-    kps::ElementwiseBinary<InT, OutT, VecSize, 1, Functor>(
-        result, args[0], args[1], func, read_lens);
-  }
-};
-
-template <typename InT, typename OutT, int VecSize, typename Functor>
-struct ElementwisePrimitiveCaller<InT, OutT, VecSize, Functor, 3, false> {
-  __device__ inline void operator()(Functor func,
-                                    InT (*args)[VecSize],
-                                    OutT *result,
-                                    int read_lens) {
-    kps::ElementwiseTernary<InT, OutT, VecSize, 1, Functor>(
-        result, args[0], args[1], args[2], func);
-  }
-};
-
 namespace detail {
 template <class F, class Tuple, std::size_t... Index>
 // GCC/Clang need the decltype() return type
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index ac27916196105..ca9de686d817f 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -255,6 +255,18 @@ __device__ __forceinline__ void Init(ArgsT* dst, T init_data, int read_lens) {
   }
 }
 
+/**
+ * The difference from the above function is that
+ * it supports different data types of inputs.
+ */
+template <typename T, typename ArgsT, int Index, int NX>
+__device__ __forceinline__ void Init(ArgsT* dst, T init_data) {
+#pragma unroll
+  for (int i = 0; i < NX; i++) {
+    std::get<Index>(dst[i]) = init_data;
+  }
+}
+
 /**
  * @brief Read 1D data from global memory to register. When IsBoundary = true
  * and (NX % 4 == 0 or Nx % 2 == 0), vectorized load data will be used to
@@ -307,6 +319,23 @@ __device__ __forceinline__ void ReadData(T* dst,
   }
 }
 
+/**
+ * @brief Read 1D data from global memory to register.
+ * @template paraments
+ * T: The type of data.
+ * NX: Each thread load NX data from global memory continuously.
+ * NY: Each thread need to load NY rows, only NY = 1 was supported.
+ * IsBoundary: Whether to make an out-of-bounds judgment on access to memory.
+ * When the number of data processed by this block is less than
+ * NX x NY x blockDim.x, boundary judgment is required to avoid memory access
+ * crossing the boundary.
+ *
+ * @param：
+ * dst: The register pointer of the thread, the size is NX * NY.
+ * src: The data pointer of the current block.
+ * size: The current block needs to load size data continuously.
+ */
+
 template <typename T, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadData(T* dst,
                                          const T* __restrict__ src,
@@ -349,7 +378,6 @@ __device__ __forceinline__ void ReadData(T* dst,
  * NY: Each thread need to load NY rows, only NY = 1 was supported.
  * ArgsT: The Type if dst, ArgsT can be std::tuple<T> or std::tuple<Args>
  * Index: The index of data stored in dst.
- * threadIdx.x is used as the thread index. Currently only GPU was supported.
  * IsBoundary: Whether to make an out-of-bounds judgment on access to memory.
  * When the number of data processed by this block is less than
  * NX x NY x blockDim.x, boundary judgment is required to avoid memory access
@@ -743,7 +771,6 @@ __device__ __forceinline__ void Init(T* dst, T* init_data, int num) {
  * NX: The number of data continuously loaded by each thread.
  * NY: The number of data rows loaded by each thread, only NY = 1 was supported.
  * threadIdx.x is used as the thread index. Currently only GPU was supported.
- * Rank: The shape size of out. eg in[1, 35], out[32, 35] then shape size is 2.
  * IsBoundary: Indicates whether to perform block access storage out-of-bounds
  * judgment. When the number of data processed by the block is less than
  * NX x NY x blockDim.x, boundary judgment is required to avoid memory access
@@ -788,6 +815,67 @@ __device__ __forceinline__ void ReadDataBc(
   }
 }
 
+/**
+ * @brief Read 1D data from global memory to register with broadcast form.
+ * The difference from the above function is that it supports different data
+ * types of inputs.
+ *
+ * @template paraments
+ * T: The type of data stored in the global memory.
+ * NX: The number of data continuously loaded by each thread.
+ * NY: The number of data rows loaded by each thread, only NY = 1 was supported.
+ * ArgsT: The Type if dst, ArgsT can be std::tuple<T> or std::tuple<Args>
+ * Index: The index of data stored in dst.
+ * IsBoundary: Indicates whether to perform block access storage out-of-bounds
+ * judgment. When the number of data processed by the block is less than
+ * NX x NY x blockDim.x, boundary judgment is required to avoid memory access
+ * crossing the boundary.
+ *
+ * @param：
+ * dst: The register pointer of the thread, the size is NX * NY.
+ * src: The original input data pointer of kernel.
+ * block_offset: The data offset of this block, blockDim.x * blockIdx.x * NX;
+ * config: Calculation configuration of broadcast. It is used to calculate the
+ * coordinate mapping relationship between output data and input data.
+ * total_num_output: Total number of original output.
+ */
+
+template <typename T,
+          int NX,
+          int NY,
+          typename ArgsT,
+          int Index,
+          bool IsBoundary = false>
+__device__ __forceinline__ void ReadDataBc(
+    ArgsT* dst,
+    const T* __restrict__ src,
+    uint32_t block_offset,
+    const details::BroadcastConfig& config,
+    int total_num_output,
+    int read_lens = NX) {
+  uint32_t thread_offset = block_offset + threadIdx.x * NX;
+  uint32_t index_src = 0;
+
+#pragma unroll
+  for (uint32_t nx = 0; nx < NX; ++nx) {
+    uint32_t index_output = thread_offset + nx;
+    index_src = 0;
+    if (IsBoundary) {
+      if (index_output >= total_num_output) {
+        break;
+      }
+    }
+#pragma unroll
+    for (int i = 0; i < phi::DDim::kMaxRank; ++i) {
+      if (i >= config.rank) break;
+      auto fast_divmoder = config.divmoders[i].Divmod(index_output);
+      index_output = fast_divmoder.val[0];
+      index_src += fast_divmoder.val[1] * config.strides[i];
+    }
+    std::get<Index>(dst[nx]) = src[index_src];
+  }
+}
+
 /**
  * @brief Initialize register with data index.
  *

From 0cc4c1e07194891caf8dc27644fd23c48f27e9ed Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Fri, 7 Apr 2023 16:48:28 +0000
Subject: [PATCH 4/8] use same unroller with XPU

---
 paddle/phi/kernels/funcs/broadcast_function.h | 82 ++++---------------
 1 file changed, 14 insertions(+), 68 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index cbb6190bbe465..5b630dbc2251d 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -100,8 +100,6 @@ struct LoaderTypeClassifier {
   }
 };
 
-#ifndef PADDLE_WITH_XPU_KP
-
 // Common broadcast/elementwise Loader.
 template <int Index, int VecSize, bool IsBoundary, int LoadType>
 struct BroadcastDataLoader {
@@ -133,8 +131,9 @@ struct BroadcastDataLoader {
     }
   }
 };
-// FIXME: add BroadcastDataLoaders Partial specialization here
 
+#ifndef PADDLE_WITH_XPU_KP
+// FIXME: add BroadcastDataLoaders Partial specialization here
 #endif
 
 // static broadcast unroller
@@ -166,8 +165,7 @@ struct BcUnroller<Func, IsBoundary, LoadType, VecSize, End, End> {
   static HOSTDEVICE inline void step(Args &&...args) {}
 };
 
-template <typename InT,
-          typename OutT,
+template <typename OutT,
           typename Functor,
           int Arity,
           int NumOuts,
@@ -189,50 +187,20 @@ __device__ void VectorizedBroadcastKernelImpl(
   __simd__ ArgsT args[VecSize];
   __simd__ ConditionalT<OutT, NumOuts> result[VecSize];
 
-// FIXME: rm xpu InT
-#ifdef PADDLE_WITH_XPU_KP
-  __simd__ InT args[Arity][VecSize];
-#pragma unroll
-  for (int i = 0; i < Arity; ++i) {
-    kps::Init<InT, VecSize>(args[i], static_cast<InT>(1.0f), read_lens);
-    if (use_broadcast[i]) {
-      kps::ReadDataBc<InT, VecSize, 1, IsBoundary>(
-          args[i], ins[i], block_offset, configs[i], numel, read_lens);
-    } else {
-      kps::ReadData<InT, VecSize, 1, IsBoundary>(
-          args[i], ins[i] + block_offset, num, read_lens);
-    }
-  }
-#else
-
   BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
       ins, args, configs, use_broadcast, block_offset, num, numel);
-#endif
-
-  // constexpr bool kCallElementwiseAny =
-  //     phi::funcs::FunctionTraits<Functor>::has_pointer_args;
 
   SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
                                      VecSize,
                                      Functor,
                                      ArgsT,
                                      Arity>()(func, args, result, read_lens);
-
-  // FIXME: clean this
-  // phi::funcs::ElementwisePrimitiveCaller<InT,
-  //                                        ConditionalT<OutT, NumOuts>,
-  //                                        VecSize,
-  //                                        Functor,
-  //                                        Arity,
-  //                                        kCallElementwiseAny>()(
-  //     func, args, result, read_lens);
   phi::funcs::
       ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, NumOuts>()(
           outs, result, block_offset, num, read_lens);
 }
 
 template <typename Functor,
-          typename InT,
           typename OutT,
           int Arity,
           int NumOuts,
@@ -252,8 +220,7 @@ __global__ void VectorizedBroadcastKernel(
   int block_offset = BLOCK_ID_X * BLOCK_NUM_X * read_lens;
   int stride = BLOCK_NUM_X * GRID_NUM_X * read_lens;
   for (; block_offset < main_offset; block_offset += stride) {
-    VectorizedBroadcastKernelImpl<InT,
-                                  OutT,
+    VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
                                   NumOuts,
@@ -271,8 +238,7 @@ __global__ void VectorizedBroadcastKernel(
   }
   int num = numel - block_offset;
   if (num > 0) {
-    VectorizedBroadcastKernelImpl<InT,
-                                  OutT,
+    VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
                                   NumOuts,
@@ -291,8 +257,7 @@ __global__ void VectorizedBroadcastKernel(
 #else
   int block_offset = BLOCK_ID_X * BLOCK_NUM_X * VecSize;
   if (block_offset < main_offset) {
-    VectorizedBroadcastKernelImpl<InT,
-                                  OutT,
+    VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
                                   NumOuts,
@@ -308,8 +273,7 @@ __global__ void VectorizedBroadcastKernel(
                                             read_lens,
                                             func);
   } else {
-    VectorizedBroadcastKernelImpl<InT,
-                                  OutT,
+    VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
                                   NumOuts,
@@ -328,12 +292,7 @@ __global__ void VectorizedBroadcastKernel(
 #endif
 }
 
-template <typename InT,
-          typename OutT,
-          typename Functor,
-          int Arity,
-          int NumOuts,
-          int VecSize>
+template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
 void LaunchBroadcastKernel(
     const KPDevice &ctx,
     const std::vector<const DenseTensor *> &ins,
@@ -355,7 +314,7 @@ void LaunchBroadcastKernel(
   int main_offset = (numel / (read_lens * threads)) * read_lens * threads;
   int tail_tid = numel % (read_lens * threads);
 
-  VectorizedBroadcastKernel<Functor, InT, OutT, Arity, NumOuts, VecSize, false>
+  VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                        outs_data,
                                        loader_classifier.use_broadcast,
@@ -377,7 +336,6 @@ void LaunchBroadcastKernel(
 
   if (loader_classifier.all_elementwise) {
     VectorizedBroadcastKernel<Functor,
-                              InT,
                               OutT,
                               Arity,
                               NumOuts,
@@ -394,13 +352,7 @@ void LaunchBroadcastKernel(
                                          func);
   } else if (loader_classifier.broadcast_num > (Arity >> 1)) {
     constexpr BroadcastLoadType type_ = (Arity > 1) ? kBroadcast : kMixed;
-    VectorizedBroadcastKernel<Functor,
-                              InT,
-                              OutT,
-                              Arity,
-                              NumOuts,
-                              VecSize,
-                              type_>
+    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
         <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                          outs_data,
                                          loader_classifier.use_broadcast,
@@ -411,13 +363,7 @@ void LaunchBroadcastKernel(
                                          VecSize,
                                          func);
   } else {
-    VectorizedBroadcastKernel<Functor,
-                              InT,
-                              OutT,
-                              Arity,
-                              NumOuts,
-                              VecSize,
-                              kMixed>
+    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, kMixed>
         <<<blocks, threads, 0, stream>>>(loader_classifier.ins_data,
                                          outs_data,
                                          loader_classifier.use_broadcast,
@@ -944,17 +890,17 @@ void BroadcastKernelForDifferentVecSize(
 #endif
   switch (loader_classifier.vec_size) {
     case VecSizeL: {
-      LaunchBroadcastKernel<InT, OutT, Functor, kArity, NumOuts, VecSizeL>(
+      LaunchBroadcastKernel<OutT, Functor, kArity, NumOuts, VecSizeL>(
           ctx, ins, outs, func, configs, loader_classifier);
       break;
     }
     case VecSizeM: {
-      LaunchBroadcastKernel<InT, OutT, Functor, kArity, NumOuts, VecSizeM>(
+      LaunchBroadcastKernel<OutT, Functor, kArity, NumOuts, VecSizeM>(
           ctx, ins, outs, func, configs, loader_classifier);
       break;
     }
     case VecSizeS: {
-      LaunchBroadcastKernel<InT, OutT, Functor, kArity, NumOuts, VecSizeS>(
+      LaunchBroadcastKernel<OutT, Functor, kArity, NumOuts, VecSizeS>(
           ctx, ins, outs, func, configs, loader_classifier);
       break;
     }

From b26a77707362aeccceb0af5da0422d12cba177ad Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Tue, 11 Apr 2023 06:11:36 +0000
Subject: [PATCH 5/8] BroadcastWithInt64Index

---
 paddle/phi/kernels/funcs/broadcast_function.h | 203 +++++-------------
 .../kernels/primitive/datamover_primitives.h  |   6 +-
 2 files changed, 61 insertions(+), 148 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 5b630dbc2251d..4dcc27117207e 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -396,94 +396,49 @@ HOSTDEVICE static int64_t ConvertSrcIdxToDstIdx(
   return dst_idx;
 }
 
-template <typename T, int VecSize, bool IsBoundary>
-HOSTDEVICE static void ReadVecDataWithInt64Index(
-    const T *in,
-    int64_t idx,
-    bool need_broadcast,
-    const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &src_strides,
-    const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &dst_strides,
-    int rank,
-    int n,
-    phi::AlignedVector<T, VecSize> *out) {
-  if (IsBoundary) {
-    for (int i = 0; i < n; ++i) {
-      (*out)[i] =
-          in[ConvertSrcIdxToDstIdx(idx + i, src_strides, dst_strides, rank)];
-    }
-  } else {
-    if (!need_broadcast) {
-      phi::Load<T, VecSize>(in + idx, out);
-    } else {
+template <int N>
+struct MaxWithOne {
+  static constexpr auto kValue = (N >= 1 ? N : 1);
+};
+
+template <int Index, int VecSize>
+struct ReadVecDataWithInt64Index {
+  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
+  static __device__ __forceinline__ void Apply(
+      const Array1 &in,
+      ArgsT *args,
+      int64_t idx,
+      const Array2 &need_broadcast,
+      const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &src_strides,
+      const Array3 &dst_strides,
+      int rank,
+      bool is_boundary) {
+    using Type = std::tuple_element_t<Index, ArgsT>;
+    if (is_boundary) {
 #pragma unroll
       for (int i = 0; i < VecSize; ++i) {
-        (*out)[i] =
-            in[ConvertSrcIdxToDstIdx(idx + i, src_strides, dst_strides, rank)];
+        std::get<Index>(args[i]) = in[Index][ConvertSrcIdxToDstIdx(
+            idx + i, src_strides, dst_strides[Index], rank)];
+      }
+    } else {
+      if (!need_broadcast[Index]) {
+        kps::ReadData<Type, VecSize, 1, ArgsT, Index, false>(
+            args, reinterpret_cast<const _ptr_ Type *>(in[Index]) + idx, 1);
+      } else {
+#pragma unroll
+        for (int i = 0; i < VecSize; ++i) {
+          std::get<Index>(args[i]) = in[Index][ConvertSrcIdxToDstIdx(
+              idx + i, src_strides, dst_strides[Index], rank)];
+        }
       }
     }
   }
-}
-
-template <typename InT,
-          typename OutT,
-          typename Functor,
-          int VecSize,
-          int NumIns>
-struct ApplyFunctorWithInt64IndexHelper {
-  HOSTDEVICE static OutT Run(const phi::AlignedVector<InT, VecSize> *ins_vec,
-                             Functor functor,
-                             int i);
 };
 
-template <typename InT, typename OutT, typename Functor, int VecSize>
-struct ApplyFunctorWithInt64IndexHelper<InT, OutT, Functor, VecSize, 0> {
-  HOSTDEVICE static OutT Run(const phi::AlignedVector<InT, VecSize> *ins_vec,
-                             Functor functor,
-                             int i) {
-    return static_cast<OutT>(functor());
-  }
-};
-
-template <typename InT, typename OutT, typename Functor, int VecSize>
-struct ApplyFunctorWithInt64IndexHelper<InT, OutT, Functor, VecSize, 1> {
-  HOSTDEVICE static OutT Run(const phi::AlignedVector<InT, VecSize> *ins_vec,
-                             Functor functor,
-                             int i) {
-    return static_cast<OutT>(functor(ins_vec[0][i]));
-  }
-};
-
-template <typename InT, typename OutT, typename Functor, int VecSize>
-struct ApplyFunctorWithInt64IndexHelper<InT, OutT, Functor, VecSize, 2> {
-  HOSTDEVICE static OutT Run(const phi::AlignedVector<InT, VecSize> *ins_vec,
-                             Functor functor,
-                             int i) {
-    return static_cast<OutT>(functor(ins_vec[0][i], ins_vec[1][i]));
-  }
-};
-
-template <typename InT, typename OutT, typename Functor, int VecSize>
-struct ApplyFunctorWithInt64IndexHelper<InT, OutT, Functor, VecSize, 3> {
-  HOSTDEVICE static OutT Run(const phi::AlignedVector<InT, VecSize> *ins_vec,
-                             Functor functor,
-                             int i) {
-    return static_cast<OutT>(
-        functor(ins_vec[0][i], ins_vec[1][i], ins_vec[2][i]));
-  }
-};
-
-template <int N>
-struct MaxWithOne {
-  static constexpr auto kValue = (N >= 1 ? N : 1);
-};
-
-template <typename InT,
-          typename OutT,
-          typename Functor,
-          int VecSize,
-          int NumIns>
+template <typename OutT, typename Functor, int VecSize, int NumIns>
 __global__ void BroadcastKernelWithInt64Index(
-    phi::Array<const InT *, MaxWithOne<NumIns>::kValue> ins,
+    const phi::Array<const _ptr_ char *__restrict__, MaxWithOne<NumIns>::kValue>
+        &ins,
     OutT *out,
     phi::Array<phi::Array<int64_t, phi::DDim::kMaxRank + 1>,
                MaxWithOne<NumIns>::kValue> ins_strides,
@@ -497,70 +452,34 @@ __global__ void BroadcastKernelWithInt64Index(
   int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x * VecSize;
   int64_t limit = numel - VecSize;
 
-  phi::Array<phi::AlignedVector<InT, VecSize>, MaxWithOne<NumIns>::kValue>
-      ins_vec;
+  using Traits = phi::funcs::FunctionTraits<Functor>;
+  using ArgsT = typename Traits::ArgsTuple;
+
+  ArgsT args[VecSize];
   phi::AlignedVector<OutT, VecSize> out_vec;
   for (; idx <= limit; idx += stride) {
-#pragma unroll
-    for (int i = 0; i < NumIns; ++i) {
-      ReadVecDataWithInt64Index<InT, VecSize, false>(ins[i],
-                                                     idx,
-                                                     need_broadcasts[i],
-                                                     out_strides,
-                                                     ins_strides[i],
-                                                     rank,
-                                                     VecSize,
-                                                     &ins_vec[i]);
-    }
+    Unroller<ReadVecDataWithInt64Index, VecSize, NumIns>::step(
+        ins, args, idx, need_broadcasts, out_strides, ins_strides, rank, false);
 
 #pragma unroll
     for (int i = 0; i < VecSize; ++i) {
-      out_vec[i] = ApplyFunctorWithInt64IndexHelper<InT,
-                                                    OutT,
-                                                    Functor,
-                                                    VecSize,
-                                                    NumIns>::Run(ins_vec.Get(),
-                                                                 functor,
-                                                                 i);
+      out_vec[i] = static_cast<OutT>(Apply(functor, args[i]));
     }
-
     phi::Store<OutT, VecSize>(out_vec, out + idx);
   }
 
   if (idx < numel) {
     int remain = numel - idx;  // remain is always less than VecSize, therefore
                                // `int` is enough here
-#pragma unroll
-    for (int i = 0; i < NumIns; ++i) {
-      ReadVecDataWithInt64Index<InT, VecSize, true>(ins[i],
-                                                    idx,
-                                                    need_broadcasts[i],
-                                                    out_strides,
-                                                    ins_strides[i],
-                                                    rank,
-                                                    remain,
-                                                    &ins_vec[i]);
-    }
-
+    Unroller<ReadVecDataWithInt64Index, VecSize, NumIns>::step(
+        ins, args, idx, need_broadcasts, out_strides, ins_strides, rank, true);
     for (int i = 0; i < remain; ++i) {
-      out[idx + i] =
-          ApplyFunctorWithInt64IndexHelper<InT,
-                                           OutT,
-                                           Functor,
-                                           VecSize,
-                                           NumIns>::Run(ins_vec.Get(),
-                                                        functor,
-                                                        i);
+      out_vec[idx + i] = static_cast<OutT>(Apply(functor, args[i]));
     }
   }
 }
 
-template <typename InT,
-          typename OutT,
-          typename Functor,
-          int Arity,
-          int NumOuts,
-          int VecSize>
+template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
 struct LaunchBroadcastKernelWithInt64IndexHelper {
   static void Run(const KPDevice &ctx,
                   const std::vector<const DenseTensor *> &ins,
@@ -572,9 +491,8 @@ struct LaunchBroadcastKernelWithInt64IndexHelper {
   }
 };
 
-template <typename InT, typename OutT, typename Functor, int Arity, int VecSize>
-struct LaunchBroadcastKernelWithInt64IndexHelper<InT,
-                                                 OutT,
+template <typename OutT, typename Functor, int Arity, int VecSize>
+struct LaunchBroadcastKernelWithInt64IndexHelper<OutT,
                                                  Functor,
                                                  Arity,
                                                  /*NumOuts=*/1,
@@ -584,10 +502,9 @@ struct LaunchBroadcastKernelWithInt64IndexHelper<InT,
                   std::vector<DenseTensor *> *outs,
                   int axis,
                   Functor functor) {
-    phi::Array<const InT *, MaxWithOne<Arity>::kValue> ins_ptrs;
-    for (int i = 0; i < Arity; ++i) {
-      ins_ptrs[i] = ins[i]->data<InT>();
-    }
+    phi::Array<const _ptr_ char *__restrict__, MaxWithOne<Arity>::kValue>
+        ins_ptrs;
+    UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, &ins_ptrs);
     auto *out_tensor = (*outs)[0];
     auto *out_ptr = ctx.Alloc<OutT>(out_tensor);
 
@@ -659,7 +576,7 @@ struct LaunchBroadcastKernelWithInt64IndexHelper<InT,
     auto gpu_config =
         phi::backends::gpu::GetGpuLaunchConfig1D(ctx, numel, VecSize);
 
-    BroadcastKernelWithInt64Index<InT, OutT, Functor, VecSize, Arity>
+    BroadcastKernelWithInt64Index<OutT, Functor, VecSize, Arity>
         <<<gpu_config.block_per_grid,
            gpu_config.thread_per_block,
            0,
@@ -768,9 +685,8 @@ struct LaunchBroadcastKernelWithInt64IndexHelper<InT,
 };
 #endif
 
-// FIXME: delete ElementwiseType InT
+// FIXME: delete ElementwiseType
 template <ElementwiseType ET,
-          typename InT,
           typename OutT,
           typename Functor,
           int kArity,
@@ -791,8 +707,7 @@ void BroadcastKernelForDifferentVecSize(
         LoaderTypeClassifier<OutT, kArity, Functor>(ins, outs);
     switch (loader_classifier.vec_size) {
       case VecSizeL: {
-        LaunchBroadcastKernelWithInt64IndexHelper<InT,
-                                                  OutT,
+        LaunchBroadcastKernelWithInt64IndexHelper<OutT,
                                                   Functor,
                                                   kArity,
                                                   NumOuts,
@@ -804,8 +719,7 @@ void BroadcastKernelForDifferentVecSize(
         break;
       }
       case VecSizeM: {
-        LaunchBroadcastKernelWithInt64IndexHelper<InT,
-                                                  OutT,
+        LaunchBroadcastKernelWithInt64IndexHelper<OutT,
                                                   Functor,
                                                   kArity,
                                                   NumOuts,
@@ -817,8 +731,7 @@ void BroadcastKernelForDifferentVecSize(
         break;
       }
       case VecSizeS: {
-        LaunchBroadcastKernelWithInt64IndexHelper<InT,
-                                                  OutT,
+        LaunchBroadcastKernelWithInt64IndexHelper<OutT,
                                                   Functor,
                                                   kArity,
                                                   NumOuts,
@@ -912,7 +825,7 @@ void BroadcastKernelForDifferentVecSize(
   }
 }
 
-// FIXME: delete (ElementwiseType ET & typename InT)
+// FIXME: delete (ElementwiseType ET)
 // default: axis = -1
 template <ElementwiseType ET,
           typename InT,
@@ -961,7 +874,7 @@ void BroadcastKernel(const KPDevice &ctx,
     max_rank = std::max(max_rank, (*outs)[0]->dims().size());
   }
   axis = axis == -1 ? max_rank - min_rank : axis;
-  BroadcastKernelForDifferentVecSize<ET, InT, OutT, Functor, kArity, NumOuts>(
+  BroadcastKernelForDifferentVecSize<ET, OutT, Functor, kArity, NumOuts>(
       ctx, ins, outs, axis, func);
 }
 
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index ca9de686d817f..2a3579d99cfe6 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -376,7 +376,7 @@ __device__ __forceinline__ void ReadData(T* dst,
  * T: The type of data.
  * NX: Each thread load NX data from global memory continuously.
  * NY: Each thread need to load NY rows, only NY = 1 was supported.
- * ArgsT: The Type if dst, ArgsT can be std::tuple<T> or std::tuple<Args>
+ * ArgsT: The Type of dst, ArgsT can be std::tuple<T> or std::tuple<Args>
  * Index: The index of data stored in dst.
  * IsBoundary: Whether to make an out-of-bounds judgment on access to memory.
  * When the number of data processed by this block is less than
@@ -397,7 +397,7 @@ template <typename T,
 __device__ __forceinline__ void ReadData(ArgsT* dst,
                                          const T* __restrict__ src,
                                          int num,
-                                         int read_lens) {
+                                         int read_lens = 0) {
   if (IsBoundary) {  // blockDim.x * NX > num
     int thread_offset = threadIdx.x * NX;
 #pragma unroll
@@ -824,7 +824,7 @@ __device__ __forceinline__ void ReadDataBc(
  * T: The type of data stored in the global memory.
  * NX: The number of data continuously loaded by each thread.
  * NY: The number of data rows loaded by each thread, only NY = 1 was supported.
- * ArgsT: The Type if dst, ArgsT can be std::tuple<T> or std::tuple<Args>
+ * ArgsT: The Type of dst, ArgsT can be std::tuple<T> or std::tuple<Args>
  * Index: The index of data stored in dst.
  * IsBoundary: Indicates whether to perform block access storage out-of-bounds
  * judgment. When the number of data processed by the block is less than

From df71250431762d1431f50b99e55b94ca7f09d813 Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Wed, 12 Apr 2023 04:24:03 +0000
Subject: [PATCH 6/8] BroadcastDataLoader template partial specialization

---
 paddle/phi/kernels/funcs/broadcast_function.h | 102 ++++++++++++++++--
 paddle/phi/kernels/funcs/elementwise_base.h   |   1 -
 2 files changed, 96 insertions(+), 7 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 4dcc27117207e..719124530a307 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -132,8 +132,102 @@ struct BroadcastDataLoader {
   }
 };
 
+/* BroadcastDataLoaders Partial specialization */
 #ifndef PADDLE_WITH_XPU_KP
-// FIXME: add BroadcastDataLoaders Partial specialization here
+// Scalar elementwise Loader with consideration of IsBoundary.
+template <int Index, int VecSize>
+struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
+  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
+  static __device__ __forceinline__ void Apply(const Array1 &ins,
+                                               ArgsT *args,
+                                               const Array2 &configs,
+                                               const Array3 &use_broadcast,
+                                               const int block_offset,
+                                               const int num,
+                                               const uint32_t numel) {
+    using Type = std::tuple_element_t<Index, ArgsT>;
+    int thread_offset = threadIdx.x * VecSize + block_offset;
+#pragma unroll
+    for (int idx = 0; idx < VecSize; ++idx) {
+      std::get<Index>(args[idx]) = static_cast<Type>(1);
+      int index = thread_offset + idx;
+      if (index < numel) {
+        std::get<Index>(args[idx]) =
+            reinterpret_cast<const _ptr_ Type *>(ins[Index])[index];
+      }
+    }
+  }
+};
+
+// Vectorized elementwise Loader without consideration of IsBoundary.
+template <int Index, int VecSize>
+struct BroadcastDataLoader<Index, VecSize, false, kElementwise> {
+  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
+  static __device__ __forceinline__ void Apply(const Array1 &ins,
+                                               ArgsT *args,
+                                               const Array2 &configs,
+                                               const Array3 &use_broadcast,
+                                               const int block_offset,
+                                               const int num,
+                                               const uint32_t numel) {
+    using Type = std::tuple_element_t<Index, ArgsT>;
+    using VecType = phi::kps::details::VectorType<Type, VecSize>;
+    VecType vec_temp;
+
+    int thread_offset = threadIdx.x + blockIdx.x * blockDim.x;
+    const VecType *__restrict__ vec_input =
+        reinterpret_cast<const VecType *__restrict__>(ins[Index]);
+    vec_temp = vec_input[thread_offset];
+#pragma unroll
+    for (int idx = 0; idx < VecSize; ++idx) {
+      std::get<Index>(args[idx]) = vec_temp.val[idx];
+    }
+  }
+};
+
+// Common broadcast data loader.
+template <int Index, int VecSize, bool IsBoundary>
+struct BroadcastDataLoader<Index, VecSize, IsBoundary, kBroadcast> {
+  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
+  static __device__ __forceinline__ void Apply(const Array1 &ins,
+                                               ArgsT *args,
+                                               const Array2 &configs,
+                                               const Array3 &use_broadcast,
+                                               const int block_offset,
+                                               const int num,
+                                               const uint32_t numel) {
+    using Type = std::tuple_element_t<Index, ArgsT>;
+    uint32_t index_bc[VecSize];
+#pragma unroll
+    for (int k = 0; k < VecSize; ++k) {
+      index_bc[k] = 0;
+      std::get<Index>(args[k]) = static_cast<Type>(1);
+    }
+
+    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
+#pragma unroll
+    for (int k = 0; k < VecSize; ++k) {
+      uint32_t idx = thread_offset + k;
+      if (IsBoundary && idx == numel) {
+        break;
+      }
+#pragma unroll
+      for (int i = 0; i < phi::DDim::kMaxRank; ++i) {
+        if (i == configs[0].rank) break;
+        auto fast_divmoder = configs[0].divmoders[i].Divmod(idx);
+        idx = fast_divmoder.val[0];
+        index_bc[k] += fast_divmoder.val[1] * configs[Index].strides[i];
+      }
+    }
+
+#pragma unroll
+    for (int k = 0; k < VecSize; ++k) {
+      std::get<Index>(args[k]) =
+          reinterpret_cast<const _ptr_ Type *>(ins[Index])[index_bc[k]];
+    }
+  }
+};
+
 #endif
 
 // static broadcast unroller
@@ -685,7 +779,6 @@ struct LaunchBroadcastKernelWithInt64IndexHelper<OutT,
 };
 #endif
 
-// FIXME: delete ElementwiseType
 template <ElementwiseType ET,
           typename OutT,
           typename Functor,
@@ -825,8 +918,6 @@ void BroadcastKernelForDifferentVecSize(
   }
 }
 
-// FIXME: delete (ElementwiseType ET)
-// default: axis = -1
 template <ElementwiseType ET,
           typename InT,
           typename OutT,
@@ -839,7 +930,6 @@ void BroadcastKernel(const KPDevice &ctx,
                      Functor func) {
   // When there are multiple inputs, the outputs's rank should be equal the
   // maximum rank of all inputs.
-  // FIXME: delete ET ？
   using Traits = phi::funcs::FunctionTraits<Functor>;
   const int kArity = Traits::arity;
   PADDLE_ENFORCE_EQ(
@@ -888,7 +978,7 @@ void ElementwiseCompute(const GPUContext &dev_ctx,
   std::vector<const DenseTensor *> ins = {&x, &y};
   std::vector<DenseTensor *> outs = {z};
   dev_ctx.template Alloc<OutType>(z);
-  // FIXME: delete ElementwiseType
+
   BroadcastKernel<ElementwiseType::kBinary, T, OutType, Functor, 1>(
       dev_ctx, ins, &outs, axis, func);
 }
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
index 14f0cecf42d75..a2cb0f8b2af8c 100644
--- a/paddle/phi/kernels/funcs/elementwise_base.h
+++ b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -35,7 +35,6 @@ namespace kps = phi::kps;
 
 namespace phi {
 
-// FIXME: delete this enum
 enum ElementwiseType { kUnary = 1, kBinary = 2, kTernary = 3 };
 /* Packing scalar type T(float, int etc.) into Array<T, NumOuts> type
    for supporting multiple-output feature in elementwise system.*/

From a6119ccb90db64cee4bb590c8b336cf7ec28724e Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Wed, 12 Apr 2023 11:58:34 +0000
Subject: [PATCH 7/8] fix compile errs in ROCms

---
 paddle/phi/kernels/funcs/broadcast_function.h | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 719124530a307..b4125ab555079 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -122,7 +122,11 @@ struct BroadcastDataLoader {
           configs[Index],
           numel,
           VecSize);
-    } else {
+    }
+    // NOTE: If use if...else... with condition `use_broadcast[Index]` here,
+    // there will be some errs with clang12 while compiling in ROCm.
+    // When the compiler is upgraded, if...else... may be used.
+    if (!use_broadcast[Index]) {
       kps::ReadData<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
           args,
           reinterpret_cast<const _ptr_ Type *>(ins[Index]) + block_offset,

From a5c47c5166e5c73621cef965ab8e501c4d0cce4c Mon Sep 17 00:00:00 2001
From: zhangbopd <1299246947@qq.com>
Date: Thu, 27 Apr 2023 07:43:28 +0000
Subject: [PATCH 8/8] PR comment

---
 paddle/phi/kernels/funcs/dropout_impl.cu.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/paddle/phi/kernels/funcs/dropout_impl.cu.h b/paddle/phi/kernels/funcs/dropout_impl.cu.h
index 8b12535bb0e00..2ae9d3c02f63d 100644
--- a/paddle/phi/kernels/funcs/dropout_impl.cu.h
+++ b/paddle/phi/kernels/funcs/dropout_impl.cu.h
@@ -311,7 +311,7 @@ void DropoutFwGPUKernelDriver(
     auto* mask_data = mask->data<uint8_t>();
     size_t size = phi::product(mask->dims());
 
-    if (abs(dropout_prob / 1.0f - 1) < 5.96e-08) {
+    if (dropout_prob == 1.0f) {
 #ifdef PADDLE_WITH_HIP
       PADDLE_ENFORCE_GPU_SUCCESS(
           hipMemsetAsync(y_data, 0, x_numel * sizeof(T), stream));
@@ -459,7 +459,7 @@ void DropoutGradGPUKernelDriver(const phi::GPUContext& dev_ctx,
     // y = factor * x
     ScaleByDropoutFactor<T, MT>(dev_ctx, grad_y, grad_x, factor);
   } else {
-    if (upscale_in_train && abs(dropout_prob / 1.0f - 1) < 5.96e-08) {
+    if (upscale_in_train && dropout_prob == 1.0f) {
 #ifdef PADDLE_WITH_HIP
       hipMemset(grad_x->data<T>(), 0, grad_x->numel() * sizeof(T));
 #else