// 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.

#include "paddle/phi/kernels/lerp_grad_kernel.h"

#include "paddle/common/enforce.h"
#include "paddle/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/kernel_registry.h"

#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/kernels/broadcast_tensors_kernel.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/common_shape.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/phi/kernels/funcs/reduce_function.h"
#include "paddle/phi/kernels/gpu/reduce.h"
#include "paddle/phi/kernels/reduce_sum_kernel.h"

COMMON_DECLARE_bool(use_accuracy_compatible_kernel);

namespace phi {

template <typename T>
__global__ void LerpGradKernelImpl(const T* weight,
                                   const T* dout,
                                   T* dx,
                                   T* dy,
                                   const int64_t out_size,
                                   const int64_t x_size,
                                   const int64_t y_size) {
  using MT = typename MPTypeTrait<T>::Type;
  CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
    MT temp_dx = static_cast<MT>(weight[idx]) * static_cast<MT>(dout[idx]);
    if (dx) {
      if (idx < x_size) {
        dx[idx] = static_cast<T>(static_cast<MT>(dout[idx]) - temp_dx);
      }
    }
    if (dy) {
      if (idx < y_size) {
        dy[idx] = static_cast<T>(temp_dx);
      }
    }
  }
}

template <typename T>
__global__ void LerpGradKernelCompatibleImpl(const T* weight,
                                             const T* dout,
                                             T* dx,
                                             T* dy,
                                             const int64_t out_size,
                                             const int64_t x_size,
                                             const int64_t y_size) {
  CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
    T weight_value = weight[idx];
    T remaining_weight_value = static_cast<T>(1) - weight[idx];
    if (dx) {
      if (idx < x_size) {
        dx[idx] = remaining_weight_value * dout[idx];
      }
    }
    if (dy) {
      if (idx < y_size) {
        dy[idx] = weight_value * dout[idx];
      }
    }
  }
}

template <typename T, typename WeightT = T>
__global__ void LerpGradScalarKernelImpl(const WeightT* weight,
                                         const T* dout,
                                         T* dx,
                                         T* dy,
                                         const int64_t out_size,
                                         const int64_t x_size,
                                         const int64_t y_size) {
  double weight_scalar = static_cast<double>(weight[0]);
  CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
    double temp_dx = weight_scalar * static_cast<double>(dout[idx]);
    if (dx) {
      if (idx < x_size) {
        dx[idx] = static_cast<T>(static_cast<double>(dout[idx]) - temp_dx);
      }
    }
    if (dy) {
      if (idx < y_size) {
        dy[idx] = static_cast<T>(temp_dx);
      }
    }
  }
}

template <typename T, typename WeightT = T>
__global__ void LerpGradScalarKernelCompatibleImpl(const WeightT* weight,
                                                   const T* dout,
                                                   T* dx,
                                                   T* dy,
                                                   const int64_t out_size,
                                                   const int64_t x_size,
                                                   const int64_t y_size) {
  T weight_scalar = static_cast<T>(weight[0]);
  T remaining_weight_scalar =
      static_cast<T>(1 - static_cast<double>(weight[0]));
  CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
    if (dx) {
      if (idx < x_size) {
        dx[idx] = remaining_weight_scalar * dout[idx];
      }
    }
    if (dy) {
      if (idx < y_size) {
        dy[idx] = weight_scalar * dout[idx];
      }
    }
  }
}

bool XYNeedReduce(const DenseTensor& x,
                  const DenseTensor& y,
                  const DenseTensor& out) {
  auto x_dims =
      x.dims().size() ? x.dims() : make_ddim(std::vector<int64_t>(1, 1));
  auto y_dims =
      y.dims().size() ? y.dims() : make_ddim(std::vector<int64_t>(1, 1));

  auto out_dims = out.dims();
  if (out_dims.size() == 0) {
    return false;
  }
  int x_rank = x_dims.size();
  int y_rank = y_dims.size();
  int out_rank = out_dims.size();
  int smaller_rank = std::min(x_rank, y_rank);
  if (std::max(x_rank, y_rank) < out_rank) {
    return true;
  }
  for (int i = 1; i <= smaller_rank; ++i) {
    int x_idx = x_rank - i;
    int y_idx = y_rank - i;
    int out_idx = out_rank - i;
    if (x_dims[x_idx] != y_dims[y_idx]) {
      return true;
    }
    if (x_dims[x_idx] == 1 && y_dims[y_idx] == 1 && out_dims[out_idx] != 1) {
      return true;
    }
  }
  return false;
}

