// 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/affine_channel_grad_kernel.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_primitives.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/cub.h"

namespace phi {

template <typename T, DataLayout layout, bool HasBias>
__global__ static inline void KeAffineChannelCUDA(const T* x,
                                                  const T* scale,
                                                  const T* bias,
                                                  const int C,
                                                  const int64_t HxW,
                                                  const int64_t num,
                                                  T* y) {
  int64_t gid =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  int stride = blockDim.x * gridDim.x;
  for (int64_t i = gid; i < num; i += stride) {
    const int c = layout == DataLayout::NCHW ? i / HxW % C : i % C;
    if (HasBias) {
      y[i] = scale[c] * x[i] + bias[c];
    } else {
      y[i] = scale[c] * x[i];
    }
  }
}

template <typename T, int BlockDim, DataLayout layout>
__global__ void AffineChannelScaleBiasGradientCUDAKernel(const T* dy,
                                                         const T* x,
                                                         const int N,
                                                         const int C,
                                                         const int64_t HxW,
                                                         T* dscale,
                                                         T* dbias) {
  const int outer_size = C;
  const int64_t inner_size = HxW * N;
  typedef cub::BlockReduce<double, BlockDim> BlockReduce;
  __shared__ typename BlockReduce::TempStorage ds_storage;
  __shared__ typename BlockReduce::TempStorage db_storage;

  for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
    T ds_sum = 0;
    T db_sum = 0;
    for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
      const int64_t index = layout == DataLayout::NCHW
                                ? (j / HxW * C + i) * HxW + j % HxW
                                : j * outer_size + i;
      ds_sum += dy[index] * x[index];
      db_sum += dy[index];
    }
    __syncthreads();
    auto ds_out =
        BlockReduce(ds_storage).Reduce(static_cast<double>(ds_sum), cub::Sum());
    auto db_out =
        BlockReduce(db_storage).Reduce(static_cast<double>(db_sum), cub::Sum());
    __syncthreads();
    if (threadIdx.x == 0) {
      dscale[i] = ds_out;
      dbias[i] = db_out;
    }
  }
}

template <typename T, typename Context>
void AffineChannelGradCUDAKernel(const Context& dev_ctx,
                                 const DenseTensor& x_in,
                                 const DenseTensor& scale_in,
                                 const DenseTensor& bias_in,
                                 const DenseTensor& out_grad,
                                 const std::string& data_layout,
                                 DenseTensor* x_grad,
                                 DenseTensor* scale_grad,
                                 DenseTensor* bias_grad) {
  auto* x = &x_in;
  auto* scale = &scale_in;
  auto* bias = &bias_in;
  auto* dy = &out_grad;

  auto* dx = x_grad;
  auto* dscale = scale_grad;
  auto* dbias = bias_grad;

  const DataLayout layout = StringToDataLayout(data_layout);

  auto dims = dy->dims();
  const int64_t num = dy->numel();
  int64_t N = dims[0];
  int64_t C = layout == DataLayout::NCHW ? dims[1] : dims[dims.size() - 1];
  int64_t HxW = num / N / C;

  const T* dy_d = dy->data<T>();
  const T* s_d = scale->data<T>();

  T* dx_d = dx ? dev_ctx.template Alloc<T>(dx) : nullptr;
  T* ds_d = dscale ? dev_ctx.template Alloc<T>(dscale) : nullptr;
  T* db_d = dbias ? dev_ctx.template Alloc<T>(dbias) : nullptr;

#ifdef PADDLE_WITH_HIP
  const int block = 256;
#else
  const int block = 1024;
#endif  // PADDLE_WITH_HIP
  int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
  const int max_blocks = std::max(max_threads / block, 1);
  int grid1 = (num + block - 1) / block;
  int grid2 = std::min(static_cast<int>(C), max_blocks);

  // NOTE(large-tensor): Kernel functions expect int for N and C parameters
  PADDLE_ENFORCE_LE_INT_MAX(N, "N");
  PADDLE_ENFORCE_LE_INT_MAX(C, "C");
  if (layout == DataLayout::NCHW) {
    if (dscale && dbias) {
      const T* x_d = x->data<T>();
      AffineChannelScaleBiasGradientCUDAKernel<T, block, DataLayout::NCHW>
          <<<grid2, block, 0, dev_ctx.stream()>>>(dy_d,
                                                  x_d,
                                                  static_cast<int>(N),
                                                  static_cast<int>(C),
                                                  HxW,
                                                  ds_d,
                                                  db_d);
    }
    if (dx) {
      KeAffineChannelCUDA<T, DataLayout::NCHW, false>
          <<<grid1, block, 0, dev_ctx.stream()>>>(
              dy_d, s_d, nullptr, static_cast<int>(C), HxW, num, dx_d);
    }
  } else {
    if (dscale && dbias) {
      const T* x_d = x->data<T>();
      AffineChannelScaleBiasGradientCUDAKernel<T, block, DataLayout::NHWC>
          <<<grid2, block, 0, dev_ctx.stream()>>>(dy_d,
                                                  x_d,
                                                  static_cast<int>(N),
                                                  static_cast<int>(C),
                                                  HxW,
                                                  ds_d,
                                                  db_d);
    }

    if (dx) {
      KeAffineChannelCUDA<T, DataLayout::NHWC, false>
          <<<grid1, block, 0, dev_ctx.stream()>>>(
              dy_d, s_d, nullptr, static_cast<int>(C), HxW, num, dx_d);
    }
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(affine_channel_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::AffineChannelGradCUDAKernel,
                   float,
                   double) {}