// Copyright (c) 2023 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.

/*
 * Copyright (c) 2022-2023, NVIDIA CORPORATION.  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/apply_per_channel_scale_kernel.h"

#include <assert.h>
#include <stdint.h>
#include <cmath>
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/common/datatype_traits.h"
#include "paddle/phi/core/kernel_registry.h"

namespace phi {

namespace {
#ifdef PADDLE_WITH_CUDA
template <typename T>
struct CUDA_HALF_2_TYPE_TARIS {};

template <>
struct CUDA_HALF_2_TYPE_TARIS<half> {
  using type = half2;
};

#ifdef PADDLE_CUDA_BF16
template <>
struct CUDA_HALF_2_TYPE_TARIS<__nv_bfloat16> {
  using type = __nv_bfloat162;
};
#endif

template <typename T>
struct HalfMul2 {};

template <>
struct HalfMul2<half2> {
  static __device__ __forceinline__ half2 apply(const half2& x,
                                                const half2& y) {
    return __hmul2(x, y);
  }
};

#ifdef PADDLE_CUDA_BF16
template <>
struct HalfMul2<__nv_bfloat162> {
  static __device__ __forceinline__ __nv_bfloat162
  apply(const __nv_bfloat162& x, const __nv_bfloat162& y) {
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
    return __hmul2(x, y);
#else
    float fxl, fxh, fyl, fyh;
    fxl = __low2float(x);
    fxh = __high2float(x);
    fyl = __low2float(y);
    fyh = __high2float(y);
    return __floats2bfloat162_rn(fxl * fyl, fxh * fyh);
#endif
  }
};
#endif

template <typename T, int kProcessRows, typename AccessType>
__global__ void apply_per_channel_scale(
    const T* act, const T* scales, int rows, int cols, T* out) {
  using HALF_2_TYPE = typename CUDA_HALF_2_TYPE_TARIS<T>::type;
  static constexpr int kElems = sizeof(AccessType) / sizeof(T);
  T scale[kElems], act_vec[kElems];
  int64_t col_offset =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  int row_offset = blockIdx.y;
  if (col_offset * kElems >= cols || row_offset * kProcessRows >= rows) return;
  act += row_offset * kProcessRows * cols;
  out += row_offset * kProcessRows * cols;
  *reinterpret_cast<AccessType*>(scale) =
      reinterpret_cast<const AccessType*>(scales)[col_offset];
#pragma unroll
  for (int i = 0; i < kProcessRows; ++i) {
    *reinterpret_cast<AccessType*>(act_vec) =
        reinterpret_cast<const AccessType*>(act + i * cols)[col_offset];
    if constexpr (kElems % 2 == 0 && (std::is_same<T, half>::value
#ifdef PADDLE_CUDA_BF16
                                      || std::is_same<T, __nv_bfloat16>::value
#endif
                                      )) {
#pragma unroll
      for (int j = 0; j < kElems; j += 2) {
        *reinterpret_cast<HALF_2_TYPE*>(act_vec + j) =
            HalfMul2<HALF_2_TYPE>::apply(
                *reinterpret_cast<HALF_2_TYPE*>(act_vec + j),
                *reinterpret_cast<HALF_2_TYPE*>(scale + j));
      }
    } else {
#pragma unroll
      for (int j = 0; j < kElems; ++j) {
        act_vec[j] *= scale[j];
      }
    }
    reinterpret_cast<AccessType*>(out + i * cols)[col_offset] =
        *reinterpret_cast<AccessType*>(act_vec);
  }
}

template <typename T, int kProcessRows, typename AccessType = float4>
void apply_per_channel_scale_launcher(const T* act,
                                      const T* scales,
                                      int rows,
                                      int cols,
                                      T* out,
                                      cudaStream_t stream = 0) {
  static constexpr int kElems = sizeof(AccessType) / sizeof(T);
  dim3 block(128);
  dim3 grid((cols / kElems + block.x - 1) / block.x,
            (rows + kProcessRows - 1) / kProcessRows);
  apply_per_channel_scale<T, kProcessRows, AccessType>
      <<<grid, block, 0, stream>>>(act, scales, rows, cols, out);
}

}  // namespace
#endif

template <typename T, typename Context>
void ApplyPerChannelScaleKernel(const Context& dev_ctx,
                                const DenseTensor& x,
                                const DenseTensor& scales,
                                DenseTensor* out) {
#ifdef PADDLE_WITH_CUDA
  using DataType = typename PDDataTypeTraits<T>::DataType;
  int64_t rows = x.dims()[0];
  int64_t cols = x.dims()[1];
  int64_t elems = rows * cols;
  const T* x_data = x.data<T>();
  const T* scales_data = scales.data<T>();
  T* out_data = dev_ctx.template Alloc<T>(out);
  if (elems < 2048 * 2048) {
    apply_per_channel_scale_launcher<DataType, 1, float4>(
        reinterpret_cast<const DataType*>(x_data),
        reinterpret_cast<const DataType*>(scales_data),
        rows,
        cols,
        reinterpret_cast<DataType*>(out_data),
        dev_ctx.stream());
  } else if (elems < 4096 * 4096) {
    apply_per_channel_scale_launcher<DataType, 4, float4>(
        reinterpret_cast<const DataType*>(x_data),
        reinterpret_cast<const DataType*>(scales_data),
        rows,
        cols,
        reinterpret_cast<DataType*>(out_data),
        dev_ctx.stream());
  } else if (elems < 8192 * 8192) {
    apply_per_channel_scale_launcher<DataType, 8, float4>(
        reinterpret_cast<const DataType*>(x_data),
        reinterpret_cast<const DataType*>(scales_data),
        rows,
        cols,
        reinterpret_cast<DataType*>(out_data),
        dev_ctx.stream());
  } else {
    PADDLE_ENFORCE_LE_INT_MAX(rows, "rows");
    PADDLE_ENFORCE_LE_INT_MAX(cols, "cols");
    apply_per_channel_scale_launcher<DataType, 16, float4>(
        reinterpret_cast<const DataType*>(x_data),
        reinterpret_cast<const DataType*>(scales_data),
        rows,
        cols,
        reinterpret_cast<DataType*>(out_data),
        dev_ctx.stream());
  }
#endif
}

}  // namespace phi

PD_REGISTER_KERNEL(apply_per_channel_scale,
                   GPU,
                   ALL_LAYOUT,
                   phi::ApplyPerChannelScaleKernel,
                   phi::float16,
                   phi::bfloat16) {}