/* Copyright (c) 2022 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 <sstream>
#include "paddle/common/enforce.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"

#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
#include "paddle/phi/kernels/funcs/dims_simplifier.h"

#endif

namespace phi {
namespace funcs {

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

enum BroadcastType { kMixed = 1, kBroadcast = 2, kElementwise = 3 };

template <typename OutT, typename Functor, int Arity, int NumOuts>
struct BroadcastTypeClassifier {
  int64_t numel{0};
  int broadcast_num{0};              // Not used for XPU
  bool all_elementwise{true};        // Not used for XPU
  Array<bool, Arity> use_broadcast;  // Not used for XPU
  Array<kps::details::BroadcastConfig, Arity> configs;
  Array<const _ptr_ char *__restrict__, Arity> ins_data;
  Array<_ptr_ OutT *, NumOuts> outs_data;

  BroadcastTypeClassifier() {}
  BroadcastTypeClassifier(const std::vector<const DenseTensor *> &ins,
                          std::vector<DenseTensor *> *outs,
                          int axis) {
    numel = (*outs)[0]->numel();

#ifndef PADDLE_WITH_XPU_KP
    for (size_t i = 0; i < ins.size(); ++i) {
      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;
    }
#endif

    InitBroadcastConfigs(ins, outs, axis);

    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 < NumOuts; ++i) {
      outs_data[i] = (*outs)[i]->data<OutT>();
    }
  }

  void InitBroadcastConfigs(const std::vector<const DenseTensor *> &ins,
                            std::vector<DenseTensor *> *outs,
                            int axis) {
#ifdef PADDLE_WITH_XPU_KP
    const auto dims_simplifier =
        BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
    if (VLOG_IS_ON(6)) {
      DimsSimplifiedLogger<int64_t>::Log(
          ins, outs, dims_simplifier, "BroadcastKernel");
    }
    configs[0] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
                                               dims_simplifier.in_dims[0],
                                               dims_simplifier.in_dims[1],
                                               dims_simplifier.rank);
    configs[1] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
                                               dims_simplifier.in_dims[1],
                                               dims_simplifier.in_dims[0],
                                               dims_simplifier.rank);
#else
    if (!all_elementwise) {
      const auto dims_simplifier =
          BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
      if (VLOG_IS_ON(6)) {
        DimsSimplifiedLogger<int64_t>::Log(
            ins, outs, dims_simplifier, "BroadcastKernel");
      }
      for (int i = 0; i < Arity; ++i) {
        // if data shape is[m, n], then you should set data_dim = {n, m}
        // eg: out's shape [3, 45, 1]. then out_dims = {1, 45, 3}
        // if (ins[i]->numel() != (*outs)[0]->numel()) {
        if (ins[i]->numel()) {
          configs[i] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
                                                     dims_simplifier.in_dims[i],
                                                     dims_simplifier.rank);
        }
      }
    }
#endif
  }
};

// Common broadcast/elementwise Loader.
template <int Index, int VecSize, bool IsBoundary, int LoadType>
struct BroadcastDataLoader {
  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,
                                               int read_lens) {
    using Type = std::tuple_element_t<Index, ArgsT>;
#ifdef PADDLE_WITH_XPU_KP
    kps::Init<Type, ArgsT, Index, VecSize>(
        args, static_cast<Type>(1.0f), read_lens);
    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,
          read_lens);
    } else {
      kps::ReadData<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
          args,
          reinterpret_cast<const _ptr_ Type *>(ins[Index]) + block_offset,
          num,
          read_lens);
    }
#else
    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);
    }
    // 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,
          num,
          VecSize);
    }
#endif
  }
};

/* BroadcastDataLoaders Partial specialization */
#ifndef PADDLE_WITH_XPU_KP
// 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,
                                               int read_lens) {
    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,
                                               int read_lens) {
    using Type = std::tuple_element_t<Index, ArgsT>;
    using VecType = phi::kps::details::VectorType<Type, VecSize>;
    VecType vec_temp;

    int64_t thread_offset =
        static_cast<int64_t>(threadIdx.x) +
        static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(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];
    }
  }
};

