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

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

#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"

namespace phi {
template <typename T, typename MT>
__global__ void RAdamGPUKernel(const T* param,
                               const T* grad,
                               const MT* learning_rate,
                               const MT* beta1_pow,
                               const MT* beta2_pow,
                               const MT* rho,
                               const MT* moment1,
                               const MT* moment2,
                               const MT* master_param,
                               MT beta1,
                               MT beta2,
                               MT epsilon,
                               MT rho_inf,
                               int64_t num,
                               T* param_out,
                               MT* beta1_pow_out,
                               MT* beta2_pow_out,
                               MT* rho_out,
                               MT* moment1_out,
                               MT* moment2_out,
                               MT* master_param_out) {
  MT lr_scalar = static_cast<MT>(learning_rate[0]);

  int64_t idx =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);

  for (int64_t index = idx; index < num; index += gridDim.x * blockDim.x) {
    // load and cast input to MT
    MT d_param =
        master_param ? master_param[index] : static_cast<MT>(param[index]);
    MT d_grad = static_cast<MT>(grad[index]);
    MT d_beta1_pow = static_cast<MT>(beta1_pow[index]);
    MT d_beta2_pow = static_cast<MT>(beta2_pow[index]);
    MT d_rho = static_cast<MT>(rho[index]);
    MT d_moment1 = static_cast<MT>(moment1[index]);
    MT d_moment2 = static_cast<MT>(moment2[index]);

    // compute
    MT beta1_pow_scalar = d_beta1_pow * beta1;
    MT beta2_pow_scalar = d_beta2_pow * beta2;
    MT rho_scalar = (d_rho * (beta2 - beta2_pow_scalar) + beta2_pow_scalar) /
                    (static_cast<MT>(1) - beta2_pow_scalar);

    MT m1_out = beta1 * d_moment1 + (static_cast<MT>(1) - beta1) * d_grad;
    MT m2_out =
        beta2 * d_moment2 + (static_cast<MT>(1) - beta2) * d_grad * d_grad;

    MT m1_hat = m1_out / (static_cast<MT>(1) - beta1_pow_scalar);

    MT rho_t = rho_inf - static_cast<MT>(2) * rho_scalar;
    MT p_out = static_cast<MT>(0);

    if (rho_t > static_cast<MT>(5)) {
      MT l_t = std::sqrt((static_cast<MT>(1) - beta2_pow_scalar)) /
               (std::sqrt(m2_out) + epsilon);
      MT r_t = std::sqrt(((rho_t - static_cast<MT>(4)) *
                          (rho_t - static_cast<MT>(2)) * rho_inf) /
                         ((rho_inf - static_cast<MT>(4)) *
                          (rho_inf - static_cast<MT>(2)) * rho_t));

      p_out = d_param - lr_scalar * m1_hat * r_t * l_t;
    } else {
      p_out = d_param - lr_scalar * m1_hat;
    }

    // store
    param_out[index] = static_cast<T>(p_out);
    beta1_pow_out[index] = static_cast<T>(beta1_pow_scalar);
    beta2_pow_out[index] = static_cast<T>(beta2_pow_scalar);
    rho_out[index] = static_cast<T>(rho_scalar);
    moment1_out[index] = static_cast<MT>(m1_out);
    moment2_out[index] = static_cast<MT>(m2_out);

    if (master_param_out) {
      master_param_out[index] = p_out;
    }
  }
}

template <typename T, typename Context>
void RAdamKernel(const Context& dev_ctx,
                 const DenseTensor& param,
                 const DenseTensor& grad,
                 const DenseTensor& learning_rate,
                 const DenseTensor& beta1_pow,
                 const DenseTensor& beta2_pow,
                 const DenseTensor& rho,
                 const DenseTensor& moment1,
                 const DenseTensor& moment2,
                 const optional<DenseTensor>& master_param,
                 float beta1,
                 float beta2,
                 float epsilon,
                 bool multi_precision,
                 DenseTensor* param_out,
                 DenseTensor* beta1_pow_out,
                 DenseTensor* beta2_pow_out,
                 DenseTensor* rho_out,
                 DenseTensor* moment1_out,
                 DenseTensor* moment2_out,
                 DenseTensor* master_param_out) {
  using MT = typename dtype::template MPTypeTrait<T>::Type;
  T* param_out_data = dev_ctx.template Alloc<T>(param_out);

  MT* beta1_pow_out_data = dev_ctx.template Alloc<MT>(beta1_pow_out);
  MT* beta2_pow_out_data = dev_ctx.template Alloc<MT>(beta2_pow_out);
  MT* rho_out_data = dev_ctx.template Alloc<MT>(rho_out);

  MT* moment1_out_data = dev_ctx.template Alloc<MT>(moment1_out);
  MT* moment2_out_data = dev_ctx.template Alloc<MT>(moment2_out);

  const MT* master_in_data =
      multi_precision ? master_param->data<MT>() : nullptr;
  MT* master_out_data =
      multi_precision ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;

  MT beta1_ = static_cast<MT>(beta1);
  MT beta2_ = static_cast<MT>(beta2);
  MT epsilon_ = static_cast<MT>(epsilon);

  MT rho_inf =
      static_cast<MT>(2) / (static_cast<MT>(1) - beta2_) - static_cast<MT>(1);

  int64_t numel = param.numel();
  int block = 512;
  int64_t grid_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
  int grid = std::min((param.numel() + block - 1) / block, grid_max);
  auto stream = dev_ctx.stream();

  RAdamGPUKernel<T, MT><<<block, grid, 0, stream>>>(param.data<T>(),
                                                    grad.data<T>(),
                                                    learning_rate.data<MT>(),
                                                    beta1_pow.data<MT>(),
                                                    beta2_pow.data<MT>(),
                                                    rho.data<MT>(),
                                                    moment1.data<MT>(),
                                                    moment2.data<MT>(),
                                                    master_in_data,
                                                    beta1_,
                                                    beta2_,
                                                    epsilon_,
                                                    rho_inf,
                                                    numel,
                                                    param_out_data,
                                                    beta1_pow_out_data,
                                                    beta2_pow_out_data,
                                                    rho_out_data,
                                                    moment1_out_data,
                                                    moment2_out_data,
                                                    master_out_data);
}
}  // namespace phi

PD_REGISTER_KERNEL(
    radam, GPU, ALL_LAYOUT, phi::RAdamKernel, float, double, phi::float16) {}