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

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

#include <glog/logging.h>

#include "dgc/dgc.h"
#include "paddle/common/enforce.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/broadcast_function.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/math_function.h"

namespace phi {

inline float get_period_sparcity(const std::vector<float>& sparsity,
                                 float cur_step,
                                 float rampup_steps) {
  PADDLE_ENFORCE_GE(
      static_cast<int>(cur_step),
      0,
      common::errors::InvalidArgument("DGC current step=%d, but it must >= 0, "
                                      "please submit issue in github",
                                      static_cast<int>(cur_step)));

  size_t idx = static_cast<int>(cur_step * sparsity.size() / rampup_steps);
  if (idx >= sparsity.size()) {
    idx = sparsity.size() - 1;
  }

  PADDLE_ENFORCE_LT(
      idx,
      sparsity.size(),
      common::errors::OutOfRange(
          "sparsity index out of bounds. idx=%d >= sparsity.size=%d",
          idx,
          sparsity.size()));
  return sparsity[idx];
}

template <typename T, typename Context>
void DGCKernel(const Context& dev_ctx,
               const DenseTensor& u,
               const DenseTensor& v,
               const DenseTensor& grad,
               const DenseTensor& param,
               const DenseTensor& current_step_tensor,
               const DenseTensor& nranks_tensor,
               float m,
               bool use_nesterov,
               const std::vector<float>& sparsity,
               float rampup_begin_step,
               float rampup_step,
               float regular_coeff,
               int regular_type,
               DenseTensor* u_out,
               DenseTensor* v_out,
               DenseTensor* encode_grad_out,
               DenseTensor* grad_out,
               DenseTensor* k_out,
               DenseTensor* gather_buff) {
  // nranks
  const int nranks = static_cast<int>(*nranks_tensor.data<float>());
  PADDLE_ENFORCE_GT(nranks,
                    1,
                    common::errors::PreconditionNotMet(
                        "DGC is not useful when num_trainers <= 1. Please "
                        "use multi card or multi machine GPU"));

  auto param_e = EigenVector<T>::Flatten(param);
  auto grad_e = EigenVector<T>::Flatten(grad);
  auto grad_out_e = EigenVector<T>::Flatten(*grad_out);

  auto& eigen_ctx = *dev_ctx.eigen_device();

  // NOTE. In paddle, loss has divided by nranks. Because dgc_op is before
  // allreduce, so local regular_coeff need div nranks too. But now we
  // multi grad with nranks in dgc_op, in that case regular_coeff don't
  // need to /nranks, can prevent precision loss. For coeff often equal
  // with 1e-4, if nranks=32, coeff/nranks will be 3.125e-6, the numerical
  // accuracy of coeff/nranks will be too low.
  PADDLE_ENFORCE_EQ(regular_type >= 0 && regular_type <= 2,
                    true,
                    common::errors::InvalidArgument(
                        "DGC only support one of None|L1Decay|L2Decay "
                        "Regularization for now."));
  if (regular_type == 0) {
    grad_out_e.device(eigen_ctx) = (1.0 * nranks) * grad_e;
  } else if (regular_type == 1) {
    // L1Decay. grad = grad + coeff * sign(param)
    grad_out_e.device(eigen_ctx) =
        (1.0 * nranks) * grad_e + regular_coeff * param_e.sign();
  } else if (regular_type == 2) {
    // L2Decay. grad = grad + coeff * param
    grad_out_e.device(eigen_ctx) =
        (1.0 * nranks) * grad_e + regular_coeff * param_e;
  }

  // current step
  const float* current_step = current_step_tensor.data<float>();

  if (static_cast<int>(*current_step) < static_cast<int>(rampup_begin_step)) {
    VLOG(10) << "current_step:" << *current_step
             << " < rampup_begin_step:" << rampup_begin_step
             << " so doesn't use dgc";
    return;
  }

  float ratio = 1 - get_period_sparcity(
                        sparsity,
                        static_cast<float>(*current_step - rampup_begin_step),
                        rampup_step);
  PADDLE_ENFORCE_GE(
      ratio,
      0.0,
      common::errors::InvalidArgument("DGC sparsity ratio must >= 0"));
  PADDLE_ENFORCE_LT(
      ratio,
      1.0,
      common::errors::InvalidArgument("DGC sparsity ratio must < 1"));
  int64_t k_64 = static_cast<int64_t>(grad.numel() * ratio);
  PADDLE_ENFORCE_LE_INT_MAX(k_64, "dgc k");
  int k = static_cast<int>(k_64);

  VLOG(10) << "m:" << m << ", use_nesterov:" << use_nesterov
           << ", rampup_begin_step:" << rampup_begin_step
           << ", rampup_step:" << rampup_step
           << ",  current_step:" << *current_step << ", ratio:" << ratio
           << ", k:" << k << ", nranks:" << nranks;

  T* k_out_data = k_out->data<T>();
  *k_out_data = k;

  // FIXME(gongwb): use cublas.
  auto u_out_e = EigenVector<T>::Flatten(*u_out);
  auto u_e = EigenVector<T>::Flatten(u);

  // calc local momentum from global momentum
  // NOTE. If grad not multi nranks, need add below code.
  // if (static_cast<int>(*current_step) ==
  //     static_cast<int>(rampup_begin_step)) {
  //   u_out_e.device(eigen_ctx) = (1.0 / nranks) * u_e;
  // }

  if (use_nesterov) {
    // u = m * (u + grad)
    u_out_e.device(eigen_ctx) = m * (u_e + grad_out_e);

    // v = u + v + grad
    dev_ctx.template Alloc<T>(v_out);
    funcs::ElementwiseCompute<funcs::AddFunctor<T>, T>(
        dev_ctx, u, v, funcs::AddFunctor<T>(), v_out, 0);

    funcs::ElementwiseCompute<funcs::AddFunctor<T>, T>(
        dev_ctx, grad, v, funcs::AddFunctor<T>(), v_out, 0);
  } else {
    // u = m * u + grad
    u_out_e.device(eigen_ctx) = m * u_e + grad_out_e;

    // v = u + v
    dev_ctx.template Alloc<T>(v_out);
    funcs::ElementwiseCompute<funcs::AddFunctor<T>, T>(
        dev_ctx, u, v, funcs::AddFunctor<T>(), v_out, 0);
  }

  T* v_out_data = dev_ctx.template Alloc<T>(v_out);
  T* u_out_data = dev_ctx.template Alloc<T>(u_out);

  encode_grad_out->Resize(DDim{2 * k});
  T* encode_grad_out_data = dev_ctx.template Alloc<T>(encode_grad_out);

  gather_buff->Resize(DDim{2 * k * nranks});
  dev_ctx.template Alloc<T>(gather_buff);

  int buf_size = paddle::communication::dgc::get_buffer_size(k);
  Allocator::AllocationPtr tmp_ious_data;
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  if (dev_ctx.GetPlace().GetType() == AllocationType::GPU ||
      dev_ctx.GetPlace().GetType() == AllocationType::CUSTOM) {
    tmp_ious_data = memory_utils::Alloc(
        dev_ctx.GetPlace(),
        buf_size,
        Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
  }
#endif
  if (dev_ctx.GetPlace().GetType() == AllocationType::CPU) {
    tmp_ious_data = memory_utils::Alloc(dev_ctx.GetPlace(), buf_size);
  }

  void* buf = reinterpret_cast<void*>(tmp_ious_data->ptr());

  int64_t v_out_numel_64 = v_out->numel();
  PADDLE_ENFORCE_LE_INT_MAX(v_out_numel_64, "dgc v_out numel");
  int v_out_numel = static_cast<int>(v_out_numel_64);
  if (!paddle::communication::dgc::k_select(
          static_cast<void*>(encode_grad_out_data),
          k,
          v_out_data,
          v_out_numel,
          buf,
          dev_ctx.stream(),
          u_out_data)) {
    // TODO(weihang): owner should polish this error message
    PADDLE_THROW(common::errors::InvalidArgument(
        "V_out numel error, V_out numel is %d.", v_out->numel()));
  }

  funcs::SetConstant<Context, T> tset;
  tset(dev_ctx, grad_out, static_cast<T>(0));
}

}  // namespace phi

PD_REGISTER_KERNEL(dgc, GPU, ALL_LAYOUT, phi::DGCKernel, float) {}