// 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/layer_norm_kernel.h"
#include "paddle/common/enforce.h"
#include "paddle/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
#include "paddle/phi/kernels/funcs/fast_ln_v2.h"
#endif
#include "paddle/phi/kernels/funcs/layer_norm_impl.cu.h"
#include "paddle/phi/kernels/funcs/layer_norm_util.h"
#ifdef PADDLE_WITH_CUDA
#include "paddle/phi/kernels/gpu/rms_norm_cuda_kernel.h"
#endif
COMMON_DECLARE_bool(use_fast_math);
COMMON_DECLARE_bool(use_apex_layer_norm_kernel);

namespace phi {

enum class LayerNormKernelVariant { FAST_LN_V1, FAST_LN_V2, GENERIC };

#ifdef PADDLE_WITH_CUDA
template <typename U>
__device__ inline void WelfordOnline(U val, U* mean, U* square, U* count) {
  *count += 1;
  U delta1 = val - *mean;
  *mean += delta1 / (*count);
  U delta2 = val - *mean;
  *square += delta1 * delta2;
}

template <typename U>
__device__ inline void WelfordOnline(
    U b_mean, U b_square, U b_cnt, U* mean, U* square, U* count) {
  if (b_cnt == 0) {
    return;
  }

  U new_cnt = *count + b_cnt;
  U nb_n = b_cnt / new_cnt;
  U delta = b_mean - *mean;
  *mean += delta * nb_n;
  *square += b_square + delta * delta * (*count) * nb_n;
  *count = new_cnt;
}

template <typename U>
__device__ inline void WelfordWarpAllReduce(U* mean, U* square, U* count) {
  constexpr int kWarpSize = 32;
#pragma unroll
  for (int mask = 1; mask < kWarpSize; mask *= 2) {
    U b_mean = __shfl_down_sync(0xffffffff, *mean, mask);
    U b_square = __shfl_down_sync(0xffffffff, *square, mask);
    U b_cnt = __shfl_down_sync(0xffffffff, *count, mask);
    WelfordOnline<U>(b_mean, b_square, b_cnt, mean, square, count);
  }

  *mean = __shfl_sync(0xffffffff, *mean, 0, kWarpSize);
  *square = __shfl_sync(0xffffffff, *square, 0, kWarpSize);
  *count = __shfl_sync(0xffffffff, *count, 0, kWarpSize);
}

template <int VecSize>
struct ThreadAssigner {
  __device__ __forceinline__ int operator()(const int cols,
                                            const int cols_per_thread,
                                            int32_t* last_tid_idx) {
    return cols_per_thread;
  }
};

template <>
struct ThreadAssigner<1> {
  __device__ inline int operator()(const int cols,
                                   const int cols_per_thread,
                                   int* last_tid_idx) {
    int cols_this_thread = cols_per_thread;
    int last_tid = (cols / cols_per_thread);
    *last_tid_idx = last_tid;
    if (threadIdx.x == last_tid) {
      cols_this_thread = cols - cols_per_thread * last_tid;
    } else if (threadIdx.x > last_tid) {
      cols_this_thread = 0;
    }
    return cols_this_thread;
  }
};

template <typename T, typename U, int VecSize>
struct LayerNormDataReader {
  __device__ inline void operator()(const T* __restrict__ row_src,
                                    U* buffer,
                                    const int last_tid_idx,
                                    const int read_times,
                                    const int cols_this_thread) {
    using VecT = AlignedVector<T, VecSize>;
    const VecT* __restrict__ v_src =
        reinterpret_cast<const VecT* __restrict__>(row_src);

    for (int i = 0; i < read_times; ++i) {
      VecT temp_src = v_src[threadIdx.x + i * blockDim.x];
#pragma unroll
      for (int j = 0; j < VecSize; ++j) {
        buffer[i * VecSize + j] = static_cast<U>(temp_src[j]);
      }
    }
  }
};

template <typename T, typename U>
struct LayerNormDataReader<T, U, 1> {
  __device__ inline void operator()(const T* __restrict__ row_src,
                                    U* buffer,
                                    const int last_tid_idx,
                                    const int read_times,
                                    const int cols_this_thread) {
    // read_time is just cols_per_thread while VecSize is 1.
    if (threadIdx.x < last_tid_idx) {
      for (int i = 0; i < cols_this_thread; ++i) {
        buffer[i] = static_cast<U>(row_src[threadIdx.x + last_tid_idx * i]);
      }
    } else {
      for (int i = 0; i < cols_this_thread; ++i) {
        buffer[i] = static_cast<U>(row_src[i + read_times * last_tid_idx]);
      }
    }
  }
};

template <typename T, typename U, bool IsSameType, int VecSize>
struct LayerNormDataWriter {
  __device__ inline void operator()(
      T* __restrict__ row_dst,
      const U* __restrict__ buffer,
      const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ scale,
      const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ bias,
      const U row_mean,
      const U row_inv_var,
      const int write_times,
      const int cols_this_thread,
      const int last_tid_idx,
      const bool valid_scale,
      const bool valid_bias) {
    using VecT = AlignedVector<T, VecSize>;
    using ScaleT = funcs::LayerNormScaleBiasT<T, U, IsSameType>;
    using VecScaleT = AlignedVector<ScaleT, VecSize>;
    VecT* v_dst = reinterpret_cast<VecT*>(row_dst);

    // cols_this_thread is just cols_per_thread
    if ((!valid_scale) && (!valid_bias)) {
      for (int i = 0; i < write_times; ++i) {
        VecT temp_dst;
#pragma unroll
        for (int j = 0; j < VecSize; ++j) {
          temp_dst[j] = static_cast<T>((buffer[i * VecSize + j] - row_mean) *
                                       row_inv_var);
        }
        v_dst[threadIdx.x + static_cast<int64_t>(blockDim.x) * i] = temp_dst;
      }
    } else {
      const VecScaleT* __restrict__ v_scale =
          reinterpret_cast<const VecScaleT* __restrict__>(scale);
      const VecScaleT* __restrict__ v_bias =
          reinterpret_cast<const VecScaleT* __restrict__>(bias);
      if (valid_scale && valid_bias) {
        for (int i = 0; i < write_times; ++i) {
          int64_t idx = threadIdx.x + static_cast<int64_t>(blockDim.x) * i;
          VecT temp_dst;
          VecScaleT temp_v_scale = v_scale[idx];
          VecScaleT temp_v_bias = v_bias[idx];
#pragma unroll
          for (int j = 0; j < VecSize; ++j) {
            temp_dst[j] = static_cast<T>(
                static_cast<U>(temp_v_scale[j]) *
                    (buffer[i * VecSize + j] - row_mean) * row_inv_var +
                static_cast<U>(temp_v_bias[j]));
          }
          v_dst[idx] = temp_dst;
        }
      } else {
        if (valid_scale) {
          for (int i = 0; i < write_times; ++i) {
            int64_t idx = threadIdx.x + static_cast<int64_t>(blockDim.x) * i;
            VecT temp_dst;
            VecScaleT temp_v_scale = v_scale[idx];
#pragma unroll
            for (int j = 0; j < VecSize; ++j) {
              temp_dst[j] = static_cast<T>(
                  static_cast<U>(temp_v_scale[j]) *
                  (buffer[i * VecSize + j] - row_mean) * row_inv_var);
            }
            v_dst[idx] = temp_dst;
          }
        } else {
          for (int i = 0; i < write_times; ++i) {
            int idx = threadIdx.x + blockDim.x * i;
            VecT temp_dst;
            VecScaleT temp_v_bias = v_bias[idx];
#pragma unroll
            for (int j = 0; j < VecSize; ++j) {
              temp_dst[j] = static_cast<T>(
                  (buffer[i * VecSize + j] - row_mean) * row_inv_var +
                  static_cast<U>(temp_v_bias[j]));
            }
            v_dst[idx] = temp_dst;
          }
        }
      }
    }
  }
};

template <typename T, typename U, bool IsSameType>
struct LayerNormDataWriter<T, U, IsSameType, 1> {
  __device__ __forceinline__ void operator()(
      T* __restrict__ row_dst,
      U* __restrict__ buffer,
      const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ scale,
      const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ bias,
      const U row_mean,
      const U row_inv_var,
      const int write_times,
      const int cols_this_thread,
      const int last_tid_idx,
      const bool valid_scale,
      const bool valid_bias) {
    // write_times is just col_per_thread.
    if ((!valid_scale) && (!valid_bias)) {
      if (threadIdx.x < last_tid_idx) {
        for (int i = 0; i < cols_this_thread; ++i) {
          row_dst[threadIdx.x + static_cast<int64_t>(last_tid_idx) * i] =
              (buffer[i] - row_mean) * row_inv_var;
        }
      } else {
        for (int i = 0; i < cols_this_thread; ++i) {
          row_dst[static_cast<int64_t>(last_tid_idx) * write_times + i] =
              (buffer[i] - row_mean) * row_inv_var;
        }
      }
    } else if (valid_scale && valid_bias) {
      if (threadIdx.x < last_tid_idx) {
        for (int i = 0; i < cols_this_thread; ++i) {
          int64_t idx = threadIdx.x + static_cast<int64_t>(last_tid_idx) * i;
          row_dst[idx] =
              static_cast<T>(static_cast<U>(scale[idx]) *
                                 (buffer[i] - row_mean) * row_inv_var +
                             static_cast<U>(bias[idx]));
        }
      } else {
        for (int i = 0; i < cols_this_thread; ++i) {
          int64_t idx = static_cast<int64_t>(last_tid_idx) * write_times + i;
          row_dst[idx] =
              static_cast<T>(static_cast<U>(scale[idx]) *
                                 (buffer[i] - row_mean) * row_inv_var +
                             static_cast<U>(bias[idx]));
        }
      }
    } else {
      if (valid_scale) {
        if (threadIdx.x < last_tid_idx) {
          for (int i = 0; i < cols_this_thread; ++i) {
            int64_t idx = threadIdx.x + static_cast<int64_t>(last_tid_idx) * i;
            row_dst[idx] = static_cast<T>(static_cast<U>(scale[idx]) *
                                          (buffer[i] - row_mean) * row_inv_var);
          }
        } else {
          for (int i = 0; i < cols_this_thread; ++i) {
            int64_t idx = static_cast<int64_t>(last_tid_idx) * write_times + i;
            row_dst[idx] = static_cast<T>(static_cast<U>(scale[idx]) *
                                          (buffer[i] - row_mean) * row_inv_var);
          }
        }
      } else {
        if (threadIdx.x < last_tid_idx) {
          for (int i = 0; i < cols_this_thread; ++i) {
            int64_t idx = threadIdx.x + static_cast<int64_t>(last_tid_idx) * i;
            row_dst[idx] = static_cast<T>((buffer[i] - row_mean) * row_inv_var +
                                          static_cast<U>(bias[idx]));
          }
        } else {
          for (int i = 0; i < cols_this_thread; ++i) {
            int64_t idx = static_cast<int64_t>(last_tid_idx) * write_times + i;
            row_dst[idx] = static_cast<T>((buffer[i] - row_mean) * row_inv_var +
                                          static_cast<U>(bias[idx]));
          }
        }
      }
    }
  }
};

template <typename IndexT, typename T, typename U, bool IsSameType, int VecSize>
__global__ void LayerNormFwdWithWelford(
    const T* __restrict__ src_data,
    T* dst_data,
    const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ scale,
    const funcs::LayerNormScaleBiasT<T, U, IsSameType>* __restrict__ bias,
    U* mean,
    U* var,
    const U epsilon,
    const IndexT rows,
    const int32_t cols,
    const int32_t cols_per_thread,
    const bool valid_scale,
    const bool valid_bias) {
  constexpr int kWarpSize = 32;
  int last_tid_idx = 0;  // For condition once vecSize is 1.
  IndexT row_offset = blockIdx.x * blockDim.y + threadIdx.y;
  int cols_this_thread =
      ThreadAssigner<VecSize>()(cols, cols_per_thread, &last_tid_idx);
  int read_times = cols_per_thread / VecSize;

  if (row_offset < rows) {
    U buffer[kWarpSize];
    U tid_cnt = static_cast<U>(0);
    U tid_mean = static_cast<U>(0);
    U tid_square = static_cast<U>(0);

    const T* __restrict__ row_src = src_data + row_offset * cols;
    T* row_dst = dst_data + row_offset * cols;
    LayerNormDataReader<T, U, VecSize>()(
        row_src, buffer, last_tid_idx, read_times, cols_this_thread);

    for (int i = 0; i < cols_this_thread; i++) {
      WelfordOnline<U>(buffer[i], &tid_mean, &tid_square, &tid_cnt);
    }

    U warp_cnt = tid_cnt;
    U warp_mean = tid_mean;
    U warp_square = tid_square;
    WelfordWarpAllReduce<U>(&warp_mean, &warp_square, &warp_cnt);

    U row_variance = max(warp_square / warp_cnt, 0.f);
    U row_inv_var = funcs::rsqrt_(row_variance + epsilon);

    // TODO(limingshu): make code below vectorization.
    if (threadIdx.x == 0) {
      // warp_mean is just row_mean here.
      mean[row_offset] = warp_mean;
      var[row_offset] = row_variance;
    }
    LayerNormDataWriter<T, U, IsSameType, VecSize>()(row_dst,
                                                     buffer,
                                                     scale,
                                                     bias,
                                                     warp_mean,
                                                     row_inv_var,
                                                     read_times,
                                                     cols_this_thread,
                                                     last_tid_idx,
                                                     valid_scale,
                                                     valid_bias);
  }
}

template <typename Context, typename T, typename U>
void LaunchLayerNormKernel(const Context& dev_ctx,
                           const T* x_data,
                           T* y_data,
                           const void* void_scale_data,
                           const void* void_bias_data,
                           U* mean_data,
                           U* var_data,
                           double epsilon,
                           const int64_t rows,
                           const int cols,
                           const bool valid_scale,
                           const bool valid_bias,
                           const bool is_same_type) {
  constexpr int WarpSize = 32;
  constexpr int RowPerBlock = 4;
  int64_t block_size = (rows + (RowPerBlock - 1)) / RowPerBlock;
  dim3 threads(WarpSize, RowPerBlock, 1);

  int vec_size = 1;
  int cols_per_thread = (cols + (WarpSize - 1)) / WarpSize;
  if (cols_per_thread > 1 && (cols % WarpSize == 0)) {
    int data_vec_size = 0;
    uint64_t addr = (reinterpret_cast<uint64_t>(x_data) |
                     reinterpret_cast<uint64_t>(y_data));
    if (valid_bias || valid_scale) {
      if (is_same_type) {
        addr = valid_scale
                   ? (addr | reinterpret_cast<uint64_t>(void_scale_data))
                   : addr;
        addr = valid_bias ? (addr | reinterpret_cast<uint64_t>(void_bias_data))
                          : addr;
        data_vec_size =
            std::min(4, GetVectorizedSize<T>(reinterpret_cast<T*>(addr)));
      } else {
        uint64_t bias_addr = reinterpret_cast<uint64_t>(void_bias_data);
        uint64_t attr_addr = valid_scale
                                 ? reinterpret_cast<uint64_t>(void_scale_data)
                                 : bias_addr;
        attr_addr = valid_bias
                        ? (valid_scale ? (attr_addr | bias_addr) : attr_addr)
                        : attr_addr;
        data_vec_size =
            std::min(GetVectorizedSize<T>(reinterpret_cast<T*>(addr)),
                     GetVectorizedSize<U>(reinterpret_cast<U*>(attr_addr)));
        data_vec_size = std::min(4, data_vec_size);
      }
    }
    for (int size = data_vec_size; size > 0; size /= 2) {
      if (cols_per_thread % size == 0) {
        vec_size = size;
        break;
      }
    }
  }

#define IMPL_LAYER_NORM_WELFORD_CASE(index_t, scale_t, is_same_, vec_size_)    \
  case (vec_size_): {                                                          \
    PADDLE_ENFORCE_LE_UINT32_MAX(block_size, "grid.x");                        \
    LayerNormFwdWithWelford<index_t, T, U, is_same_, vec_size_>                \
        <<<static_cast<uint32_t>(block_size), threads, 0, dev_ctx.stream()>>>( \
            x_data,                                                            \
            y_data,                                                            \
            static_cast<const scale_t*>(void_scale_data),                      \
            static_cast<const scale_t*>(void_bias_data),                       \
            mean_data,                                                         \
            var_data,                                                          \
            static_cast<const U>(epsilon),                                     \
            rows,                                                              \
            cols,                                                              \
            cols_per_thread,                                                   \
            valid_scale,                                                       \
            valid_bias);                                                       \
  } break

#define IMPL_LAYER_NORM_WELFORD(index_t, scale_t, is_same_)    \
  IMPL_LAYER_NORM_WELFORD_CASE(index_t, scale_t, is_same_, 4); \
  IMPL_LAYER_NORM_WELFORD_CASE(index_t, scale_t, is_same_, 2); \
  IMPL_LAYER_NORM_WELFORD_CASE(index_t, scale_t, is_same_, 1);

  if (rows < std::numeric_limits<int32_t>::max()) {
    if (is_same_type) {
      switch (vec_size) { IMPL_LAYER_NORM_WELFORD(int32_t, T, true); }
    } else {
      switch (vec_size) { IMPL_LAYER_NORM_WELFORD(int32_t, U, false); }
    }
  } else {
    if (is_same_type) {
      switch (vec_size) { IMPL_LAYER_NORM_WELFORD(int64_t, T, true); }
    } else {
      switch (vec_size) { IMPL_LAYER_NORM_WELFORD(int64_t, U, false); }
    }
  }
#undef IMPL_LAYER_NORM_WELFORD_CASE
#undef IMPL_LAYER_NORM_WELFORD
}
#endif  // PADDLE_WITH_CUDA

template <typename T, typename U>
void LayerNormDirectCUDAFunctor<T, U>::operator()(
    gpuStream_t stream,
    const T* input,
    std::vector<int64_t> input_shape,
    const U* bias,
    const U* scale,
    T* output,
    U* mean,
    U* variance,
    int begin_norm_axis,
    float eps) {
  const auto x_dims = make_ddim(input_shape);
  auto matrix_dim = common::flatten_to_2d(x_dims, begin_norm_axis);
  int64_t batch_size = matrix_dim[0];
  int64_t feature_size = matrix_dim[1];
  // TODO(large-tensor): generic kernel launch uses uint32 grid dim
  PADDLE_ENFORCE_LE_UINT32_MAX(batch_size,
                               "Kernel launch requires uint32 for grid dim");
  switch (funcs::GetDesiredBlockDim(feature_size)) {
    FIXED_BLOCK_DIM_CASE(
        funcs::LayerNormForward<T, U, kBlockDim>
        <<<static_cast<uint32_t>(batch_size), kBlockDim, 0, stream>>>(
            input, scale, bias, output, mean, variance, eps, feature_size));
    default:
      PADDLE_THROW(common::errors::InvalidArgument(
          "Product from begin_norm_axis to end in layer_norm must be larger "
          "than 1"));
      break;
  }
}

template class PADDLE_API LayerNormDirectCUDAFunctor<float, float>;
template class PADDLE_API LayerNormDirectCUDAFunctor<double, double>;
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP)
template class PADDLE_API LayerNormDirectCUDAFunctor<half, float>;
#endif
static inline LayerNormKernelVariant LayerNormKernelDispatch(
    const DataType weight_type,
    const DataType input_type,
    const DataType output_type,
    const DataType compute_type,
    const int64_t hidden_size,
    const int64_t x_numel,
    const DenseTensor* scale,
    const DenseTensor* bias) {
  if (scale == nullptr || bias == nullptr) {
    return LayerNormKernelVariant::GENERIC;
  }
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
  if (FLAGS_use_apex_layer_norm_kernel) {
    if (funcs::fast_ln_v2::has_fast_ln_v2_fwd_kernel(
            weight_type, input_type, output_type, compute_type, hidden_size)) {
      return LayerNormKernelVariant::FAST_LN_V2;
    }
    PADDLE_THROW(common::errors::InvalidArgument(
        "FLAGS_use_apex_layer_norm_kernel requires inputs supported by "
        "fast_ln_v2 forward kernel."));
  }
#endif
#ifdef PADDLE_WITH_CUDA
  if (FLAGS_use_accuracy_compatible_kernel) {
    return LayerNormKernelVariant::GENERIC;
  }
#endif
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
  if (input_type != DataType::FLOAT32 && hidden_size != 4096 &&
      hidden_size > 1024 && hidden_size <= 10240 &&
      x_numel <= std::numeric_limits<uint32_t>::max()) {
    // using fast_ln_v2 only sm > 70 and x_numel <= uint32_max
    auto prop = funcs::fast_ln_v2::GetDeviceProp();
    uint32_t hidden_size_32 = static_cast<uint32_t>(hidden_size);
    if (prop->major > 7 &&
        funcs::fast_ln_v2::has_fast_ln_v2_fwd_kernel(weight_type,
                                                     input_type,
                                                     output_type,
                                                     compute_type,
                                                     hidden_size_32)) {
      return LayerNormKernelVariant::FAST_LN_V2;
    }
  }
#endif
  if ((hidden_size >= 768 && hidden_size <= 2048 && hidden_size % 256 == 0 ||
       hidden_size == 4096) &&
      x_numel <= std::numeric_limits<int>::max() && scale != nullptr &&
      bias != nullptr) {
    return LayerNormKernelVariant::FAST_LN_V1;
  }

  return LayerNormKernelVariant::GENERIC;
}