template <int Index, int VecSize>
struct BroadcastDataInit {
  template <typename ArgsT>
  static __device__ __forceinline__ void Apply(ArgsT *args) {
    using Type = std::tuple_element_t<Index, ArgsT>;
#pragma unroll
    for (int k = 0; k < VecSize; ++k) {
      std::get<Index>(args[k]) = static_cast<Type>(1);
    }
  }
};

template <int Index, int VecSize>
struct BroadcastDataSetter {
  template <typename Array, typename ArgsT>
  static __device__ __forceinline__ void Apply(const Array &ins,
                                               ArgsT *args,
                                               uint32_t index_bc[][VecSize]) {
    using Type = std::tuple_element_t<Index, ArgsT>;
#pragma unroll
    for (int k = 0; k < VecSize; ++k) {
      std::get<Index>(args[k]) =
          reinterpret_cast<const _ptr_ Type *>(ins[Index])[index_bc[Index][k]];
    }
  }
};

#endif

// 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...);
  }
};

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) {}
};

template <typename OutT,
          typename Functor,
          int Arity,
          int NumOuts,
          int VecSize,
          bool IsBoundary,
          int LoadType>
__device__ void VectorizedBroadcastKernelImpl(
    const Array<const _ptr_ char *__restrict__, Arity> &ins,
    Array<_ptr_ OutT *, NumOuts> outs,
    const Array<bool, Arity> &use_broadcast,
    const uint32_t numel,
    const Array<kps::details::BroadcastConfig, Arity> &configs,
    uint32_t num,
    uint32_t block_offset,
    int read_lens,
    Functor func) {
  using Traits = funcs::FunctionTraits<Functor>;
  using ArgsT = typename Traits::ArgsTuple;
  __simd__ ArgsT args[VecSize];
  __simd__ ConditionalT<OutT, NumOuts> result[VecSize];

#ifdef PADDLE_WITH_XPU_KP
  BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
      ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
#else
  if (LoadType == kBroadcast) {
    uint32_t index_bc[Arity][VecSize] = {0};
    Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
    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 < 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];
        }
      }
    }
    Unroller<BroadcastDataSetter, VecSize, Arity>::step(ins, args, index_bc);
  } else {
    BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
        ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
  }
#endif
  SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
                                     VecSize,
                                     Functor,
                                     ArgsT,
                                     Arity>()(func, args, result, read_lens);
  funcs::ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, NumOuts>()(
      outs, result, block_offset, num, read_lens);
}

template <typename Functor,
          typename OutT,
          int Arity,
          int NumOuts,
          int VecSize,
          int LoadType>
__global__ void VectorizedBroadcastKernel(
    Array<const _ptr_ char *__restrict__, Arity> ins,
    Array<_ptr_ OutT *, NumOuts> outs,
    Array<bool, Arity> use_broadcast,
    uint32_t numel,
    Array<kps::details::BroadcastConfig, Arity> configs,
    uint32_t main_offset,
    uint32_t tail_tid,
    int read_lens,
    Functor func) {
#ifdef PADDLE_WITH_XPU_KP
  int64_t block_offset =
      static_cast<int64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
  int64_t stride = static_cast<int64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
  for (; block_offset < main_offset; block_offset += stride) {
    VectorizedBroadcastKernelImpl<OutT,
                                  Functor,
                                  Arity,
                                  NumOuts,
                                  VecSize,
                                  false,
                                  LoadType>(ins,
                                            outs,
                                            use_broadcast,
                                            numel,
                                            configs,
                                            BLOCK_NUM_X * read_lens,
                                            block_offset,
                                            read_lens,
                                            func);
  }
  int64_t num = numel - block_offset;
  if (num > 0) {
    VectorizedBroadcastKernelImpl<OutT,
                                  Functor,
                                  Arity,
                                  NumOuts,
                                  VecSize,
                                  true,
                                  LoadType>(ins,
                                            outs,
                                            use_broadcast,
                                            numel,
                                            configs,
                                            num,
                                            block_offset,
                                            read_lens,
                                            func);
  }
#else
  int64_t block_offset =
      static_cast<int64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
  if (block_offset < main_offset) {
    VectorizedBroadcastKernelImpl<OutT,
                                  Functor,
                                  Arity,
                                  NumOuts,
                                  VecSize,
                                  false,
                                  LoadType>(ins,
                                            outs,
                                            use_broadcast,
                                            numel,
                                            configs,
                                            BLOCK_NUM_X * VecSize,
                                            block_offset,
                                            read_lens,
                                            func);
  } else {
    VectorizedBroadcastKernelImpl<OutT,
                                  Functor,
                                  Arity,
                                  NumOuts,
                                  VecSize,
                                  true,
                                  LoadType>(ins,
                                            outs,
                                            use_broadcast,
                                            numel,
                                            configs,
                                            tail_tid,
                                            block_offset,
                                            read_lens,
                                            func);
  }
#endif
}

