/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/common/transform.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/common_shape.h"
#include "paddle/phi/kernels/funcs/elementwise_utils.h"
#include "paddle/phi/kernels/funcs/math_function.h"

#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/kernels/funcs/aligned_vector.h"
#include "paddle/phi/kernels/funcs/function_traits.h"
#include "paddle/phi/kernels/primitive/kernel_primitives.h"

#define HOSTDEVICE __host__ __device__

#endif

namespace phi {

/* 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>
using ConditionalT = typename std::conditional_t<Num == 1, T, Array<T, Num>>;

namespace funcs {
using DDim = DDim;

template <typename T, typename DeviceContext>
class RowwiseTransformIterator;

template <typename T, typename DeviceContext>
class MidWiseTransformIterator;

// NOTE(dzhwinter): ptrdiff_t in iterator is deprecated in c++17
template <typename T>
class RowwiseTransformIterator<T, CPUContext> {
 public:
  using iterator_category = std::random_access_iterator_tag;
  using value_type = T;
  using difference_type = std::ptrdiff_t;
  using pointer = T *;
  using reference = T &;

  RowwiseTransformIterator(const T *ptr, int n) : ptr_(ptr), i_(0), n_(n) {}

  RowwiseTransformIterator<T, CPUContext> &operator++() {
    ++i_;
    if (UNLIKELY(i_ == n_)) {
      i_ = 0;
    }
    return *this;
  }

  RowwiseTransformIterator<T, CPUContext> &operator+(int n) {
    while (n-- > 0) {
      ++i_;
      if (UNLIKELY(i_ == n_)) {
        i_ = 0;
      }
    }

    return *this;
  }

  bool operator==(const RowwiseTransformIterator<T, CPUContext> &rhs) const {
    return (ptr_ + i_) == &(*rhs);
  }

  bool operator!=(const RowwiseTransformIterator<T, CPUContext> &rhs) const {
    return (ptr_ + i_) != &(*rhs);
  }

  const T &operator*() { return ptr_[i_]; }

 private:
  const T *ptr_;
  int i_;
  int64_t n_;
};

template <typename T>
class MidWiseTransformIterator<T, CPUContext> {
 public:
  using iterator_category = std::random_access_iterator_tag;
  using value_type = T;
  using difference_type = std::ptrdiff_t;
  using pointer = T *;
  using reference = T &;

  MidWiseTransformIterator(const T *ptr, int n, int post)
      : ptr_(ptr), i_(0), j_(0), n_(n), post_(post) {}

  MidWiseTransformIterator<T, CPUContext> &operator++() {
    ++j_;
    if (UNLIKELY(j_ == post_)) {
      ++i_;
      j_ = 0;
      if (UNLIKELY(i_ == n_)) {
        i_ = 0;
      }
    }
    return *this;
  }

  MidWiseTransformIterator<T, CPUContext> &operator+(int n) {
    while (n-- > 0) {
      ++j_;
      if (UNLIKELY(j_ == post_)) {
        ++i_;
        j_ = 0;
        if (UNLIKELY(i_ == n_)) {
          i_ = 0;
        }
      }
    }
    return *this;
  }

  bool operator==(const MidWiseTransformIterator<T, CPUContext> &rhs) const {
    return (ptr_ + i_) == &(*rhs);
  }

  bool operator!=(const MidWiseTransformIterator<T, CPUContext> &rhs) const {
    return (ptr_ + i_) != &(*rhs);
  }

  const T &operator*() { return ptr_[i_]; }

 private:
  const T *ptr_;
  int64_t i_;
  int64_t j_;
  int64_t n_;
  int64_t post_;
};

#if defined(__NVCC__) || defined(__HIPCC__)
template <typename T>
class RowwiseTransformIterator<T, GPUContext>
    : public thrust::iterator_adaptor<RowwiseTransformIterator<T, GPUContext>,
                                      const T *> {
 public:
  typedef thrust::iterator_adaptor<RowwiseTransformIterator<T, GPUContext>,
                                   const T *>
      super_t;
  HOSTDEVICE RowwiseTransformIterator(const T *x, int n)
      : super_t(x), begin_(x), n_(n) {}
  friend class thrust::iterator_core_access;

 private:
  unsigned int n_;
  const T *begin_;
  HOSTDEVICE typename super_t::reference dereference() const {
    return *(begin_ + (this->base() - begin_) % n_);
  }
};

template <typename T>
class MidWiseTransformIterator<T, GPUContext>
    : public thrust::iterator_adaptor<MidWiseTransformIterator<T, GPUContext>,
                                      const T *> {
 public:
  typedef thrust::iterator_adaptor<MidWiseTransformIterator<T, GPUContext>,
                                   const T *>
      super_t;
  HOSTDEVICE MidWiseTransformIterator(const T *x, int n, int post)
      : super_t(x), begin_(x), n_(n), post_(post) {}
  friend class thrust::iterator_core_access;

 private:
  unsigned int post_;
  unsigned int n_;
  const T *begin_;
  HOSTDEVICE typename super_t::reference dereference() const {
    return *(begin_ + (((this->base() - begin_) / post_) % n_));
  }
};
#endif

template <typename Functor,
          typename T,
          typename DeviceContext,
          typename OutType = T>
class TransformFunctor {
 public:
  TransformFunctor(const DenseTensor &x,
                   const DenseTensor &y,
                   DenseTensor *z,
                   const DeviceContext &dev_ctx,
                   Functor func,
                   const bool is_xsize_larger = true)
      : x_(x.data<T>()),
        y_(y.data<T>()),
        z_(dev_ctx.template Alloc<OutType>(z)),
        nx_(x.numel()),
        dev_ctx_(dev_ctx),
        func_(func),
        is_xsize_larger_(is_xsize_larger) {
    if (is_xsize_larger_ == false) {
      nx_ = y.numel();
    }
  }

  inline void Run() const {
    phi::Transform<DeviceContext> trans;
    trans(dev_ctx_, x_, x_ + nx_, y_, z_, func_);
  }

  inline void RunRowWise(int n) const {
    phi::Transform<DeviceContext> trans;
    if (is_xsize_larger_) {
      trans(dev_ctx_,
            x_,
            x_ + nx_,
            RowwiseTransformIterator<T, DeviceContext>(y_, n),
            z_,
            func_);
    } else {
      trans(dev_ctx_,
            y_,
            y_ + nx_,
            RowwiseTransformIterator<T, DeviceContext>(x_, n),
            z_,
            func_);
    }
  }

  inline void RunMidWise(int n, int post) const {
    phi::Transform<DeviceContext> trans;
    if (is_xsize_larger_) {
      trans(dev_ctx_,
            x_,
            x_ + nx_,
            MidWiseTransformIterator<T, DeviceContext>(y_, n, post),
            z_,
            func_);
    } else {
      trans(dev_ctx_,
            y_,
            y_ + nx_,
            MidWiseTransformIterator<T, DeviceContext>(x_, n, post),
            z_,
            func_);
    }
  }

 private:
  const T *x_;
  const T *y_;
  OutType *z_;
  int64_t nx_;
  const DeviceContext &dev_ctx_;
  Functor func_;
  bool is_xsize_larger_;
};

template <typename Functor, typename T, typename OutType = T>
void CommonForwardBroadcastCPU(const DenseTensor &x,
                               const DenseTensor &y,
                               DenseTensor *z,
                               int64_t *x_dims_array,
                               int64_t *y_dims_array,
                               int64_t *out_dims_array,
                               int max_dim,
                               const CPUContext &dev_ctx,
                               Functor func,
                               const bool is_xsize_larger = true) {
  std::vector<int64_t> index_array(max_dim, 0);
  const T *x_data = x.data<T>();
  const T *y_data = y.data<T>();
  if (z && z->numel() == 0) {
    dev_ctx.Alloc<OutType>(z);
    return;
  }
  OutType *out_data = dev_ctx.Alloc<OutType>(z);

  const int64_t out_size = std::accumulate(out_dims_array,
                                           out_dims_array + max_dim,
                                           1ll,
                                           std::multiplies<int64_t>());
  int64_t x_index, y_index;
  for (int64_t out_index = 0; out_index < out_size; ++out_index) {
    x_index =
        GetElementwiseIndex<int64_t>(x_dims_array, max_dim, index_array.data());
    y_index =
        GetElementwiseIndex<int64_t>(y_dims_array, max_dim, index_array.data());
    if (is_xsize_larger) {
      out_data[out_index] = func(x_data[x_index], y_data[y_index]);
    } else {
      out_data[out_index] = func(y_data[y_index], x_data[x_index]);
    }

    UpdateElementwiseIndexArray<int64_t>(
        out_dims_array, max_dim, index_array.data());
  }
}

template <typename Functor, typename T, typename OutType = T>
void CommonElementwiseBroadcastForward(const CPUContext &dev_ctx,
                                       const DenseTensor &x,
                                       const DenseTensor &y,
                                       DenseTensor *z,
                                       const DDim &x_dims,
                                       const DDim &y_dims,
                                       Functor func,
                                       int axis,
                                       const bool is_xsize_larger = true) {
  int max_dim = (std::max)(x_dims.size(), y_dims.size());
  axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
  PADDLE_ENFORCE_GE(
      axis,
      0,
      common::errors::InvalidArgument(
          "Axis should be great than or equal to 0, but received axis is %d.",
          axis));
  PADDLE_ENFORCE_LE(
      axis,
      max_dim,
      common::errors::InvalidArgument(
          "Axis should be less than or equal to %d, but received axis is %d.",
          max_dim,
          axis));
  std::vector<int64_t> x_dims_array(max_dim);
  std::vector<int64_t> y_dims_array(max_dim);
  std::vector<int64_t> out_dims_array(max_dim);
  GetBroadcastDimsArrays(x_dims,
                         y_dims,
                         x_dims_array.data(),
                         y_dims_array.data(),
                         out_dims_array.data(),
                         max_dim,
                         axis);

  CommonForwardBroadcastCPU<Functor, T, OutType>(x,
                                                 y,
                                                 z,
                                                 x_dims_array.data(),
                                                 y_dims_array.data(),
                                                 out_dims_array.data(),
                                                 max_dim,
                                                 dev_ctx,
                                                 func,
                                                 is_xsize_larger);
}

// It is a common CPU implementation to compute binary calculation with the
// support of broadcast. Note:
// 1. CPU implementation cannot support the case when x needs broadcast, thus
//    this function need to be called with XxxFunctor and XxxInverseFunctor,
//    like AddFunctor and InverseAddFunctor.
// 2. The corresponding GPU implementation supports all the broadcast cases,
//    thus there is no need to define and call with XxxInverseFunctor.
// TODO(liuyiqun): optimize the CPU implementation to support all broadcast
// cases and avoid the need of XxxInverseFunctor.
template <typename Functor, typename T, typename OutType = T>
void ElementwiseCompute(const CPUContext &dev_ctx,
                        const DenseTensor &x,
                        const DenseTensor &y,
                        Functor func,
                        DenseTensor *z,
                        int axis = -1) {
  dev_ctx.Alloc<OutType>(z);
  if (z && z->numel() == 0) {
    return;
  }
  auto x_dims = x.dims();
  auto y_dims = y.dims();
  bool is_xsize_larger = true;
  int max_dim = x_dims.size();
  if (x_dims.size() < y_dims.size()) {
    is_xsize_larger = false;
    max_dim = y_dims.size();
  }
  TransformFunctor<Functor, T, CPUContext, OutType> functor(
      x, y, z, dev_ctx, func, is_xsize_larger);
  if (x_dims == y_dims) {
    functor.Run();
    return;
  }

  axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
  PADDLE_ENFORCE_GE(
      axis,
      0,
      errors::InvalidArgument(
          "Axis should be great than or equal to 0, but received axis is %d.",
          axis));
  PADDLE_ENFORCE_LE(
      axis,
      max_dim,
      errors::InvalidArgument(
          "Axis should be less than or equal to %d, but received axis is %d.",
          max_dim,
          axis));

  size_t pre, n, post;
  int is_run_common_broadcast, axis_trim = 0;
  if (is_xsize_larger) {
    auto y_dims_trimmed = TrimTrailingSingularDims(y_dims);
    axis_trim = (y_dims_trimmed.size() == 0) ? x_dims.size() : axis;
    GetMidDims(x_dims,
               y_dims_trimmed,
               axis_trim,
               &pre,
               &n,
               &post,
               &is_run_common_broadcast);
  } else {
    auto x_dims_trimmed = TrimTrailingSingularDims(x_dims);
    axis_trim = (x_dims_trimmed.size() == 0) ? y_dims.size() : axis;
    GetMidDims(y_dims,
               x_dims_trimmed,
               axis_trim,
               &pre,
               &n,
               &post,
               &is_run_common_broadcast);
  }
  // special case for common implementation.
  // case 1: x=[2,3,1,5], y=[2,1,4,1]
  // case 2: x=[2,3,4], y=[1,1,4]
  if (is_run_common_broadcast == 1) {
    CommonElementwiseBroadcastForward<Functor, T, OutType>(
        dev_ctx, x, y, z, x_dims, y_dims, func, axis, is_xsize_larger);
    return;
  }

  if (post == 1) {
    functor.RunRowWise(n);
    return;
  } else {
    functor.RunMidWise(n, post);
    return;
  }
}

// for broadcast backwards
static inline std::vector<int> GetReduceDim(const DDim &in,
                                            const DDim &out,
                                            int axis) {
  axis =
      (axis == -1 ? std::abs(static_cast<int>(out.size() - in.size())) : axis);
  std::vector<int> dims;
  for (int i = 0; i < axis; ++i) {
    dims.push_back(i);
  }
  for (int i = 0; i < in.size(); ++i) {
    if (out[i + axis] != in[i]) {
      dims.push_back(i + axis);
    }
  }
  for (int i = axis + in.size(); i < out.size(); ++i) {
    dims.push_back(i);
  }
  return dims;
}

template <typename DeviceContext, typename T>
static inline void GetDoubleGradSafeTensor(const DeviceContext &dev_ctx,
                                           const DenseTensor &x,
                                           const DenseTensor *ddx,
                                           DenseTensor *ddx_safe) {
  if (ddx) {
    *ddx_safe = *ddx;
  } else {
    auto meta = DenseTensorMeta(x.dtype(), x.dims(), x.layout());
    *ddx_safe = Empty(dev_ctx, std::move(meta));
    dev_ctx.template Alloc<T>(ddx_safe);
    SetConstant<DeviceContext, T> set_zero;
    set_zero(dev_ctx, ddx_safe, static_cast<T>(0));
  }
}

inline void ElementwiseGradPreProcess(const DenseTensor &dout,
                                      DenseTensor *dx) {
  if (dx != nullptr) {
    dx->set_lod(dout.lod());
  }
}

#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)

// static unroller
template <template <int Index, int VecSize> typename Func,
          int VecSize,
          int End,
          int Begin = 0>
struct Unroller {
  template <typename... Args>
  static HOSTDEVICE inline void step(Args &&...args) {
    Func<Begin, VecSize>::Apply(std::forward<Args>(args)...);
    Unroller<Func, VecSize, End, Begin + 1>::step(args...);
  }
};

template <template <int Index, int VecSize> typename Func, int VecSize, int End>
struct Unroller<Func, VecSize, End, End> {
  template <typename... Args>
  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>
  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);
    if (is_boundary) {
      kps::ReadData<Type, VecSize, 1, ArgsT, Index, true>(
          args,
          reinterpret_cast<const _ptr_ Type *>(in[Index]) + offset,
          num,
          read_lens);
    } else {
      kps::ReadData<Type, VecSize, 1, ArgsT, Index, false>(
          args,
          reinterpret_cast<const _ptr_ Type *>(in[Index]) + offset,
          num,
          read_lens);
    }
  }
};

template <int Index>
struct InputSetter {
  template <typename Array, typename ArgsT>
  static void Apply(const std::vector<const DenseTensor *> &ins_tensor,
                    const ArgsT &args,
                    Array *ins_data) {
    using Type = std::tuple_element_t<Index, ArgsT>;
    (*ins_data)[Index] = (const _ptr_ char *)(ins_tensor[Index]->data<Type>());
  }
};

static inline int GetVectorizedSizeWithDtype(const DenseTensor *tensor) {
  int element_size = phi::SizeOf(tensor->dtype());
  if (element_size > sizeof(float)) {
    return 1;
  }
  constexpr int max_load_bits = 128;
  int vec_size = max_load_bits / CHAR_BIT / element_size;
  return vec_size;
}
static inline int GetVectorizedSizeWithAddress(const DenseTensor *tensor) {
  int element_size = phi::SizeOf(tensor->dtype());
  if (element_size > sizeof(float)) {
    return 1;
  }
  uint64_t address = reinterpret_cast<uint64_t>(tensor->data());

  // Currently, decide to deal with no more than 4 data once while adopting
  // vectorization load/store, if performance test shows that dealing with
  // 8 data once in vectorization load/store does get optimized, code below
  // can begin with :
  if (address % (element_size * 8) == 0) {
    return 8;
  } else if (address % (element_size * 4) == 0) {
    return 4;
  } else if (address % (element_size * 2) == 0) {
    return 2;
  } else {
    return 1;
  }
}

static int GetVectorizedSizeForTensors(
    const std::vector<const DenseTensor *> &ins,
    const std::vector<DenseTensor *> &outs,
    bool only_consider_outs_dtype = false) {
#ifdef PADDLE_WITH_XPU_KP
  int vec_size = 256;
#else
  constexpr int max_vec_size = 8;
  int vec_size = 1;
  if (!only_consider_outs_dtype) {
    for (size_t i = 0; i < ins.size(); ++i) {
      vec_size = std::max(vec_size, GetVectorizedSizeWithDtype(ins[i]));
    }
  }
  for (size_t i = 0; i < outs.size(); ++i) {
    vec_size = std::max(vec_size, GetVectorizedSizeWithDtype(outs[i]));
  }

  for (size_t i = 0; i < ins.size(); ++i) {
    vec_size = std::min(vec_size, GetVectorizedSizeWithAddress(ins[i]));
  }
  for (size_t i = 0; i < outs.size(); ++i) {
    vec_size = std::min(vec_size, GetVectorizedSizeWithAddress(outs[i]));
  }
  vec_size = std::min(vec_size, max_vec_size);
#endif
  return vec_size;
}

namespace detail {
template <class F, class Tuple, std::size_t... Index>
// GCC/Clang need the decltype() return type
HOSTDEVICE constexpr decltype(auto) ApplyImpl(F &&f,
                                              Tuple &&t,
                                              std::index_sequence<Index...>) {
  return std::forward<F>(f)(std::get<Index>(std::forward<Tuple>(t))...);
}
}  // namespace detail

template <class F, class Tuple>
HOSTDEVICE constexpr decltype(auto) Apply(F &&f, Tuple &&t) {
  return detail::ApplyImpl(
      std::forward<F>(f),
      std::forward<Tuple>(t),
      std::make_index_sequence<
          std::tuple_size<std::remove_reference_t<Tuple>>::value>{});
}

template <typename OutT,
          int VecSize,
          typename Functor,
          typename ArgsT,
          int Arity>
struct SameDimsElementwisePrimitiveCaller {
  __device__ inline void operator()(Functor func,
                                    ArgsT *args,
                                    OutT *result,
                                    int read_lens) {
#ifdef PADDLE_WITH_XPU_KP
    for (int idx = 0; idx < read_lens; ++idx) {
      result[idx] = static_cast<OutT>(Apply(func, args[idx]));
    }
#else
#pragma unroll
    for (int idx = 0; idx < VecSize; ++idx) {
      result[idx] = static_cast<OutT>(Apply(func, args[idx]));
    }
#endif
  }
};

template <typename OutT, int VecSize, bool IsBoundary, int NumOuts>
struct ElementwiseWriteDataCallerBc {
  __device__ __forceinline__ void operator()(
      Array<_ptr_ OutT *, NumOuts> outs,
      ConditionalT<OutT, NumOuts> src[VecSize],
      kps::IndexType block_offset,
      int num,
      int read_lens) {
    OutT dst[NumOuts][VecSize];
#pragma unroll
    for (int i = 0; i < read_lens; ++i) {
#pragma unroll
      for (int j = 0; j < NumOuts; ++j) {
        dst[j][i] = (src[i])[j];
      }
    }
#pragma unroll
    for (int i = 0; i < NumOuts; ++i) {
      kps::WriteData<OutT, VecSize, 1, IsBoundary>(
          outs[i] + block_offset, dst[i], num, read_lens);
    }
  }
};

template <typename OutT, int VecSize, bool IsBoundary>
struct ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, 1> {
  __device__ __forceinline__ void operator()(Array<_ptr_ OutT *, 1> outs,
                                             OutT src[VecSize],
                                             kps::IndexType block_offset,
                                             int num,
                                             int read_lens) {
    kps::WriteData<OutT, VecSize, 1, IsBoundary>(
        outs[0] + block_offset, src, num, read_lens);
  }
};

template <typename OutT,
          typename Functor,
          int Arity,
          int NumOuts,
          int VecSize,
          bool IsBoundary>
__device__ void VectorizedElementwiseKernelImpl(
    const Array<const _ptr_ char *__restrict__, Arity> &in,
    Array<_ptr_ OutT *, NumOuts> outs,
    kps::IndexType offset,
    int num,
    int read_lens,
    Functor func) {
  using Traits = funcs::FunctionTraits<Functor>;
  using ArgsT = typename Traits::ArgsTuple;
  ArgsT args[VecSize];
  ConditionalT<OutT, NumOuts> result[VecSize];

  Unroller<Loader, VecSize, Arity>::step(
      in, args, offset, num, read_lens, IsBoundary);

  SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
                                     VecSize,
                                     Functor,
                                     ArgsT,
                                     Arity>()(func, args, result, read_lens);

  ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, NumOuts>()(
      outs, result, offset, num, read_lens);
}

template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
__global__ void VectorizedElementwiseKernel(
    Array<const _ptr_ char *__restrict__, Arity> ins,
    Array<_ptr_ OutT *, NumOuts> outs,
    kps::IndexType numel,
    kps::IndexType main_offset,
    int read_lens,
    Functor func) {
  kps::IndexType data_offset =
      static_cast<kps::IndexType>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
  kps::IndexType stride =
      static_cast<kps::IndexType>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
  for (; data_offset < main_offset; data_offset += stride) {
    VectorizedElementwiseKernelImpl<OutT,
                                    Functor,
                                    Arity,
                                    NumOuts,
                                    VecSize,
                                    false>(
        ins, outs, data_offset, read_lens * BLOCK_NUM_X, read_lens, func);
  }

  kps::IndexType remain = numel - data_offset;
  if (remain > 0) {
    VectorizedElementwiseKernelImpl<OutT,
                                    Functor,
                                    Arity,
                                    NumOuts,
                                    VecSize,
                                    true>(
        ins, outs, data_offset, static_cast<int>(remain), read_lens, func);
  }
}

template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
void LaunchElementwiseKernel(const KPDevice &dev_ctx,
                             const std::vector<const DenseTensor *> &ins,
                             std::vector<DenseTensor *> *outs,
                             Functor func) {
  // There are at least 1 output, but maybe 0 input (ins.size() == 0).
  // For large tensor numel * sizeof(T) > 2^31, we must use int64_t as index
  // type.
  int64_t numel = (*outs)[0]->numel();
  Array<const _ptr_ char *__restrict__, Arity> ins_data;
  Array<_ptr_ OutT *, NumOuts> outs_data;

  using Traits = funcs::FunctionTraits<Functor>;
  using ArgsT = typename Traits::ArgsTuple;
  ArgsT arg;
  UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, arg, &ins_data);
  for (int i = 0; i < outs->size(); ++i) {
    outs_data[i] = (*outs)[i]->data<OutT>();
  }

#ifdef PADDLE_WITH_XPU_KP
  int block_size = 64;
  int grid_size = 8;
  int read_lens = kps::details::GetXpuReadLens(numel, block_size, grid_size);
  auto stream = dev_ctx.x_context()->xpu_stream;
  int64_t main_offset =
      (numel / (read_lens * block_size)) * read_lens * block_size;
  VectorizedElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSize>
      <<<grid_size, block_size, 0, stream>>>(
          ins_data, outs_data, numel, main_offset, read_lens, func);
#else
  auto gpu_config =
      phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, VecSize);
  int64_t main_offset = (numel / (VecSize * gpu_config.GetBlockSize())) *
                        VecSize * gpu_config.GetBlockSize();
  auto stream = dev_ctx.stream();
  VectorizedElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSize>
      <<<gpu_config.block_per_grid, gpu_config.thread_per_block, 0, stream>>>(
          ins_data, outs_data, numel, main_offset, VecSize, func);
#endif
}

template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
typename std::enable_if<!NeedVectorized<OutT>::value, void>::type
ElementwiseKernelForDifferentVecSize(
    const KPDevice &dev_ctx,
    const std::vector<const DenseTensor *> &ins,
    std::vector<DenseTensor *> *outs,
    Functor func) {
  LaunchElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSizeS>(
      dev_ctx, ins, outs, func);
}

template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
typename std::enable_if<NeedVectorized<OutT>::value, void>::type
ElementwiseKernelForDifferentVecSize(
    const KPDevice &dev_ctx,
    const std::vector<const DenseTensor *> &ins,
    std::vector<DenseTensor *> *outs,
    Functor func) {
  static int capability = dev_ctx.GetComputeCapability();
  // For Hopper and Blackwell, max vectorized size is 8.
  static int max_vec_size = capability >= 90 ? VecSizeVL : VecSizeL;
  // calculate the max vec_size for all ins and outs
  int vec_size = GetVectorizedSizeForTensors(ins, *outs);
  vec_size = std::min(vec_size, max_vec_size);

  switch (vec_size) {
    case VecSizeVL:
      LaunchElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSizeVL>(
          dev_ctx, ins, outs, func);
      break;
    case VecSizeL:
      LaunchElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSizeL>(
          dev_ctx, ins, outs, func);
      break;
    case VecSizeM:
      LaunchElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSizeM>(
          dev_ctx, ins, outs, func);
      break;
    case VecSizeS:
      LaunchElementwiseKernel<OutT, Functor, Arity, NumOuts, VecSizeS>(
          dev_ctx, ins, outs, func);
      break;
    default: {
      PADDLE_THROW(common::errors::Unimplemented(
          "Unsupported vectorized size: %d !", vec_size));
      break;
    }
  }
}

template <typename OutT, typename Functor, int NumOuts = 1>
void ElementwiseKernel(const KPDevice &dev_ctx,
                       const std::vector<const DenseTensor *> &ins,
                       std::vector<DenseTensor *> *outs,
                       Functor func) {
  using Traits = funcs::FunctionTraits<Functor>;
  const int kArity = Traits::arity;
  PADDLE_ENFORCE_EQ(ins.size(),
                    kArity,
                    common::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,
                    common::errors::InvalidArgument(
                        "Number of outputs shall equal to number of functions, "
                        "but number of outputs is %d, of functions is %d.",
                        outs->size(),
                        NumOuts));

  bool have_0_size = false;
  for (int i = 0; i < outs->size(); ++i) {
    if (outs->at(i)->numel() == 0) {
      have_0_size = true;
    }
    if (i > 0) {
      PADDLE_ENFORCE_EQ(
          (*outs)[i]->dims(),
          (*outs)[0]->dims(),
          common::errors::InvalidArgument(
              "The shape of each output tensor shall be identical yet, "
              "but %dth output tensor`s shape is not.",
              i));
    }
    dev_ctx.template Alloc<OutT>((*outs)[i]);
  }
  if (have_0_size) {
    return;
  }

  ElementwiseKernelForDifferentVecSize<OutT, Functor, kArity, NumOuts>(
      dev_ctx, ins, outs, func);
}

#endif

}  // namespace funcs
}  // namespace phi