template <typename T, typename Context>
void LayerNormKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const optional<DenseTensor>& scale_opt,
                     const optional<DenseTensor>& bias_opt,
                     double epsilon,
                     int begin_norm_axis,
                     DenseTensor* y,
                     DenseTensor* mean,
                     DenseTensor* var) {
  using U = funcs::LayerNormParamType<T>;
  auto* scale = scale_opt.get_ptr();
  auto* bias = bias_opt.get_ptr();

  const auto x_dims = x.dims();
  auto* x_data = x.data<T>();
  auto* y_data = dev_ctx.template Alloc<T>(y);
  auto* mean_data = dev_ctx.template Alloc<U>(mean);
  auto* var_data = dev_ctx.template Alloc<U>(var);
  if (x.numel() == 0) return;

  bool valid_scale = (scale != nullptr);
  bool valid_bias = (bias != nullptr);
  auto* void_scale_data = valid_scale ? scale->data() : nullptr;
  auto* void_bias_data = valid_bias ? bias->data() : nullptr;

  auto x_dtype = x.dtype();
  auto y_dtype = y->dtype();
  DataType scale_bias_dtype;
  if (valid_scale) {
    scale_bias_dtype = scale->dtype();
    if (valid_bias) {
      PADDLE_ENFORCE_EQ(scale->dtype(),
                        bias->dtype(),
                        common::errors::InvalidArgument(
                            "This Scale and Bias of layer_norm op "
                            "should have the same data type."));
    }
  } else {
    scale_bias_dtype = valid_bias ? bias->dtype() : x_dtype;
  }

  bool is_scale_bias_same_dtype_with_x = x_dtype == scale_bias_dtype;
  if (!is_scale_bias_same_dtype_with_x) {
    PADDLE_ENFORCE_EQ(scale_bias_dtype,
                      CppTypeToDataType<U>::Type(),
                      common::errors::InvalidArgument(
                          "Unsupported data type of Scale and Bias"));
  }

  auto matrix_dim = common::flatten_to_2d(x_dims, begin_norm_axis);
  int64_t batch_size = matrix_dim[0];
  // TODO(large-tensor): generic kernel launch uses uint32 grid dim
  PADDLE_ENFORCE_LE_UINT32_MAX(batch_size,
                               "Kernel launch requires uint32 for grid dim");
  int64_t feature_size = matrix_dim[1];
  auto stream = dev_ctx.stream();
  auto place = x.place();

#define PADDLE_LAUNCH_LAYERNORM_FWD(ScaleBiasT, IsScaleBiasSameDTypeWithX)    \
  do {                                                                        \
    switch (funcs::GetDesiredBlockDim(feature_size)) {                        \
      FIXED_BLOCK_DIM_CASE(                                                   \
          funcs::LayerNormForward<T, U, kBlockDim, IsScaleBiasSameDTypeWithX> \
          <<<static_cast<uint32_t>(batch_size), kBlockDim, 0, stream>>>(      \
              x_data,                                                         \
              static_cast<const ScaleBiasT*>(void_scale_data),                \
              static_cast<const ScaleBiasT*>(void_bias_data),                 \
              y_data,                                                         \
              mean_data,                                                      \
              var_data,                                                       \
              epsilon,                                                        \
              feature_size));                                                 \
      default:                                                                \
        PADDLE_THROW(common::errors::InvalidArgument(                         \
            "Product from begin_norm_axis to end must be larger than 1"));    \
        break;                                                                \
    }                                                                         \
  } while (0)

#define PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, feature_size)       \
  case (feature_size): {                                                     \
    constexpr int WARPS_N = feature_size < 1024 ? 1 : (feature_size / 1024); \
    constexpr int WARPS_M = 4 / WARPS_N;                                     \
    const int THREADS_PER_WARP = 32;                                         \
    const int BYTES_PER_LDG = 16;                                            \
    const int VecSize = BYTES_PER_LDG / sizeof(T);                           \
    const int THREADS_PER_CTA = WARPS_N * THREADS_PER_WARP * WARPS_M;        \
    const int ROWS_PER_CTA = WARPS_M;                                        \
    const int64_t grid = (batch_size + ROWS_PER_CTA - 1) / ROWS_PER_CTA;     \
    PADDLE_ENFORCE_LE_UINT32_MAX(grid, "layer_norm fast v1 grid");           \
    PADDLE_ENFORCE_LE_INT_MAX(batch_size, "layer_norm fast v1 batch_size");  \
    funcs::fast_ln_v1::fast_ln_v1_fwd_kernel<T,                              \
                                             U,                              \
                                             ScaleT,                         \
                                             VecSize,                        \
                                             WARPS_M,                        \
                                             WARPS_N,                        \
                                             BYTES_PER_LDG,                  \
                                             feature_size>                   \
        <<<static_cast<uint32_t>(grid), THREADS_PER_CTA, 0, stream>>>(       \
            static_cast<int>(batch_size),                                    \
            feature_size,                                                    \
            epsilon,                                                         \
            x_data,                                                          \
            static_cast<const ScaleT*>(void_scale_data),                     \
            static_cast<const ScaleT*>(void_bias_data),                      \
            mean_data,                                                       \
            var_data,                                                        \
            y_data);                                                         \
  } break

#define PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD(ScaleT)       \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 768);  \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 1024); \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 1280); \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 1536); \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 1792); \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 2048); \
  PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD_BASE(ScaleT, 4096)
  auto compute_dtype = CppTypeToDataType<U>::Type();
  auto kernel_variant = LayerNormKernelDispatch(scale_bias_dtype,
                                                x_dtype,
                                                y_dtype,
                                                compute_dtype,
                                                feature_size,
                                                x.numel(),
                                                scale,
                                                bias);

  switch (kernel_variant) {
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
    case LayerNormKernelVariant::FAST_LN_V2: {
      uint32_t hidden_size = static_cast<uint32_t>(feature_size);
      funcs::fast_ln_v2::LaunchNormFwd<T, Context>(dev_ctx,
                                                   stream,
                                                   place,
                                                   x_data,
                                                   void_scale_data,
                                                   void_bias_data,
                                                   y_data,
                                                   mean_data,
                                                   var_data,
                                                   scale_bias_dtype,
                                                   x_dtype,
                                                   y_dtype,
                                                   compute_dtype,
                                                   hidden_size,
                                                   batch_size,
                                                   feature_size,
                                                   epsilon);
      break;
    }
#endif
    case LayerNormKernelVariant::FAST_LN_V1:
      if (is_scale_bias_same_dtype_with_x) {
        switch (feature_size) {
          PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD(T);
          default:
            break;
        }
      } else {
        switch (feature_size) {
          PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD(U);
          default:
            break;
        }
      }
      break;
    case LayerNormKernelVariant::GENERIC:
    default:
#ifdef PADDLE_WITH_CUDA
      if ((x_dtype == scale_bias_dtype) &&
          (FLAGS_use_accuracy_compatible_kernel ||
           (!isPowerOfTwo(feature_size) && feature_size > 1024))) {
        LayerNormFwdCompatKernel<T, Context>(
            dev_ctx,
            x_data,
            valid_scale ? static_cast<const T*>(void_scale_data) : nullptr,
            valid_bias ? static_cast<const T*>(void_bias_data) : nullptr,
            epsilon,
            batch_size,
            feature_size,
            y_data,
            mean_data,
            var_data);
      } else if (FLAGS_use_fast_math && feature_size <= 1024 &&
                 (!std::is_same<T, int8_t>::value)) {
        // WarpShuffle intrinsics is involved in LaunchLayerNormKernel.
        LaunchLayerNormKernel<Context, T, U>(dev_ctx,
                                             x_data,
                                             y_data,
                                             void_scale_data,
                                             void_bias_data,
                                             mean_data,
                                             var_data,
                                             epsilon,
                                             batch_size,
                                             feature_size,
                                             valid_scale,
                                             valid_bias,
                                             is_scale_bias_same_dtype_with_x);
      } else {
#endif
        if (is_scale_bias_same_dtype_with_x) {
          PADDLE_LAUNCH_LAYERNORM_FWD(T, true);
        } else {
          PADDLE_LAUNCH_LAYERNORM_FWD(U, false);
        }
#ifdef PADDLE_WITH_CUDA
      }
#endif
      break;
  }

#undef PADDLE_LAUNCH_LAYERNORM_FWD
#undef PADDLE_LAUNCH_FAST_LAYERNORM_V1_FWD
}
#ifdef _WIN32
template PADDLE_API void LayerNormKernel<float, GPUContext>(
    const GPUContext& dev_ctx,
    const DenseTensor& x,
    const optional<DenseTensor>& scale_opt,
    const optional<DenseTensor>& bias_opt,
    double epsilon,
    int begin_norm_axis,
    DenseTensor* y,
    DenseTensor* mean,
    DenseTensor* var);