template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
void LaunchBroadcastKernel(
    const KPDevice &dev_ctx,
    const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
    Functor func) {
#ifdef PADDLE_WITH_XPU_KP
  const int64_t numel_64 = classifier.numel;
  PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "XPU broadcast kernel numel");
  const uint32_t numel = static_cast<uint32_t>(numel_64);
  const int threads = 64;
  const int blocks = 8;
  int read_lens = configs[0].buf_len;
  auto stream = dev_ctx.x_context()->xpu_stream;
  const int64_t block_len = static_cast<int64_t>(read_lens) * threads;
  const int64_t main_offset_64 = (numel_64 / block_len) * block_len;
  const int64_t tail_tid_64 = numel_64 % block_len;
  const uint32_t main_offset = static_cast<uint32_t>(main_offset_64);
  const uint32_t tail_tid = static_cast<uint32_t>(tail_tid_64);

  VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
      <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                       classifier.outs_data,
                                       classifier.use_broadcast,
                                       numel,
                                       classifier.configs,
                                       main_offset,
                                       tail_tid,
                                       read_lens,
                                       func);
#else
  const int64_t numel_64 = classifier.numel;
  auto gpu_config =
      phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel_64, VecSize);
  auto stream = dev_ctx.stream();
  uint32_t threads = static_cast<uint32_t>(gpu_config.GetBlockSize());
  auto blocks = gpu_config.block_per_grid;
  PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "numel");
  const uint32_t numel = static_cast<uint32_t>(numel_64);
  const int64_t block_len = static_cast<int64_t>(VecSize) * threads;
  const int64_t main_offset_64 = (numel_64 / block_len) * block_len;
  const int64_t tail_tid_64 = numel_64 % block_len;
  uint32_t main_offset = static_cast<uint32_t>(main_offset_64);
  uint32_t tail_tid = static_cast<uint32_t>(tail_tid_64);

  if (classifier.all_elementwise) {
    VectorizedBroadcastKernel<Functor,
                              OutT,
                              Arity,
                              NumOuts,
                              VecSize,
                              kElementwise>
        <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                         classifier.outs_data,
                                         classifier.use_broadcast,
                                         numel,
                                         classifier.configs,
                                         main_offset,
                                         tail_tid,
                                         VecSize,
                                         func);
  } else if (classifier.broadcast_num > (Arity >> 1)) {
    constexpr BroadcastType type_ = (Arity > 1) ? kBroadcast : kMixed;
    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
        <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                         classifier.outs_data,
                                         classifier.use_broadcast,
                                         numel,
                                         classifier.configs,
                                         main_offset,
                                         tail_tid,
                                         VecSize,
                                         func);
  } else {
    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, kMixed>
        <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                         classifier.outs_data,
                                         classifier.use_broadcast,
                                         numel,
                                         classifier.configs,
                                         main_offset,
                                         tail_tid,
                                         VecSize,
                                         func);
  }
#endif
}

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

