// 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/swiglu_grad_kernel.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/activation_functor.h"

#include "paddle/phi/kernels/funcs/aligned_vector.h"
#include "paddle/phi/kernels/primitive/kernel_primitives.h"

namespace phi {

template <typename T, int VecSize, bool IsCombine, bool HasDX, bool HasDY>
__global__ void SwiGLUGradCUDAKernel(const T *__restrict__ x,
                                     const T *__restrict__ y,
                                     const T *__restrict__ dz,
                                     T *__restrict__ dx,
                                     T *__restrict__ dy,
                                     int64_t m,
                                     int64_t n) {
  funcs::SwiGLUGradFunctor<T, HasDX, HasDY> functor;
  if constexpr (IsCombine) {
    int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
    int64_t n_vec_piece = n / VecSize;
    int64_t valid_num = m * n_vec_piece;
    while (idx < valid_num) {
      int64_t row_offset = idx / n_vec_piece * n;
      int64_t col_offset = idx % n_vec_piece * VecSize;

      int64_t dz_offset = row_offset + col_offset;
      int64_t x_offset = dz_offset + row_offset;
      AlignedVector<T, VecSize> x_vec;
      AlignedVector<T, VecSize> y_vec;
      AlignedVector<T, VecSize> dz_vec;

      Load<T, VecSize>(x + x_offset, &x_vec);
      Load<T, VecSize>(y + x_offset, &y_vec);
      Load<T, VecSize>(dz + dz_offset, &dz_vec);

#pragma unroll
      for (int i = 0; i < VecSize; ++i) {
        functor(x_vec[i], y_vec[i], dz_vec[i], &x_vec[i], &y_vec[i]);
      }

      Store<T, VecSize>(x_vec, dx + x_offset);
      Store<T, VecSize>(y_vec, dy + x_offset);
      idx += stride;
    }
  } else {
    int64_t idx =
        (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) * VecSize;
    int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x * VecSize;
    int64_t numel = m * n;
    int64_t limit = numel - VecSize;

    while (idx <= limit) {
      AlignedVector<T, VecSize> x_vec;
      AlignedVector<T, VecSize> y_vec;
      AlignedVector<T, VecSize> dz_vec;

      Load<T, VecSize>(x + idx, &x_vec);
      if constexpr (HasDX) {
        Load<T, VecSize>(y + idx, &y_vec);
      }
      Load<T, VecSize>(dz + idx, &dz_vec);

#pragma unroll
      for (int i = 0; i < VecSize; ++i) {
        functor(x_vec[i],
                HasDX ? y_vec[i] : x_vec[i],
                dz_vec[i],
                &x_vec[i],
                &dz_vec[i]);
      }
      if constexpr (HasDX) {
        Store<T, VecSize>(x_vec, dx + idx);
      }
      if constexpr (HasDY) {
        Store<T, VecSize>(dz_vec, dy + idx);
      }
      idx += stride;
    }

    while (idx < numel) {
      functor(x[idx],
              HasDX ? y[idx] : x[idx],
              dz[idx],
              HasDX ? &dx[idx] : nullptr,
              HasDY ? &dy[idx] : nullptr);
      ++idx;
    }
  }
}

template <typename T, typename Context>
void SwiGLUGradKernelImpl(const Context &dev_ctx,
                          const T *x,
                          const T *y,
                          const T *dz,
                          T *dx,
                          T *dy,
                          int64_t m,
                          int64_t n) {
  int vec_size = std::min(GetVectorizedSize<T>(x), GetVectorizedSize<T>(dz));
  if (y) {
    vec_size = std::min(vec_size, GetVectorizedSize<T>(y));
  }
  if (dx) {
    vec_size = std::min(vec_size, GetVectorizedSize<T>(dx));
  }
  if (dy) {
    vec_size = std::min(vec_size, GetVectorizedSize<T>(dy));
  }

#define PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL_BASE(                           \
    __vec_size, __is_combine, __has_dx, __has_dy)                         \
  case __vec_size: {                                                      \
    SwiGLUGradCUDAKernel<T, __vec_size, __is_combine, __has_dx, __has_dy> \
        <<<config.block_per_grid,                                         \
           config.thread_per_block,                                       \
           0,                                                             \
           dev_ctx.stream()>>>(x, y, dz, dx, dy, m, n);                   \
    break;                                                                \
  }

#define PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL(__is_combine, __has_dx, __has_dy) \
  do {                                                                      \
    switch (vec_size) {                                                     \
      PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL_BASE(                               \
          VecSizeVL, __is_combine, __has_dx, __has_dy);                     \
      PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL_BASE(                               \
          VecSizeL, __is_combine, __has_dx, __has_dy);                      \
      PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL_BASE(                               \
          VecSizeM, __is_combine, __has_dx, __has_dy);                      \
      PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL_BASE(                               \
          VecSizeS, __is_combine, __has_dx, __has_dy);                      \
      default:                                                              \
        PADDLE_THROW(common::errors::Unimplemented(                         \
            "Unsupported vectorized size: %d !", vec_size));                \
        break;                                                              \
    }                                                                       \
  } while (0)

  if (y) {
    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, m * n, vec_size);
    if (dx) {
      if (dy) {
        PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL(false, true, true);
      } else {
        PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL(false, true, false);
      }
    } else {
      PADDLE_ENFORCE_NOT_NULL(
          dy,
          common::errors::InvalidArgument(
              "Both gradients of Input(X) and Input(Y) is None."));
      PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL(false, false, true);
    }
  } else {
    PADDLE_ENFORCE_NOT_NULL(
        dx,
        common::errors::InvalidArgument(
            "Both gradients of Input(X) and Input(Y) is None."));
    while (n % vec_size != 0) {
      vec_size /= 2;
    }
    y = x + n;
    dy = dx + n;
    auto config =
        backends::gpu::GetGpuLaunchConfig1D(dev_ctx, m * n / vec_size, 1);
    PD_LAUNCH_SWIGLU_GRAD_CUDA_KERNEL(true, true, true);
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(swiglu_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::SwiGLUGradKernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {}