// Copyright (c) 2024 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 <glog/logging.h>

#include <algorithm>
#include <functional>  // for multiplies
#include <iterator>
#include <vector>

#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/common/transform.h"
#include "paddle/phi/core/dense_tensor.h"
#if !defined(PADDLE_WITH_CUDA) || !defined(PADDLE_WITH_CUSTOM_DEVICE)
#include "paddle/phi/kernels/cpu/elementwise.h"
#include "paddle/phi/kernels/cpu/elementwise_grad.h"
#endif
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"

#if defined(__NVCC__) || defined(__HIPCC__)
#ifdef __NVCC__
#include <cuda.h>
#elif defined(__HIPCC__)
#include <hip/hip_runtime.h>
#endif
#include <thrust/iterator/iterator_adaptor.h>

#include "paddle/phi/backends/gpu/gpu_device_function.h"
#include "paddle/phi/backends/gpu/gpu_primitives.h"
#include "paddle/phi/kernels/funcs/elementwise/elementwise_op_broadcast.cu.h"
#include "paddle/phi/kernels/funcs/reduce_function.h"
#include "paddle/phi/kernels/gpu/elementwise_grad.h"

#endif

#include "paddle/phi/kernels/funcs/for_range.h"
#include "paddle/phi/kernels/funcs/math_function.h"

#define DIVUP(x, y) (((x) + (y)-1) / (y))

#define ROUNDUP(x, y) (DIVUP((x), (y)) * (y))

namespace phi {
namespace funcs {

// It is a common implementation to compute binary calculation with the support
// of broadcast, supporting both CPU and GPU.
// - 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.
// - 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 DeviceContext,
          typename T,
          typename OutType = T>
void ElementwiseComputeEx(const DeviceContext &dev_ctx,
                          const DenseTensor *x,
                          const DenseTensor *y,
                          int axis,
                          Functor func,
                          DenseTensor *z) {
  dev_ctx.template Alloc<OutType>(z);
  funcs::ElementwiseCompute<Functor, T, OutType>(
      dev_ctx, *x, *y, func, z, axis);
}

// FusedElemwiseAndAct
// --- forward
template <typename T, typename CompoundFunctor, bool KeepIntermediateOut>
struct FusedElemwiseAndActNoBroadcast {
  HOSTDEVICE void operator()(size_t i) {
    T y_val = y_[i];
    T x_val = x_[i];
    if (KeepIntermediateOut) {
      T intermeidiate_out = compound_functor_.GetIntermediateOut(x_val, y_val);
      intermediate_out_[i] = intermeidiate_out;
      out_[i] =
          compound_functor_.GetOutUseIntermediateOut(x_val, intermeidiate_out);
    } else {
      out_[i] = compound_functor_.GetOut(x_val, y_val);
    }
  }