template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
typename std::enable_if<NeedVectorized<OutT>::value, void>::type
BroadcastKernelForDifferentVecSize(const KPDevice &dev_ctx,
                                   const std::vector<const DenseTensor *> &ins,
                                   std::vector<DenseTensor *> *outs,
                                   int axis,
                                   Functor func) {
  auto classifier =
      BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(ins, outs, axis);
#ifdef PADDLE_WITH_XPU_KP
  auto type = kps::details::OptType::CanNotOptimize;
  bool is_optimize = classifier.configs[0].cmp_type != type;
  int vec_size = is_optimize ? VecSizeL : VecSizeM;
#else
  static int capability = dev_ctx.GetComputeCapability();
  // For Hopper and Blackwell, max vectorized size is VecSizeL(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, true);
  vec_size = std::min(vec_size, max_vec_size);
  int64_t numel = classifier.numel;
  // For small tensor, using VecSizeL can improve performance more than
  // VecSizeVL
  constexpr int64_t large_vect_threshold = 1024 * 1024 * 4;
  if (numel < large_vect_threshold) {
    vec_size = std::min(vec_size, VecSizeL);
  }
#endif

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

static void updateStridesDims(std::vector<int64_t> *strides,
                              std::vector<int64_t> *dims) {
  for (int i = 1; i < strides->size(); i++) {
    (*strides)[i] = (*strides)[i - 1] * (*dims)[i - 1];
  }
  // reverse origin_in_dim and origin_in_stride if in's dim_size > 0
  std::reverse(strides->begin(), strides->end());
  std::reverse(dims->begin(), dims->end());
}

static void SliceTensor(DenseTensor *x,
                        const DenseTensor *share,
                        const std::vector<int64_t> &out_compute_dims,
                        int64_t offset) {
  auto new_dim = make_ddim(out_compute_dims);
  DenseTensorMeta meta(share->dtype(),
                       new_dim,
                       share->layout(),
                       offset * SizeOf(share->dtype()) + share->offset());
  x->set_meta(meta);
  x->ShareBufferWith(*(share), true);
  x->Resize(new_dim);
}

template <typename OutT, typename Functor, int kArity, int NumOuts = 1>
void BroadcastKernelSplit(const KPDevice &dev_ctx,
                          const std::vector<const DenseTensor *> &ins,
                          std::vector<DenseTensor *> *outs,
                          int axis,
                          Functor func,
                          const int64_t compute_size) {
  const auto dims_simplifier =
      BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
  if (VLOG_IS_ON(6)) {
    DimsSimplifiedLogger<int64_t>::Log(
        ins, outs, dims_simplifier, "GPU Broadcast");
  }

  int all_rank = dims_simplifier.rank;
  std::vector<int64_t> origin_out_strides(all_rank, 1);
  auto origin_in_dims = dims_simplifier.in_dims;
  auto origin_out_dims = dims_simplifier.out_dims;
  auto origin_in_strides = dims_simplifier.in_dims;

  // for split
  std::vector<int64_t> loop_num_out(all_rank, 1);
  std::vector<int64_t> loop_num_out_stride(all_rank, 1);

  // for input's offset
  std::vector<int64_t> ins_offset(kArity, 0);
  std::vector<int64_t> ins_scale_for_dim(kArity, 0);

  // init offset and check in's dim
  for (int k = 0; k < kArity; k++) {
    ins_scale_for_dim[k] = ins[k]->dims().size() == 0 ? 0 : 1;
    if (ins_scale_for_dim[k]) {
      origin_in_strides[k][0] = 1;
    }
  }

  updateStridesDims(&origin_out_strides, &origin_out_dims);
  for (int k = 0; k < kArity; k++) {
    if (ins_scale_for_dim[k]) {
      updateStridesDims(&origin_in_strides[k], &origin_in_dims[k]);
    }
  }

  // init out_split_dim and in_split_dims
  auto out_split_dim = origin_out_dims;
  auto in_split_dims = origin_in_dims;

  // init
  int64_t loop_num = 1;
  int64_t split_idx = 0;

  for (int r = 0; r < all_rank; r++) {
    // if the compute_size was too small the split_size must be 0, but the
    // dim_num must ge 1
    int64_t split_size = compute_size / origin_out_strides[r];
    out_split_dim[r] = std::max(split_size, static_cast<int64_t>(1));
    loop_num_out[r] =
        (origin_out_dims[r] + out_split_dim[r] - 1) / out_split_dim[r];
    loop_num *= loop_num_out[r];

    for (int k = 0; k < kArity; k++) {
      if (ins_scale_for_dim[k]) {
        in_split_dims[k][r] = std::min(origin_in_dims[k][r], out_split_dim[r]);
      }
    }

    // split_idx is the index for lash split dim
    if (split_size != 0) {
      split_idx = r;
      break;
    }
  }

  loop_num_out_stride[all_rank - 1] = 1;
  for (int r = all_rank - 2; r >= 0; r--) {
    loop_num_out_stride[r] = loop_num_out_stride[r + 1] * loop_num_out[r + 1];
  }

  // compute

  for (int iter = 0; iter < loop_num; iter++) {
    std::vector<const DenseTensor *> new_ins = {};
    std::vector<DenseTensor *> new_outs = {};
    DenseTensor tmp_in[kArity];
    DenseTensor tmp_out[NumOuts];

    int64_t tmp_size = iter;
    int64_t out_offset = 0;
    // compute the offset before  last split dim
    for (int i = 0; i < split_idx; i++) {
      auto repeat_times = tmp_size / loop_num_out_stride[i];
      out_offset += repeat_times * origin_out_strides[i];
      for (int k = 0; k < kArity; k++) {
        if (ins_scale_for_dim[k]) {
          ins_offset[k] +=
              (repeat_times % origin_in_dims[k][i]) * origin_in_strides[k][i];
        }
      }
      tmp_size = tmp_size % loop_num_out_stride[i];
    }
    // tmp_size is the last split_dims's repeat idx
    auto pre_deal_size = tmp_size * out_split_dim[split_idx];
    out_offset += pre_deal_size * origin_out_strides[split_idx];
    // compute_size
    auto remainder_size = origin_out_dims[split_idx] - pre_deal_size;

    // get current compute size
    auto out_compute_dims = out_split_dim;
    out_compute_dims[split_idx] =
        std::min(out_split_dim[split_idx], remainder_size);

    // in + compute_size
    auto in_compute_dims = in_split_dims;
    for (int k = 0; k < kArity; k++) {
      if (ins_scale_for_dim[k]) {
        auto split_repeat =
            origin_in_dims[k][split_idx] == origin_out_dims[split_idx]
                ? tmp_size
                : 0;
        ins_offset[k] += split_repeat * in_split_dims[k][split_idx] *
                         origin_in_strides[k][split_idx];
        in_compute_dims[k][split_idx] =
            std::min(in_split_dims[k][split_idx], out_compute_dims[split_idx]);
      }
      SliceTensor(&tmp_in[k],
                  ins[k],
                  in_compute_dims[k],
                  ins_scale_for_dim[k] * ins_offset[k]);
      new_ins.emplace_back(&tmp_in[k]);
      ins_offset[k] = 0;
    }

    for (int n = 0; n < NumOuts; n++) {
      SliceTensor(&tmp_out[n], (*outs)[n], out_compute_dims, out_offset);
      new_outs.emplace_back(&tmp_out[n]);
    }

    BroadcastKernelForDifferentVecSize<OutT, Functor, kArity, NumOuts>(
        dev_ctx, new_ins, &new_outs, axis, func);
  }
}

template <typename OutT, typename Functor, int kArity, int NumOuts = 1>
void BroadcastKernelApply(const KPDevice &dev_ctx,
                          const std::vector<const DenseTensor *> &ins,
                          std::vector<DenseTensor *> *outs,
                          int axis,
                          Functor func) {
#ifndef PADDLE_WITH_XPU_KP
  auto compute_size =
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
  bool use_int64_index_kernel = false;
  for (auto *out : *outs) {
    if (out->numel() >= compute_size) {
      use_int64_index_kernel = true;
    }
  }
  if (use_int64_index_kernel) {  // use_int64_index_kernel
    BroadcastKernelSplit<OutT, Functor, kArity, NumOuts>(
        dev_ctx, ins, outs, axis, func, compute_size);
    return;
  }
#endif
  BroadcastKernelForDifferentVecSize<OutT, Functor, kArity, NumOuts>(
      dev_ctx, ins, outs, axis, func);
}

template <typename OutT, typename Functor, int NumOuts = 1>
void BroadcastKernel(const KPDevice &dev_ctx,
                     const std::vector<const DenseTensor *> &ins,
                     std::vector<DenseTensor *> *outs,
                     Functor func,
                     int axis = -1) {
  // When there are multiple inputs, the outputs's rank should be equal the
  // maximum rank of all inputs.
  using Traits = funcs::FunctionTraits<Functor>;
  const int kArity = Traits::arity;

#ifdef PADDLE_WITH_XPU_KP
  PADDLE_ENFORCE_EQ(
      ins.size(),
      2,
      common::errors::InvalidArgument(
          "XPU only support inputs is 2, but received %d", ins.size()));
#endif

  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));

  for (auto i = 0; i < outs->size(); ++i) {
    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 "
              "%d-th output tensor`s shape is not.",
              i));
    }
    dev_ctx.template Alloc<OutT>((*outs)[i]);
  }
  if ((*outs)[0]->numel() == 0) {
    return;
  }
  int max_rank = 0;
  int min_rank = DDim::kMaxRank;
  for (auto *in : ins) {
    max_rank = std::max(max_rank, in->dims().size());
    min_rank = std::min(min_rank, in->dims().size());
  }
  if (ins.size() == 1) {
    // When there is only 1 input, the input's rank may be less than outputs'
    // rank.
    max_rank = std::max(max_rank, (*outs)[0]->dims().size());
  }
  axis = axis == -1 ? max_rank - min_rank : axis;
  BroadcastKernelApply<OutT, Functor, kArity, NumOuts>(
      dev_ctx, ins, outs, axis, func);
}