template <typename T, typename Context>
void SwitchKernel(const Context& dev_ctx,
                  const DenseTensor& weight,
                  const DenseTensor& out_grad,
                  const int64_t x_grad_size,
                  const int64_t y_grad_size,
                  T* x_grad_data,
                  T* y_grad_data) {
  if (weight.numel() == 1) {
    //    condition when weight is a scalar
    const T* out_grad_data = out_grad.data<T>();
    const int64_t out_size = out_grad.numel();
    const int64_t weight_size = weight.numel();

    auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_size);
    const size_t grid_size = gpu_config.GetGridSize();
    const size_t block_size = gpu_config.GetBlockSize();
    PADDLE_ENFORCE_LE_UINT32_MAX(grid_size, "grid");
    const uint32_t grid = static_cast<uint32_t>(grid_size);
    const uint32_t block = static_cast<uint32_t>(block_size);

    if (weight.dtype() == DataType::FLOAT64) {
      const double* weight_data = weight.data<double>();
      if (FLAGS_use_accuracy_compatible_kernel) {
        LerpGradScalarKernelCompatibleImpl<T, double>
            <<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                   out_grad_data,
                                                   x_grad_data,
                                                   y_grad_data,
                                                   out_size,
                                                   x_grad_size,
                                                   y_grad_size);
      } else {
        LerpGradScalarKernelImpl<T, double>
            <<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                   out_grad_data,
                                                   x_grad_data,
                                                   y_grad_data,
                                                   out_size,
                                                   x_grad_size,
                                                   y_grad_size);
      }
    } else {
      const T* weight_data = weight.data<T>();
      if (FLAGS_use_accuracy_compatible_kernel) {
        LerpGradScalarKernelCompatibleImpl<T>
            <<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                   out_grad_data,
                                                   x_grad_data,
                                                   y_grad_data,
                                                   out_size,
                                                   x_grad_size,
                                                   y_grad_size);
      } else {
        LerpGradScalarKernelImpl<T>
            <<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                   out_grad_data,
                                                   x_grad_data,
                                                   y_grad_data,
                                                   out_size,
                                                   x_grad_size,
                                                   y_grad_size);
      }
    }
  } else {
    //    broadcast weight with out_grad's dimensions
    const std::vector<const DenseTensor*> in_tensors = {&weight, &out_grad};
    DenseTensor b_weight = EmptyLike<T>(dev_ctx, out_grad);
    DenseTensor b_out = EmptyLike<T>(dev_ctx, out_grad);
    std::vector<DenseTensor*> out_tensors = {&b_weight, &b_out};

    BroadcastTensorsKernel<T, Context>(dev_ctx, in_tensors, out_tensors);

    const T* weight_data = b_weight.data<T>();
    const T* out_grad_data = b_out.data<T>();
    const int64_t out_size = out_grad.numel();
    const int64_t weight_size = weight.numel();
    auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_size);
    const int64_t grid_size = gpu_config.GetGridSize();
    const int64_t block_size = gpu_config.GetBlockSize();
    PADDLE_ENFORCE_LE_UINT32_MAX(grid_size, "grid");
    PADDLE_ENFORCE_LE_UINT32_MAX(block_size, "block");
    const uint32_t grid = static_cast<uint32_t>(grid_size);
    const uint32_t block = static_cast<uint32_t>(block_size);
    if (FLAGS_use_accuracy_compatible_kernel) {
      LerpGradKernelCompatibleImpl<T>
          <<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                 out_grad_data,
                                                 x_grad_data,
                                                 y_grad_data,
                                                 out_size,
                                                 x_grad_size,
                                                 y_grad_size);
    } else {
      LerpGradKernelImpl<T><<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
                                                                  out_grad_data,
                                                                  x_grad_data,
                                                                  y_grad_data,
                                                                  out_size,
                                                                  x_grad_size,
                                                                  y_grad_size);
    }
  }
}

template <typename T, typename Context>
void LerpGradKernel(const Context& dev_ctx,
                    const DenseTensor& x,
                    const DenseTensor& y,
                    const DenseTensor& weight,
                    const DenseTensor& out,
                    const DenseTensor& out_grad,
                    DenseTensor* x_grad,
                    DenseTensor* y_grad) {
  if (out_grad.numel() == 0) {
    if (x_grad) {
      Full<T, Context>(dev_ctx, x_grad->dims(), 0, x_grad);
    }
    if (y_grad) {
      Full<T, Context>(dev_ctx, y_grad->dims(), 0, y_grad);
    }
    return;
  }

  const int rank = out.dims().size();
  PADDLE_ENFORCE_GE(
      rank,
      0,
      common::errors::InvalidArgument(
          "The number of dimensions for LerpGradOp must be "
          "greater than or equal to 0, but the value received is %d.",
          rank));
  PADDLE_ENFORCE_LE(
      rank,
      6,
      common::errors::InvalidArgument(
          "The number of dimensions for LerpGradOp must be "
          "less than or equal to 6, but the value received is %d.",
          rank));

  //  check if x_grad and y_grad need to be reduced
  //  if x has a different dimension with y or weight in the middle axis, then
  //  they need to be broadcast and then reduced.
  bool reduce_flag = XYNeedReduce(x, y, out);
  if (!reduce_flag) {
    int64_t x_grad_size = 0, y_grad_size = 0;
    T* x_grad_data = NULL;
    T* y_grad_data = NULL;

    if (x_grad) {
      x_grad_data = dev_ctx.template Alloc<T>(x_grad);
      x_grad_size = x.numel();
    }

    if (y_grad) {
      y_grad_data = dev_ctx.template Alloc<T>(y_grad);
      y_grad_size = y.numel();
    }

    SwitchKernel<T, Context>(dev_ctx,
                             weight,
                             out_grad,
                             x_grad_size,
                             y_grad_size,
                             x_grad_data,
                             y_grad_data);

  } else {
    int64_t x_grad_size = 0, y_grad_size = 0;
    DenseTensor b_xgrad = EmptyLike<T, Context>(dev_ctx, out_grad);
    DenseTensor b_ygrad = EmptyLike<T, Context>(dev_ctx, out_grad);
    T* x_grad_data = NULL;
    T* y_grad_data = NULL;

    if (x_grad) {
      x_grad_data = dev_ctx.template Alloc<T>(&b_xgrad);
      x_grad_size = out.numel();
    }

    if (y_grad) {
      y_grad_data = dev_ctx.template Alloc<T>(&b_ygrad);
      y_grad_size = out.numel();
    }

    SwitchKernel<T, Context>(dev_ctx,
                             weight,
                             out_grad,
                             x_grad_size,
                             y_grad_size,
                             x_grad_data,
                             y_grad_data);

    auto zero_dim = make_ddim(std::vector<int64_t>(1, 1));
    if (x_grad) {
      std::vector<int> reduce_axis_x =
          funcs::GetReduceDim(x_grad->dims().size() ? x_grad->dims() : zero_dim,
                              b_xgrad.dims(),
                              -1);
      if (!reduce_axis_x.empty()) {
        SumKernel<T, Context>(
            dev_ctx, b_xgrad, reduce_axis_x, b_xgrad.dtype(), false, x_grad);
      } else {
        x_grad->ShareDataWith(b_xgrad);
      }
    }

    if (y_grad) {
      std::vector<int> reduce_axis_y =
          funcs::GetReduceDim(y_grad->dims().size() ? y_grad->dims() : zero_dim,
                              b_ygrad.dims(),
                              -1);
      if (!reduce_axis_y.empty()) {
        SumKernel<T, Context>(
            dev_ctx, b_ygrad, reduce_axis_y, b_ygrad.dtype(), false, y_grad);
      } else {
        y_grad->ShareDataWith(b_ygrad);
      }
    }
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(lerp_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::LerpGradKernel,
                   phi::float16,
                   phi::bfloat16,
                   float,
                   double) {}