  const T *x_;
  const T *y_;
  CompoundFunctor compound_functor_;
  T *out_;
  T *intermediate_out_;
};

// FusedElemwiseAndActBroadcast1:
// In this case, X and Y can be reshaped to a matrix.
// For example shape(X) = (2, 3, 4, 5), shape(Y) = (4, 5) and axis = -1 or 2,
// X can be reshaped to (6, 20) and Y can be reshaped to (1, 20)
template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActBroadcast1CPU(const T *x,
                                             const T *y,
                                             CompoundFunctor compound_functor,
                                             int h,
                                             int w,
                                             T *out,
                                             T *intermediate_out) {
  for (int i = 0; i < h; ++i) {
    for (int j = 0; j < w; ++j) {
      int64_t offset = static_cast<int64_t>(i) * w + j;

      T y_val = BcastY ? y[j] : y[offset];
      T x_val = BcastY ? x[offset] : x[j];
      int64_t intermediate_out_offset;
      if (KeepIntermediateOut) {
        T intermeidiate_out = compound_functor.GetIntermediateOut(x_val, y_val);

        if (SameShapeOfIntermediateOutAndOut) {
          // for the case of f1(f2(x, y))
          intermediate_out_offset = offset;
        } else if (BcastY) {
          intermediate_out_offset = j;
        } else {
          intermediate_out_offset = offset;
        }

        intermediate_out[intermediate_out_offset] = intermeidiate_out;
        out[offset] =
            compound_functor.GetOutUseIntermediateOut(x_val, intermeidiate_out);
      } else {
        out[offset] = compound_functor.GetOut(x_val, y_val);
      }
    }
  }
}

// FusedElemwiseAndActBroadcast2
// In this case, X and Y can be reshaped to a matrix.
// For example shape(X) = (2, 3, 4, 5), shape(Y) = (3, 4) and axis = 1,
// X can be reshaped to (2, 12, 5) and Y can be reshaped to (1, 12, 1)
// pre = 2, n = 12, post = 5
template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActBroadcast2CPU(const T *x,
                                             const T *y,
                                             int pre,
                                             int n,
                                             int post,
                                             CompoundFunctor compound_functor,
                                             T *out,
                                             T *intermediate_out) {
  for (int i = 0; i < pre; ++i) {
    for (int j = 0; j < n; ++j) {
      for (int k = 0; k < post; ++k) {
        int64_t offset = static_cast<int64_t>(i) * n * post +
                         static_cast<int64_t>(j) * post + k;

        T y_val = BcastY ? y[j] : y[offset];
        T x_val = BcastY ? x[offset] : x[j];
        int64_t intermediate_out_offset;

        if (KeepIntermediateOut) {
          T intermeidiate_out =
              compound_functor.GetIntermediateOut(x_val, y_val);

          if (SameShapeOfIntermediateOutAndOut) {
            // for the case of f1(f2(x, y))
            intermediate_out_offset = offset;
          } else if (BcastY) {
            intermediate_out_offset = j;
          } else {
            intermediate_out_offset = offset;
          }

          intermediate_out[intermediate_out_offset] = intermeidiate_out;
          out[offset] = compound_functor.GetOutUseIntermediateOut(
              x_val, intermeidiate_out);
        } else {
          out[offset] = compound_functor.GetOut(x_val, y_val);
        }
      }
    }
  }
}

#if defined(__NVCC__) || defined(__HIPCC__)
template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static __global__ void FusedElemwiseAndActBroadcast1CUDAKernel(
    const T *x,
    const T *y,
    int h,
    int w,
    CompoundFunctor compound_functor,
    T *out,
    T *intermediate_out) {
  int i = blockIdx.x;
  int j = threadIdx.x;

  while (j < w) {
    int64_t offset = static_cast<int64_t>(i) * w + j;

    T y_val = BcastY ? y[j] : y[offset];
    T x_val = BcastY ? x[offset] : x[j];
    int64_t intermediate_out_offset;

    if (KeepIntermediateOut) {
      T intermeidiate_out = compound_functor.GetIntermediateOut(x_val, y_val);

      if (SameShapeOfIntermediateOutAndOut) {
        // for the case of f1(f2(x, y))
        intermediate_out_offset = offset;
      } else if (BcastY) {
        intermediate_out_offset = j;
      } else {
        intermediate_out_offset = offset;
      }

      intermediate_out[intermediate_out_offset] = intermeidiate_out;
      out[offset] =
          compound_functor.GetOutUseIntermediateOut(x_val, intermeidiate_out);
    } else {
      out[offset] = compound_functor.GetOut(x_val, y_val);
    }

    j += ELEMWISE_MAX_BLOCK_DIM;
  }
}

template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActBroadcast1CUDA(gpuStream_t stream,
                                              const T *x,
                                              const T *y,
                                              CompoundFunctor compound_functor,
                                              int64_t h,
                                              int64_t w,
                                              T *out,
                                              T *intermediate_out) {
  int64_t block_size =
      std::min(static_cast<int64_t>(ELEMWISE_MAX_BLOCK_DIM), w);
  int64_t gird_size = h;
  FusedElemwiseAndActBroadcast1CUDAKernel<T,
                                          CompoundFunctor,
                                          BcastY,
                                          KeepIntermediateOut,
                                          SameShapeOfIntermediateOutAndOut>
      <<<gird_size, block_size, 0, stream>>>(
          x, y, h, w, compound_functor, out, intermediate_out);
}

template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static __global__ void FusedElemwiseAndActBroadcast2CUDAKernel(
    const T *x,
    const T *y,
    CompoundFunctor compound_functor,
    int pre,
    int n,
    int post,
    T *out,
    T *intermediate_out) {
  int tid = threadIdx.x;
  int j = blockIdx.x;

  while (true) {
    int i = tid / post;
    int k = tid % post;
    if (i >= pre) break;

    int64_t offset =
        static_cast<int64_t>(i) * n * post + static_cast<int64_t>(j) * post + k;

    T y_val = BcastY ? y[j] : y[offset];
    T x_val = BcastY ? x[offset] : x[j];
    int64_t intermediate_out_offset;

    if (KeepIntermediateOut) {
      T intermeidiate_out = compound_functor.GetIntermediateOut(x_val, y_val);

      if (SameShapeOfIntermediateOutAndOut) {
        // for the case of f1(f2(x, y))
        intermediate_out_offset = offset;
      } else if (BcastY) {
        intermediate_out_offset = j;
      } else {
        intermediate_out_offset = offset;
      }

      intermediate_out[intermediate_out_offset] = intermeidiate_out;
      out[offset] =
          compound_functor.GetOutUseIntermediateOut(x_val, intermeidiate_out);
    } else {
      out[offset] = compound_functor.GetOut(x_val, y_val);
    }

    tid += ELEMWISE_MAX_BLOCK_DIM;
  }
}

template <typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActBroadcast2CUDA(gpuStream_t stream,
                                              const T *x,
                                              const T *y,
                                              int64_t pre,
                                              int64_t n,
                                              int64_t post,
                                              CompoundFunctor compound_functor,
                                              T *out,
                                              T *intermediate_out) {
  int64_t block_size =
      std::min(static_cast<int64_t>(ELEMWISE_MAX_BLOCK_DIM), pre * post);
  int64_t gird_size = n;

  FusedElemwiseAndActBroadcast2CUDAKernel<T,
                                          CompoundFunctor,
                                          BcastY,
                                          KeepIntermediateOut,
                                          SameShapeOfIntermediateOutAndOut>
      <<<gird_size, block_size, 0, stream>>>(
          x, y, compound_functor, pre, n, post, out, intermediate_out);
}

#endif

template <typename DeviceContext,
          typename T,
          typename CompoundFunctor,
          bool KeepIntermediateOut>
void FusedElemwiseAndActComputeNoBroadcast(const DeviceContext &dev_ctx,
                                           const DDim &x_dim,
                                           const DenseTensor &x,
                                           const DenseTensor &y,
                                           CompoundFunctor compound_functor,
                                           DenseTensor *out,
                                           DenseTensor *intermediate_out) {
  size_t N = static_cast<size_t>(common::product(x_dim));

  funcs::ForRange<DeviceContext> for_range(dev_ctx, N);

  for_range(
      FusedElemwiseAndActNoBroadcast<T, CompoundFunctor, KeepIntermediateOut>{
          x.data<T>(),
          y.data<T>(),
          compound_functor,
          dev_ctx.template Alloc<T>(out),
          intermediate_out == nullptr
              ? nullptr
              : dev_ctx.template Alloc<T>(intermediate_out)});
}

template <typename DeviceContext,
          typename T,
          typename CompoundFunctor,
          bool BcastY,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
void FusedElemwiseAndActComputeWithBroadcast(const DeviceContext &dev_ctx,
                                             const DDim &x_dim,
                                             const DDim &y_dim_untrimed,
                                             const DenseTensor &x,
                                             const DenseTensor &y,
                                             CompoundFunctor compound_functor,
                                             int axis,
                                             DenseTensor *out,
                                             DenseTensor *intermediate_out) {
  axis = (axis == -1 ? x_dim.size() - y_dim_untrimed.size() : axis);
  auto y_dim = funcs::TrimTrailingSingularDims(y_dim_untrimed);
  axis = (y_dim.size() == 0) ? x_dim.size() : axis;

  size_t pre, n, post;
  int is_run_common_broadcast;
  funcs::GetMidDims(
      x_dim, y_dim, axis, &pre, &n, &post, &is_run_common_broadcast);
  if (post == 1) {
    int h = pre;
    int w = n;
    if (dev_ctx.GetPlace().GetType() == AllocationType::GPU) {
#if defined(__NVCC__) || defined(__HIPCC__)
      FusedElemwiseAndActBroadcast1CUDA<T,
                                        CompoundFunctor,
                                        BcastY,
                                        KeepIntermediateOut,
                                        SameShapeOfIntermediateOutAndOut>(
          dev_ctx.stream(),
          x.data<T>(),
          y.data<T>(),
          compound_functor,
          h,
          w,
          dev_ctx.template Alloc<T>(out),
          intermediate_out == nullptr
              ? nullptr
              : dev_ctx.template Alloc<T>(intermediate_out));
#endif
    } else {
      FusedElemwiseAndActBroadcast1CPU<T,
                                       CompoundFunctor,
                                       BcastY,
                                       KeepIntermediateOut,
                                       SameShapeOfIntermediateOutAndOut>(
          x.data<T>(),
          y.data<T>(),
          compound_functor,
          h,
          w,
          dev_ctx.template Alloc<T>(out),
          intermediate_out == nullptr
              ? nullptr
              : dev_ctx.template Alloc<T>(intermediate_out));
    }
  } else {
    if (dev_ctx.GetPlace().GetType() == AllocationType::GPU) {
#if defined(__NVCC__) || defined(__HIPCC__)
      FusedElemwiseAndActBroadcast2CUDA<T,
                                        CompoundFunctor,
                                        BcastY,
                                        KeepIntermediateOut,
                                        SameShapeOfIntermediateOutAndOut>(
          dev_ctx.stream(),
          x.data<T>(),
          y.data<T>(),
          pre,
          n,
          post,
          compound_functor,
          dev_ctx.template Alloc<T>(out),
          intermediate_out == nullptr
              ? nullptr
              : dev_ctx.template Alloc<T>(intermediate_out));
#endif
    } else {
      FusedElemwiseAndActBroadcast2CPU<T,
                                       CompoundFunctor,
                                       BcastY,
                                       KeepIntermediateOut,
                                       SameShapeOfIntermediateOutAndOut>(
          x.data<T>(),
          y.data<T>(),
          pre,
          n,
          post,
          compound_functor,
          dev_ctx.template Alloc<T>(out),
          intermediate_out == nullptr
              ? nullptr
              : dev_ctx.template Alloc<T>(intermediate_out));
    }
  }
}

// --- backward
template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut>
struct FusedElemwiseAndActGradNoBroadcast {
  HOSTDEVICE void operator()(size_t i) {
    T zero = static_cast<T>(0);
    T x_val = (x_ == nullptr) ? zero : x_[i];
    T y_val = (y_ == nullptr) ? zero : y_[i];
    T out_val = out_[i];
    T dout_val = dout_[i];
    T intermediate_out_val = UseIntermediateOut
                                 ? intermediate_out_[i]
                                 : dx_op_.GetIntermediateOut(x_val, y_val);
    if (dx_ != nullptr) {
      dx_[i] = dx_op_.UseIntermediateOut(
          x_val, y_val, intermediate_out_val, out_val, dout_val);
    }
    if (dy_ != nullptr) {
      dy_[i] = dy_op_.UseIntermediateOut(
          x_val, y_val, intermediate_out_val, out_val, dout_val);
    }
    if (dintermediate_ != nullptr) {
      dintermediate_[i] = dintermediate_op_.UseIntermediateOut(
          x_val, intermediate_out_val, out_val, dout_val);
    }
  }

  const T *x_;
  const T *y_;
  const T *intermediate_out_;
  const T *out_;
  const T *dout_;
  DX_OP dx_op_;
  DY_OP dy_op_;
  DIntermediate_OP dintermediate_op_;
  T *dx_;
  T *dy_;
  T *dintermediate_;
};

template <typename DeviceContext,
          typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut>
void FusedElemwiseAndActGradComputeNoBroadcast(
    const DeviceContext &dev_ctx,
    const DDim &x_dim,
    const DDim &y_dim UNUSED,
    const DenseTensor *x,
    const DenseTensor *y,
    const DenseTensor *intermediate_out,
    const DenseTensor *out,
    const DenseTensor *dout,
    int axis UNUSED,
    DenseTensor *dx,
    DenseTensor *dy,
    DenseTensor *dintermediate,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op) {
  size_t N = static_cast<size_t>(common::product(x_dim));
  funcs::ForRange<DeviceContext> for_range(dev_ctx, N);
  const T *x_data = nullptr;
  const T *y_data = nullptr;
  if (x->IsInitialized()) x_data = x->data<T>();
  if (y->IsInitialized()) y_data = y->data<T>();

  for_range(FusedElemwiseAndActGradNoBroadcast<T,
                                               DX_OP,
                                               DY_OP,
                                               DIntermediate_OP,
                                               UseIntermediateOut>{
      x_data,
      y_data,
      intermediate_out ? intermediate_out->data<T>() : nullptr,
      out->data<T>(),
      dout->data<T>(),
      dx_op,
      dy_op,
      dintermediate_op,
      dx == nullptr ? nullptr : dev_ctx.template Alloc<T>(dx),
      dy == nullptr ? nullptr : dev_ctx.template Alloc<T>(dy),
      dintermediate == nullptr ? nullptr
                               : dev_ctx.template Alloc<T>(dintermediate)});
}

template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActGradBroadcast1CPU(
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int h,
    int w,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *d_intermediate) {
  int64_t tmp_out_idx, x_idx, y_idx;
  T zero = static_cast<T>(0);
  for (int i = 0; i < h; ++i) {
    for (int j = 0; j < w; ++j) {
      int64_t offset = static_cast<int64_t>(i) * w + j;

      tmp_out_idx = BcastY ? j : offset;
      y_idx = BcastY ? j : offset;
      x_idx = BcastY ? offset : j;
      T x_val = (x == nullptr) ? zero : x[x_idx];
      T y_val = (y == nullptr) ? zero : y[y_idx];

      if (SameShapeOfIntermediateOutAndOut) {
        tmp_out_idx = offset;
      }

      if (dx != nullptr) {
        T tmp = UseIntermediateOut
                    ? dx_op.UseIntermediateOut(x_val,
                                               y_val,
                                               intermediate_out[tmp_out_idx],
                                               out[offset],
                                               dout[offset])
                    : dx_op.Recompute(x_val, y_val, out[offset], dout[offset]);

        if (BcastY) {
          dx[x_idx] = tmp;
        } else {
          if (i == 0) {
            dx[x_idx] = tmp;
          } else {
            dx[x_idx] += tmp;
          }
        }
      }
      if (dy != nullptr) {
        T tmp = UseIntermediateOut
                    ? dy_op.UseIntermediateOut(x_val,
                                               y_val,
                                               intermediate_out[tmp_out_idx],
                                               out[offset],
                                               dout[offset])
                    : dy_op.Recompute(x_val, y_val, out[offset], dout[offset]);
        if (BcastY) {
          if (i == 0) {
            dy[y_idx] = tmp;
          } else {
            dy[y_idx] += tmp;
          }
        } else {
          dy[y_idx] = tmp;
        }
      }
      if (d_intermediate != nullptr) {
        T tmp = UseIntermediateOut ? dintermediate_op.UseIntermediateOut(
                                         x_val,
                                         intermediate_out[tmp_out_idx],
                                         out[offset],
                                         dout[offset])
                                   : dintermediate_op.Recompute(
                                         x_val, y_val, out[offset], dout[i]);
        if (SameShapeOfIntermediateOutAndOut) {
          d_intermediate[tmp_out_idx] = tmp;
        } else {
          if (i == 0) {
            d_intermediate[tmp_out_idx] = tmp;
          } else {
            d_intermediate[tmp_out_idx] += tmp;
          }
        }
      }
    }
  }
}

template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActGradBroadcast2CPU(
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int pre,
    int n,
    int post,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *d_intermediate) {
  int64_t tmp_out_idx, x_idx, y_idx;
  T zero = static_cast<T>(0);
  for (int i = 0; i < pre; ++i) {
    for (int j = 0; j < n; ++j) {
      for (int k = 0; k < post; ++k) {
        int64_t offset = static_cast<int64_t>(i) * n * post +
                         static_cast<int64_t>(j) * post + k;

        tmp_out_idx = BcastY ? j : offset;
        y_idx = BcastY ? j : offset;
        x_idx = BcastY ? offset : j;

        T x_val = (x == nullptr) ? zero : x[x_idx];
        T y_val = (y == nullptr) ? zero : y[y_idx];

        if (SameShapeOfIntermediateOutAndOut) {
          tmp_out_idx = offset;
        }

        if (dx != nullptr) {
          T tmp =
              UseIntermediateOut
                  ? dx_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dx_op.Recompute(x_val, y_val, out[offset], dout[offset]);

          if (BcastY) {
            dx[x_idx] = tmp;
          } else {
            if (i == 0 && k == 0) {
              dx[x_idx] = tmp;
            } else {
              dx[x_idx] += tmp;
            }
          }
        }
        if (dy != nullptr) {
          T tmp =
              UseIntermediateOut
                  ? dy_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dy_op.Recompute(x_val, y_val, out[offset], dout[offset]);
          if (BcastY) {
            if (i == 0 && k == 0) {
              dy[y_idx] = tmp;
            } else {
              dy[y_idx] += tmp;
            }
          } else {
            dy[y_idx] = tmp;
          }
        }
        if (d_intermediate != nullptr) {
          T tmp = UseIntermediateOut ? dintermediate_op.UseIntermediateOut(
                                           x_val,
                                           intermediate_out[tmp_out_idx],
                                           out[offset],
                                           dout[offset])
                                     : dintermediate_op.Recompute(
                                           x_val, y_val, out[offset], dout[i]);
          if (SameShapeOfIntermediateOutAndOut) {
            d_intermediate[tmp_out_idx] = tmp;
          } else {
            if (i == 0) {
              d_intermediate[tmp_out_idx] = tmp;
            } else {
              d_intermediate[tmp_out_idx] += tmp;
            }
          }
        }
      }
    }
  }
}

#if defined(__NVCC__) || defined(__HIPCC__)
template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static __global__ void FusedElemwiseAndActGradBroadcast1CUDAKernel(
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int h,
    int w,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *d_intermediate) {
  __shared__ T sdata[BLOCK_Y][BLOCK_X];
  size_t idx = static_cast<size_t>(threadIdx.x) +
               BLOCK_X * static_cast<size_t>(blockIdx.x);
  size_t width_stride = static_cast<size_t>(gridDim.x) * BLOCK_X;

  size_t full_w = ROUNDUP(w, BLOCK_X);

  T zero = static_cast<T>(0);

  for (size_t j = idx; j < full_w; j += width_stride) {
    T val(0), inter_val(0);
    if (j < w) {
      for (size_t i = threadIdx.y; i < h; i += BLOCK_Y) {
        size_t offset = i * w + j;

        size_t tmp_out_idx = BcastY ? j : offset;
        size_t y_idx = BcastY ? j : offset;
        size_t x_idx = BcastY ? offset : j;
        T x_val = (x == nullptr) ? zero : x[x_idx];
        T y_val = (y == nullptr) ? zero : y[y_idx];

        if (SameShapeOfIntermediateOutAndOut) {
          tmp_out_idx = offset;
        }

        if (dx != nullptr) {
          T tmp =
              UseIntermediateOut
                  ? dx_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dx_op.Recompute(x_val, y_val, out[offset], dout[offset]);

          if (BcastY) {
            dx[x_idx] = tmp;
          } else {
            val += tmp;
          }
        }
        if (dy != nullptr) {
          T tmp =
              UseIntermediateOut
                  ? dy_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dy_op.Recompute(x_val, y_val, out[offset], dout[offset]);
          if (BcastY) {
            val += tmp;
          } else {
            dy[y_idx] = tmp;
          }
        }
        if (d_intermediate != nullptr) {
          T tmp = UseIntermediateOut
                      ? dintermediate_op.UseIntermediateOut(
                            y[y_idx],
                            intermediate_out[tmp_out_idx],
                            out[offset],
                            dout[offset])
                      : dintermediate_op.Recompute(
                            x_val, y_val, out[offset], dout[offset]);
          if (SameShapeOfIntermediateOutAndOut) {
            d_intermediate[tmp_out_idx] = tmp;
          } else {
            inter_val += tmp;
          }
        }
      }
    }

    // transpose, for ReduceSum with wrap
    sdata[threadIdx.y][threadIdx.x] = val;
    __syncthreads();
    val = sdata[threadIdx.x][threadIdx.y];
#pragma unroll
    for (int i = BLOCK_X >> 1; i > 0; i >>= 1) {
      // reduce sum with wrap
      val += phi::backends::gpu::CudaShuffleXorSync(0xFFFFFFFF, val, i);
    }

    size_t idx_j = j + static_cast<size_t>(threadIdx.y);
    if (BcastY) {
      if (dy) {
        if (threadIdx.x == 0 && (idx_j < w)) dy[idx_j] = val;
      }
    } else {
      if (dx) {
        if (threadIdx.x == 0 && (idx_j < w)) dx[idx_j] = val;
      }
    }

    if (!SameShapeOfIntermediateOutAndOut) {
      if (d_intermediate) {
        sdata[threadIdx.y][threadIdx.x] = inter_val;
        __syncthreads();
        inter_val = sdata[threadIdx.x][threadIdx.y];
#pragma unroll
        for (int i = BLOCK_X >> 1; i > 0; i >>= 1) {
          // reduce sum with wrap
          inter_val +=
              phi::backends::gpu::CudaShuffleXorSync(0xFFFFFFFF, inter_val, i);
        }
        if (threadIdx.x == 0 && (idx_j < w)) d_intermediate[idx_j] = inter_val;
      }
    }
  }  // end for
}

template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActGradBroadcast1CUDA(
    const GPUContext &dev_ctx,
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int h,
    int w,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *d_intermediate) {
  gpuStream_t stream = dev_ctx.stream();

  dim3 blocks(BLOCK_X, BLOCK_Y);
  int max_gpu_threads = dev_ctx.GetMaxPhysicalThreadCount();
  int max_blocks = std::max(max_gpu_threads / (BLOCK_X * BLOCK_Y), 1);
  int theory_block = (w + BLOCK_X - 1) / BLOCK_X;
  dim3 grids(std::min(theory_block, max_blocks));

  FusedElemwiseAndActGradBroadcast1CUDAKernel<T,
                                              DX_OP,
                                              DY_OP,
                                              DIntermediate_OP,
                                              UseIntermediateOut,
                                              BcastY,
                                              SameShapeOfIntermediateOutAndOut>
      <<<grids, blocks, 0, stream>>>(x,
                                     y,
                                     intermediate_out,
                                     out,
                                     dout,
                                     h,
                                     w,
                                     dx_op,
                                     dy_op,
                                     dintermediate_op,
                                     dx,
                                     dy,
                                     d_intermediate);
}

template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static __global__ void FusedElemwiseAndActGradBroadcast2CUDAKernel(
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int pre,
    int n,
    int post,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *d_intermediate) {
  int tid = threadIdx.x;
  int j = blockIdx.x;

  T val(0), inter_val(0);
  int ttid = tid;
  int64_t tmp_out_idx, x_idx, y_idx;
  T zero = static_cast<T>(0);
  while (true) {
    int i = ttid / post;
    int k = ttid % post;
    if (i >= pre) break;

    int64_t offset =
        static_cast<int64_t>(i) * n * post + static_cast<int64_t>(j) * post + k;

    tmp_out_idx = BcastY ? j : offset;
    y_idx = BcastY ? j : offset;
    x_idx = BcastY ? offset : j;
    T x_val = (x == nullptr) ? zero : x[x_idx];
    T y_val = (y == nullptr) ? zero : y[y_idx];

    if (SameShapeOfIntermediateOutAndOut) {
      tmp_out_idx = offset;
    }

    if (dx != nullptr) {
      T tmp = UseIntermediateOut
                  ? dx_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dx_op.Recompute(x_val, y_val, out[offset], dout[offset]);

      if (BcastY) {
        dx[x_idx] = tmp;
      } else {
        val += tmp;
      }
    }
    if (dy != nullptr) {
      T tmp = UseIntermediateOut
                  ? dy_op.UseIntermediateOut(x_val,
                                             y_val,
                                             intermediate_out[tmp_out_idx],
                                             out[offset],
                                             dout[offset])
                  : dy_op.Recompute(x_val, y_val, out[offset], dout[offset]);
      if (BcastY) {
        val += tmp;
      } else {
        dy[y_idx] = tmp;
      }
    }
    if (d_intermediate != nullptr) {
      T tmp = UseIntermediateOut ? dintermediate_op.UseIntermediateOut(
                                       y_val,
                                       intermediate_out[tmp_out_idx],
                                       out[offset],
                                       dout[offset])
                                 : dintermediate_op.Recompute(
                                       x_val, y_val, out[offset], dout[offset]);
      if (SameShapeOfIntermediateOutAndOut) {
        d_intermediate[tmp_out_idx] = tmp;
      } else {
        inter_val += tmp;
      }
    }
    ttid += ELEMWISE_MAX_BLOCK_DIM;
  }

  int h = pre * post;
  h = h > ELEMWISE_MAX_BLOCK_DIM ? ELEMWISE_MAX_BLOCK_DIM : h;
  if (BcastY) {
    if (dy) {
      val = phi::backends::gpu::reduceSum(val, tid, h);
      if (threadIdx.x == 0) {
        dy[j] = val;
      }
    }
  } else {
    if (dx) {
      val = phi::backends::gpu::reduceSum(val, tid, h);
      if (threadIdx.x == 0) {
        dx[j] = val;
      }
    }
  }
  if (!SameShapeOfIntermediateOutAndOut) {
    if (d_intermediate) {
      inter_val = phi::backends::gpu::reduceSum(inter_val, tid, h);
      if (threadIdx.x == 0) {
        d_intermediate[j] = inter_val;
      }
    }
  }
}

template <typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
static void FusedElemwiseAndActGradBroadcast2CUDA(
    gpuStream_t stream,
    const T *x,
    const T *y,
    const T *intermediate_out,
    const T *out,
    const T *dout,
    int64_t pre,
    int64_t n,
    int64_t post,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op,
    T *dx,
    T *dy,
    T *dintermediate) {
  int64_t block_size =
      std::min(static_cast<int64_t>(ELEMWISE_MAX_BLOCK_DIM), pre * post);
  int64_t gird_size = n;
  FusedElemwiseAndActGradBroadcast2CUDAKernel<T,
                                              DX_OP,
                                              DY_OP,
                                              DIntermediate_OP,
                                              UseIntermediateOut,
                                              BcastY,
                                              SameShapeOfIntermediateOutAndOut>
      <<<gird_size, block_size, 0, stream>>>(x,
                                             y,
                                             intermediate_out,
                                             out,
                                             dout,
                                             pre,
                                             n,
                                             post,
                                             dx_op,
                                             dy_op,
                                             dintermediate_op,
                                             dx,
                                             dy,
                                             dintermediate);
}
#endif

template <typename DeviceContext,
          typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool BcastY,
          bool SameShapeOfIntermediateOutAndOut>
void FusedElemwiseAndActGradComputeWithBroadcast(
    const DeviceContext &dev_ctx,
    const DDim &x_dim,
    const DDim &y_dim_untrimed,
    const DenseTensor *x,
    const DenseTensor *y,
    const DenseTensor *intermediate_out,
    const DenseTensor *out,
    const DenseTensor *dout,
    int axis,
    DenseTensor *dx,
    DenseTensor *dy,
    DenseTensor *dintermediate,
    DX_OP dx_op,
    DY_OP dy_op,
    DIntermediate_OP dintermediate_op) {
  axis = (axis == -1 ? x_dim.size() - y_dim_untrimed.size() : axis);
  auto y_dim = funcs::TrimTrailingSingularDims(y_dim_untrimed);
  axis = (y_dim.size() == 0) ? x_dim.size() : axis;

  size_t pre, n, post;
  int is_run_common_broadcast;
  funcs::GetMidDims(
      x_dim, y_dim, axis, &pre, &n, &post, &is_run_common_broadcast);
  const T *x_data = nullptr;
  const T *y_data = nullptr;
  if (x->IsInitialized()) x_data = x->data<T>();
  if (y->IsInitialized()) y_data = y->data<T>();
  if (post == 1) {
    int h = pre;
    int w = n;

    if (dev_ctx.GetPlace().GetType() == AllocationType::GPU) {
#if defined(__NVCC__) || defined(__HIPCC__)
      FusedElemwiseAndActGradBroadcast1CUDA<T,
                                            DX_OP,
                                            DY_OP,
                                            DIntermediate_OP,
                                            UseIntermediateOut,
                                            BcastY,
                                            SameShapeOfIntermediateOutAndOut>(
          reinterpret_cast<const GPUContext &>(dev_ctx),
          x_data,
          y_data,
          intermediate_out == nullptr ? nullptr : intermediate_out->data<T>(),
          out->data<T>(),
          dout->data<T>(),
          h,
          w,
          dx_op,
          dy_op,
          dintermediate_op,
          dx == nullptr ? nullptr : dev_ctx.template Alloc<T>(dx),
          dy == nullptr ? nullptr : dev_ctx.template Alloc<T>(dy),
          dintermediate == nullptr ? nullptr
                                   : dev_ctx.template Alloc<T>(dintermediate));
#endif
    } else {
      FusedElemwiseAndActGradBroadcast1CPU<T,
                                           DX_OP,
                                           DY_OP,
                                           DIntermediate_OP,
                                           UseIntermediateOut,
                                           BcastY,
                                           SameShapeOfIntermediateOutAndOut>(
          x_data,
          y_data,
          intermediate_out == nullptr ? nullptr : intermediate_out->data<T>(),
          out->data<T>(),
          dout->data<T>(),
          h,
          w,
          dx_op,
          dy_op,
          dintermediate_op,
          dx == nullptr ? nullptr : dev_ctx.template Alloc<T>(dx),
          dy == nullptr ? nullptr : dev_ctx.template Alloc<T>(dy),
          dintermediate == nullptr ? nullptr
                                   : dev_ctx.template Alloc<T>(dintermediate));
    }
  } else {
    if (dev_ctx.GetPlace().GetType() == AllocationType::GPU) {
#if defined(__NVCC__) || defined(__HIPCC__)
      FusedElemwiseAndActGradBroadcast2CUDA<T,
                                            DX_OP,
                                            DY_OP,
                                            DIntermediate_OP,
                                            UseIntermediateOut,
                                            BcastY,
                                            SameShapeOfIntermediateOutAndOut>(
          reinterpret_cast<const GPUContext &>(dev_ctx).stream(),
          x_data,
          y_data,
          intermediate_out == nullptr ? nullptr : intermediate_out->data<T>(),
          out->data<T>(),
          dout->data<T>(),
          pre,
          n,
          post,
          dx_op,
          dy_op,
          dintermediate_op,
          dx == nullptr ? nullptr : dev_ctx.template Alloc<T>(dx),
          dy == nullptr ? nullptr : dev_ctx.template Alloc<T>(dy),
          dintermediate == nullptr ? nullptr
                                   : dev_ctx.template Alloc<T>(dintermediate));
#endif
    } else {
      FusedElemwiseAndActGradBroadcast2CPU<T,
                                           DX_OP,
                                           DY_OP,
                                           DIntermediate_OP,
                                           UseIntermediateOut,
                                           BcastY,
                                           SameShapeOfIntermediateOutAndOut>(
          x_data,
          y_data,
          intermediate_out == nullptr ? nullptr : intermediate_out->data<T>(),
          out->data<T>(),
          dout->data<T>(),
          pre,
          n,
          post,
          dx_op,
          dy_op,
          dintermediate_op,
          dx == nullptr ? nullptr : dev_ctx.template Alloc<T>(dx),
          dy == nullptr ? nullptr : dev_ctx.template Alloc<T>(dy),
          dintermediate == nullptr ? nullptr
                                   : dev_ctx.template Alloc<T>(dintermediate));
    }
  }
}

template <typename DeviceContext,
          typename T,
          typename DX_OP,
          typename DY_OP,
          typename DIntermediate_OP,
          bool UseIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
void FusedElemwiseAndActGradComputeEx(const DeviceContext &dev_ctx,
                                      const DenseTensor *x,
                                      const DenseTensor *y,
                                      const DenseTensor *out,
                                      const DenseTensor *intermediate_out,
                                      const DenseTensor *dout,
                                      int axis,
                                      DenseTensor *dx,
                                      DenseTensor *dy,
                                      DenseTensor *dintermediate,
                                      DX_OP dx_op,
                                      DY_OP dy_op,
                                      DIntermediate_OP dintermediate_op) {
  const DDim &x_dim = x->dims();
  const DDim &y_dim = y->dims();
  if (UseIntermediateOut) {
    PADDLE_ENFORCE_NOT_NULL(
        intermediate_out,
        common::errors::InvalidArgument("Intermediate out is null pointer."));
  }
  if (x_dim == y_dim) {
    FusedElemwiseAndActGradComputeNoBroadcast<DeviceContext,
                                              T,
                                              DX_OP,
                                              DY_OP,
                                              DIntermediate_OP,
                                              UseIntermediateOut>(
        dev_ctx,
        x_dim,
        y_dim,
        x,
        y,
        intermediate_out,
        out,
        dout,
        axis,
        dx,
        dy,
        dintermediate,
        dx_op,
        dy_op,
        dintermediate_op);
  } else {  // Y is a scalar
    bool bcast_y = x_dim.size() >= y_dim.size();
    if (x_dim.size() == y_dim.size()) {
      for (int i = 0; i < x_dim.size(); ++i) {
        if (x_dim[i] < y_dim[i]) {
          bcast_y = false;
          break;
        }
      }
    }

    // z = f1(x, f2(y))
    // z = f1(f2(x, y))
    if (bcast_y) {  // Y should be broadcast.
      FusedElemwiseAndActGradComputeWithBroadcast<
          DeviceContext,
          T,
          DX_OP,
          DY_OP,
          DIntermediate_OP,
          UseIntermediateOut,
          true /*BcastY*/,
          SameShapeOfIntermediateOutAndOut>(dev_ctx,
                                            x_dim,
                                            y_dim,
                                            x,
                                            y,
                                            intermediate_out,
                                            out,
                                            dout,
                                            axis,
                                            dx,
                                            dy,
                                            dintermediate,
                                            dx_op,
                                            dy_op,
                                            dintermediate_op);
    } else {
      FusedElemwiseAndActGradComputeWithBroadcast<
          DeviceContext,
          T,
          DX_OP,
          DY_OP,
          DIntermediate_OP,
          UseIntermediateOut,
          false /*BcastY*/,
          SameShapeOfIntermediateOutAndOut>(dev_ctx,
                                            y_dim,
                                            x_dim,
                                            x,
                                            y,
                                            intermediate_out,
                                            out,
                                            dout,
                                            axis,
                                            dx,
                                            dy,
                                            dintermediate,
                                            dx_op,
                                            dy_op,
                                            dintermediate_op);
    }
  }
}

template <typename DeviceContext,
          typename T,
          typename CompoundFunctor,
          bool KeepIntermediateOut,
          bool SameShapeOfIntermediateOutAndOut>
void FusedElemwiseAndActComputeEx(const DeviceContext &dev_ctx,
                                  const DenseTensor &x,
                                  const DenseTensor &y,
                                  int axis,
                                  CompoundFunctor compound_functor,
                                  DenseTensor *out,
                                  DenseTensor *intermediate_out) {
  if (KeepIntermediateOut) {
    PADDLE_ENFORCE_NOT_NULL(
        intermediate_out,
        common::errors::InvalidArgument(
            "The save_intermediate_out is opened, intermediate "
            "out is null pointer."));
  }

  const DDim &x_dim = x.dims();
  const DDim &y_dim = y.dims();
  if (x.dims() == y.dims()) {
    FusedElemwiseAndActComputeNoBroadcast<DeviceContext,
                                          T,
                                          CompoundFunctor,
                                          KeepIntermediateOut>(
        dev_ctx, x_dim, x, y, compound_functor, out, intermediate_out);
  } else {
    // Whether the shape of Y is a continuous subsequence of X,
    // For more information please refer to the op's introduction.
    bool bcast_y = x.numel() >= y.numel();
    // z = f1(x, f2(y))
    // z = f1(f2(x, y))
    if (bcast_y) {  // Y should be broadcast.
      // In this case,
      // for 'f2(y)', the shape of intermediate_out should be equal to the
      // shape
      // of Y.
      // for 'f2(x, y)', the shape of intermediate_out should be equal to the
      // shape of Out.
      // the shape of Out should be equal to the shape of X.
      FusedElemwiseAndActComputeWithBroadcast<DeviceContext,
                                              T,
                                              CompoundFunctor,
                                              true /*BcastY*/,
                                              KeepIntermediateOut,
                                              SameShapeOfIntermediateOutAndOut>(
          dev_ctx,
          x_dim /*OutShape*/,
          y_dim,
          x,
          y,
          compound_functor,
          axis,
          out,
          intermediate_out);
    } else {
      // In this case,
      // for 'f2(y)', the shape of intermediate_out should be equal to the
      // shape
      // of Out.
      // for 'f2(x, y)', the shape of intermediate_out should be equal to the
      // shape of Out.
      // the shape of Out should be equal to the shape of Y.
      FusedElemwiseAndActComputeWithBroadcast<DeviceContext,
                                              T,
                                              CompoundFunctor,
                                              false /*BcastY*/,
                                              KeepIntermediateOut,
                                              SameShapeOfIntermediateOutAndOut>(
          dev_ctx,
          y_dim /*OutShape*/,
          x_dim,
          x,
          y,
          compound_functor,
          axis,
          out,
          intermediate_out);
    }
  }
}

template <typename DeviceContext, typename T>
static inline void GetDoubleGradSafeTensor(const DeviceContext &dev_ctx,
                                           const DenseTensor *x,
                                           const DenseTensor *ddx,
                                           DenseTensor *ddx_safe) {
  funcs::GetDoubleGradSafeTensor<DeviceContext, T>(dev_ctx, *x, ddx, ddx_safe);
}

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

template <typename T, typename Functor>
void GetGradXAndYOut(const GPUContext &dev_ctx,
                     const phi::Place &place,
                     int axis,
                     std::vector<const DenseTensor *> ins,
                     const DenseTensor *dout,
                     DenseTensor *dx,
                     DenseTensor *dy,
                     Functor func) {
  phi::GetGradXAndYOut<T, Functor>(
      dev_ctx, place, axis, ins, *dout, dx, dy, func);
}

template <typename T, typename Functor>
void GetGradXOrYOut(const GPUContext &dev_ctx,
                    const phi::Place &place,
                    int axis,
                    std::vector<const DenseTensor *> ins,
                    const DenseTensor *dout,
                    DenseTensor *dxy,
                    Functor func) {
  phi::GetGradXOrYOut<T, Functor>(dev_ctx, place, axis, ins, *dout, dxy, func);
}

#endif

}  // namespace funcs
}  // namespace phi