template <typename Functor, typename T, typename OutType = T>
void ElementwiseCompute(const GPUContext &dev_ctx,
                        const DenseTensor &x,
                        const DenseTensor &y,
                        Functor func,
                        DenseTensor *z,
                        int axis = -1) {
  std::vector<const DenseTensor *> ins = {&x, &y};
  std::vector<DenseTensor *> outs = {z};
  dev_ctx.template Alloc<OutType>(z);

  BroadcastKernel<OutType, Functor, 1>(dev_ctx, ins, &outs, func, axis);
}

template <typename DeviceContext,
          typename T,
          typename Functor,
          typename InverseFunctor>
void DefaultElementwiseOperator(const DeviceContext &dev_ctx,
                                const DenseTensor &x,
                                const DenseTensor &y,
                                DenseTensor *z,
                                int axis = -1) {
  auto x_dims = x.dims();
  auto y_dims = y.dims();
  dev_ctx.template Alloc<T>(z);
  funcs::ElementwiseCompute<Functor, T>(dev_ctx, x, y, Functor(), z, axis);
}

#else

template <typename DeviceContext,
          typename T,
          typename Functor,
          typename InverseFunctor>
void DefaultElementwiseOperator(const DeviceContext &dev_ctx,
                                const DenseTensor &x,
                                const DenseTensor &y,
                                DenseTensor *z,
                                int axis = -1) {
  auto x_dims = x.dims();
  auto y_dims = y.dims();
  dev_ctx.template Alloc<T>(z);
  if (x_dims.size() >= y_dims.size()) {
    funcs::ElementwiseCompute<Functor, T>(dev_ctx, x, y, Functor(), z, axis);
  } else {
    funcs::ElementwiseCompute<InverseFunctor, T>(
        dev_ctx, x, y, InverseFunctor(), z, axis);
  }
}
#endif

}  // namespace funcs
}  // namespace phi