 }
 
 template <typename T>
-__global__ void GPUSigmoidForward(const T *x_data, const T *label_data,
-                                  const int ignore_index, const int limit,
-                                  T *out_data, T *counts) {
-  CUDA_KERNEL_LOOP(i, limit) {
-    T x = x_data[i];
-    T label = label_data[i];
-    T eps = static_cast<T>(1e-5);
-    T diff = label - static_cast<T>(ignore_index);
+struct NonzeroFunctor {
+  HOSTDEVICE explicit inline NonzeroFunctor() {}
+  HOSTDEVICE inline T operator()(const T x) const {
+    return static_cast<T>(static_cast<double>(x) != 0);
+  }
+};
+
+template <typename T>
+struct SigmoidFwdFunctor {
+  T ignore_index_;
+  T eps = static_cast<T>(1e-5);
+
+  HOSTDEVICE inline SigmoidFwdFunctor(const T ignore_index)
+      : ignore_index_(ignore_index) {}
+
+  HOSTDEVICE inline phi::Array<T, 2> operator()(const T x, const T label) {
+    T counts;
+    T out_data;
+
+    T diff = label - static_cast<T>(ignore_index_);
     if ((diff > -eps) && (diff < eps)) {
-      out_data[i] = static_cast<T>(0.);
-      counts[i] = 0;
+      out_data = static_cast<T>(0.);
+      counts = 0;
     } else {
       T term1 = (x > 0) ? x : 0;
       T term2 = x * label;
       T term3 = real_log(static_cast<T>(1) + real_exp(static_cast<T>(-abs(x))));
-      out_data[i] = term1 - term2 + term3;
-      counts[i] = 1;
-    }
-  }
-}
-
-template <typename T, int BlockDim>
-__global__ void Sum(const T *counts, int num, const T eps, T *sum) {
-  typedef cub::BlockReduce<double, BlockDim> BlockReduce;
-  __shared__ typename BlockReduce::TempStorage temp_storage;
-  T in = 0;
-  for (int i = threadIdx.x; i < num; i += BlockDim) {
-    in += counts[i];
-  }
-  __syncthreads();
-  auto out =
-      BlockReduce(temp_storage).Reduce(static_cast<double>(in), cub::Sum());
-  __syncthreads();
-  if (threadIdx.x == 0) {
-    T a = out > eps ? out : eps;
-    sum[0] = a;
-  }
-}
-
-template <typename T>
-__global__ void Div(T *loss, const int num, const T *norm) {
-  CUDA_KERNEL_LOOP(i, num) { loss[i] /= norm[0]; }
-}
-
-template <typename T>
-__global__ void GPUSigmoidBackward(const T *x_data, const T *label_data,
-                                   const int ignore_index, const T *dout_data,
-                                   const int limit, T *dx_data, T *counts) {
-  CUDA_KERNEL_LOOP(i, limit) {
-    T x = x_data[i];
-    T label = label_data[i];
-    T dout = dout_data[i];
-    T eps = static_cast<T>(1e-5);
-    T diff = label - static_cast<T>(ignore_index);
+
+      out_data = term1 - term2 + term3;
+      counts = 1;
+    }
+    phi::Array<T, 2> outs;
+
+    outs[0] = out_data;
+    outs[1] = counts;
+    return outs;
+  }
+};
+
+template <typename T>
+struct SigmoidBwdFunctor {
+  T ignore_index_;
+  T eps = static_cast<T>(1e-5);
+
+  HOSTDEVICE inline SigmoidBwdFunctor(const T ignore_index)
+      : ignore_index_(ignore_index) {}
+
+  HOSTDEVICE inline phi::Array<T, 2> operator()(const T x, const T label,
+                                                const T dout) {
+    T counts;
+    T dx_data;
+
+    T diff = label - static_cast<T>(ignore_index_);
     if ((diff > -eps) && (diff < eps)) {
-      dx_data[i] = static_cast<T>(0.);
-      counts[i] = 0;
+      dx_data = static_cast<T>(0.);
+      counts = 0;
     } else {
       T simoid_x = static_cast<T>(1) / (static_cast<T>(1) + real_exp(-x));
       T diff = simoid_x - label;
-      dx_data[i] = dout * diff;
-      counts[i] = 1;
-    }
-  }
-}
+      dx_data = dout * diff;
+      counts = 1;
+    }
+    phi::Array<T, 2> outs;
+
+    outs[0] = dx_data;
+    outs[1] = counts;
+    return outs;
+  }
+};
+
+template <typename T>
+struct DivFunctor {
+  const T norm_;
+  HOSTDEVICE inline DivFunctor(const T norm) : norm_(norm) {}
+
+  HOSTDEVICE inline T operator()(T loss) {
+    loss /= norm_;
+    return loss;
+  }
+};
 
 // Out = max(X, 0) - X * Labels + log(1 + exp(-abs(X)))
 template <typename DeviceContext, typename T>

    bool normalize = context.Attr<bool>("normalize");

    // Temporary memory
    auto cnt_ptr = memory::Alloc(dev_ctx, Labels->numel() * sizeof(T));
    T *counts = reinterpret_cast<T *>(cnt_ptr->ptr());

    int limit = Out->numel();
    int blocks = NumBlocks(limit);
    int threads = kNumCUDAThreads;
    GPUSigmoidForward<T><<<blocks, threads, 0, dev_ctx.stream()>>>(
        X->data<T>(), Labels->data<T>(), ignore_index, limit, out_data, counts);
    if (normalize) {
      auto norm_ptr = memory::Alloc(dev_ctx, sizeof(T));
      T *norm = reinterpret_cast<T *>(norm_ptr->ptr());
      Sum<T, kNumCUDAThreads><<<1, kNumCUDAThreads, 0, dev_ctx.stream()>>>(
          counts, limit, static_cast<T>(1e-5), norm);
      Div<T><<<blocks, threads, 0, dev_ctx.stream()>>>(out_data, limit, norm);
    }
  }
};


    auto &dev_ctx = context.cuda_device_context();
    // Temporary memory
    Tensor *counts_tensor = new Tensor();
    counts_tensor->mutable_data<T>(context.GetPlace(),
                                   Labels->numel() * sizeof(T));
    counts_tensor->Resize(dX->dims());

    int limit = dX->numel();
    int blocks = NumBlocks(limit);
    int threads = kNumCUDAThreads;
    std::vector<const framework::Tensor *> ins = {X, Labels, dOut};
    std::vector<framework::Tensor *> outs = {dX, counts_tensor};
    auto functor = SigmoidBwdFunctor<T>(ignore_index);
    constexpr int Size = 2;
    phi::funcs::ElementwiseKernel<T, decltype(functor), Size>(dev_ctx, ins,
                                                              &outs, functor);
    bool normalize = context.Attr<bool>("normalize");
    if (normalize) {
      T *counts = counts_tensor->mutable_data<T>(context.GetPlace());
      Tensor *norm_tensor = new Tensor();
      norm_tensor->mutable_data<T>(context.GetPlace(), sizeof(T));
      auto dims = phi::vectorize(counts_tensor->dims());
      std::vector<int> reduce_dim = {};
      for (int i = 0; i < dims.size(); i++) {
        reduce_dim.push_back(i);
      }

      TensorReduceImpl<T, T, kps::AddFunctor, NonzeroFunctor<T>>(
          context.cuda_device_context(), *counts_tensor, norm_tensor,
          NonzeroFunctor<T>(), reduce_dim, dev_ctx.stream());
      T *norm = norm_tensor->mutable_data<T>(context.GetPlace());
      auto norm_cpu_mem = memory::Alloc(platform::CPUPlace(), sizeof(T));
      T *norm_cpu_ptr = reinterpret_cast<T *>(norm_cpu_mem->ptr());
      memory::Copy(platform::CPUPlace(), norm_cpu_ptr, dev_ctx.GetPlace(), norm,
                   sizeof(T), dev_ctx.stream());
      auto eps = static_cast<T>(1e-5);
      *norm_cpu_ptr = *norm_cpu_ptr > eps ? *norm_cpu_ptr : eps;

      std::vector<const framework::Tensor *> div_ins = {dX};
      std::vector<framework::Tensor *> div_outs = {dX};
      auto div_functor = DivFunctor<T>(*norm_cpu_ptr);
      phi::funcs::ElementwiseKernel<T>(dev_ctx, div_ins, &div_outs,
                                       div_functor);
      delete norm_tensor;
    }
  }
};