template PADDLE_API void LayerNormKernel<dtype::float16, GPUContext>(
    const GPUContext& dev_ctx,
    const DenseTensor& x,
    const optional<DenseTensor>& scale_opt,
    const optional<DenseTensor>& bias_opt,
    double epsilon,
    int begin_norm_axis,
    DenseTensor* y,
    DenseTensor* mean,
    DenseTensor* var);
template PADDLE_API void LayerNormKernel<double, GPUContext>(
    const GPUContext& dev_ctx,
    const DenseTensor& x,
    const optional<DenseTensor>& scale_opt,
    const optional<DenseTensor>& bias_opt,
    double epsilon,
    int begin_norm_axis,
    DenseTensor* y,
    DenseTensor* mean,
    DenseTensor* var);
#endif
}  // namespace phi

#ifdef PADDLE_WITH_HIP
// MIOPEN do not support double
PD_REGISTER_KERNEL(
    layer_norm, GPU, ALL_LAYOUT, phi::LayerNormKernel, float, phi::float16) {
  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
}
#elif CUDNN_VERSION_MIN(8, 1, 0)
PD_REGISTER_KERNEL(layer_norm,
                   GPU,
                   ALL_LAYOUT,
                   phi::LayerNormKernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {
  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
}
#else
PD_REGISTER_KERNEL(layer_norm,
                   GPU,
                   ALL_LAYOUT,
                   phi::LayerNormKernel,
                   float,
                   double,
                   phi::float16) {
  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
}
#endif