/**
 * Copyright (c) 2016-present, Facebook, Inc.
 *
 * 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) 2017, NVIDIA CORPORATION. All rights reserved.
// NVIDIA/apex is licensed under the
// BSD 3 - Clause "New" or "Revised" License
//

/* Modifications Copyright (c) Microsoft. */

#include "core/providers/cuda/cu_inc/common.cuh"
#include "layer_norm_impl.h"
#include "core/providers/cpu/nn/layer_norm_helper.h"

namespace onnxruntime {
namespace cuda {

template <typename U, bool simplified>
__device__ void cuWelfordOnlineSum(
    const U curr,
    U& mu,
    U& sigma2,
    U& count) {
  count = count + U(1);
  U delta = curr - mu;
  U lmean = mu + delta / count;
  mu = lmean;
  if (simplified) {
    sigma2 = sigma2 + curr * curr;
  } else {
    U delta2 = curr - lmean;
    sigma2 = sigma2 + delta * delta2;
  }
}

template <typename U, bool simplified>
__device__ void cuChanOnlineSum(
    const U muB,
    const U sigma2B,
    const U countB,
    U& mu,
    U& sigma2,
    U& count) {
  U delta = muB - mu;
  U nA = count;
  U nB = countB;
  count = count + countB;
  U nX = count;
  if (nX > U(0)) {
    nA = nA / nX;
    nB = nB / nX;
    mu = nA * mu + nB * muB;
    if (simplified) {
      sigma2 = sigma2 + sigma2B;
    } else {
      sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;
    }
  } else {
    mu = U(0);
    sigma2 = U(0);
  }
}

template <typename T, typename U, bool simplified>
__device__ void cuWelfordMuSigma2(
    const T* __restrict__ vals,
    const int n1,
    const int n2,
    const int i1,
    U& mu,
    U& sigma2,
    U* buf,
    const T* __restrict__ skip,
    const T* __restrict__ bias) {
  // Assumptions:
  // 1) blockDim.x == GPU_WARP_SIZE
  // 2) Tensor is contiguous
  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
  //
  // compute variance and mean over n2
  U count = U(0);
  mu = U(0);
  sigma2 = U(0);
  if (i1 < n1) {
    // one warp normalizes one n1 index,
    // synchronization is implicit
    // initialize with standard Welford algorithm
    const int numx = blockDim.x * blockDim.y;
    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
    const T* lvals = vals + i1 * n2;
    const T* skip_vals = (skip != nullptr) ? skip + i1 * n2 : nullptr;
    int l = 4 * thrx;
    for (; l + 3 < n2; l += 4 * numx) {
      for (int k = 0; k < 4; ++k) {
        U curr = static_cast<U>(lvals[l + k]);

        if (bias != nullptr) {
          curr += static_cast<U>(bias[l + k]);
        }

        if (skip_vals != nullptr) {
          curr += static_cast<U>(skip_vals[l + k]);
        }

        cuWelfordOnlineSum<U, simplified>(curr, mu, sigma2, count);
      }
    }
    for (; l < n2; ++l) {
      U curr = static_cast<U>(lvals[l]);

      if (bias != nullptr) {
        curr += static_cast<U>(bias[l]);
      }

      if (skip_vals != nullptr) {
        curr += static_cast<U>(skip_vals[l]);
      }

      cuWelfordOnlineSum<U, simplified>(curr, mu, sigma2, count);
    }
// intra-warp reductions
#pragma unroll
    for (int stride = GPU_WARP_SIZE / 2; stride > 0; stride /= 2) {
      U muB = WARP_SHFL_DOWN(mu, stride);
      U countB = WARP_SHFL_DOWN(count, stride);
      U sigma2B = WARP_SHFL_DOWN(sigma2, stride);
      cuChanOnlineSum<U, simplified>(muB, sigma2B, countB, mu, sigma2, count);
    }

    // threadIdx.x == 0 has correct values for each warp
    // inter-warp reductions
    if (blockDim.y > 1) {
      U* ubuf = (U*)buf;
      U* ibuf = (U*)(ubuf + blockDim.y);
      for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {
        // upper half of warps write to shared
        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {
          const int wrt_y = threadIdx.y - offset;
          ubuf[2 * wrt_y] = mu;
          ubuf[2 * wrt_y + 1] = sigma2;
          ibuf[wrt_y] = count;
        }
        __syncthreads();
        // lower half merges
        if (threadIdx.x == 0 && threadIdx.y < offset) {
          U muB = ubuf[2 * threadIdx.y];
          U sigma2B = ubuf[2 * threadIdx.y + 1];
          U countB = ibuf[threadIdx.y];
          cuChanOnlineSum<U, simplified>(muB, sigma2B, countB, mu, sigma2, count);
        }
        __syncthreads();
      }
      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
      if (threadIdx.x == 0 && threadIdx.y == 0) {
        ubuf[0] = mu;
        ubuf[1] = sigma2;
      }
      __syncthreads();
      mu = ubuf[0];
      sigma2 = ubuf[1] / U(n2);
      // don't care about final value of count, we know count == n2
    } else {
      mu = WARP_SHFL(mu, 0);
      sigma2 = WARP_SHFL(sigma2 / U(n2), 0);
    }
  }
}

template <bool simplified>
__device__ void cuWelfordMuSigma2(
    const half* __restrict__ vals,
    const int n1,
    const int n2,
    const int i1,
    float& mu,
    float& sigma2,
    float* buf) {
  // Assumptions:
  // 1) blockDim.x == GPU_WARP_SIZE
  // 2) Tensor is contiguous
  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
  //
  // compute variance and mean over n2
  float count = 0.0f;
  mu = float(0);
  sigma2 = float(0);
  if (i1 < n1) {
    // one warp normalizes one n1 index,
    // synchronization is implicit
    // initialize with standard Welford algorithm
    const int numx = blockDim.x * blockDim.y;
    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
    const half* lvals = vals + i1 * n2;
    int l = 8 * thrx;
    if ((((size_t)lvals) & 3) != 0) {
      // 16 bit alignment
      // first thread consumes first point
      if (thrx == 0) {
        float curr = static_cast<float>(lvals[0]);
        cuWelfordOnlineSum<float, simplified>(curr, mu, sigma2, count);
      }
      ++l;
    }
    // at this point, lvals[l] are 32 bit aligned for all threads.
    for (; l + 7 < n2; l += 8 * numx) {
      for (int k = 0; k < 8; k += 2) {
        float2 curr = __half22float2(*((__half2*)(lvals + l + k)));
        cuWelfordOnlineSum<float, simplified>(curr.x, mu, sigma2, count);
        cuWelfordOnlineSum<float, simplified>(curr.y, mu, sigma2, count);
      }
    }
    for (; l < n2; ++l) {
      float curr = static_cast<float>(lvals[l]);
      cuWelfordOnlineSum<float, simplified>(curr, mu, sigma2, count);
    }
// intra-warp reductions
#pragma unroll
    for (int stride = GPU_WARP_SIZE / 2; stride > 0; stride /= 2) {
      float muB = WARP_SHFL_DOWN(mu, stride);
      float countB = WARP_SHFL_DOWN(count, stride);
      float sigma2B = WARP_SHFL_DOWN(sigma2, stride);
      cuChanOnlineSum<float, simplified>(muB, sigma2B, countB, mu, sigma2, count);
    }

    // threadIdx.x == 0 has correct values for each warp
    // inter-warp reductions
    if (blockDim.y > 1) {
      float* ubuf = (float*)buf;
      float* ibuf = (float*)(ubuf + blockDim.y);
      for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {
        // upper half of warps write to shared
        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {
          const int wrt_y = threadIdx.y - offset;
          ubuf[2 * wrt_y] = mu;
          ubuf[2 * wrt_y + 1] = sigma2;
          ibuf[wrt_y] = count;
        }
        __syncthreads();
        // lower half merges
        if (threadIdx.x == 0 && threadIdx.y < offset) {
          float muB = ubuf[2 * threadIdx.y];
          float sigma2B = ubuf[2 * threadIdx.y + 1];
          float countB = ibuf[threadIdx.y];
          cuChanOnlineSum<float, simplified>(muB, sigma2B, countB, mu, sigma2, count);
        }
        __syncthreads();
      }
      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
      if (threadIdx.x == 0 && threadIdx.y == 0) {
        ubuf[0] = mu;
        ubuf[1] = sigma2;
      }
      __syncthreads();
      mu = ubuf[0];
      sigma2 = ubuf[1] / float(n2);
      // don't care about final value of count, we know count == n2
    } else {
      mu = WARP_SHFL(mu, 0);
      sigma2 = WARP_SHFL(sigma2 / float(n2), 0);
    }
  }
}

template <typename U>
__device__ U rsqrt(U v) {
  return U(1) / sqrt(v);
}
template <>
__device__ float rsqrt(float v) {
  return rsqrtf(v);
}
template <>
__device__ double rsqrt(double v) {
  return rsqrt(v);
}

namespace {
// This is the un-specialized struct.  Note that we prevent instantiation of this
// struct by putting an undefined symbol in the function body so it won't compile.
//  template <typename T>
//  struct SharedMemory
//  {
//      // Ensure that we won't compile any un-specialized types
//      __device__ T *getPointer()
//      {
//          extern __device__ void error(void);
//          error();
//          return nullptr;
//      }
//  };
// https://github.com/NVIDIA/apex/issues/246
template <typename T>
struct SharedMemory;

template <>
struct SharedMemory<float> {
  __device__ float* getPointer() {
    extern __shared__ float s_float[];
    return s_float;
  }
};

template <>
struct SharedMemory<double> {
  __device__ double* getPointer() {
    extern __shared__ double s_double[];
    return s_double;
  }
};
}  // namespace

template <typename T, typename U, typename V, bool simplified>
__global__ void cuApplyLayerNorm(
    V* __restrict__ output_vals,
    U* __restrict__ mean,
    U* __restrict__ inv_std_dev,
    const T* __restrict__ vals,
    const int n1,
    const int n2,
    const U epsilon,
    const V* __restrict__ gamma,
    const V* __restrict__ beta,
    int broadcast_param,
    const T* __restrict__ skip,
    const T* __restrict__ bias,
    T* __restrict__ skip_input_bias_add_output) {
  // Assumptions:
  // 1) blockDim.x == GPU_WARP_SIZE
  // 2) Tensors are contiguous
  //
  for (int i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {
    SharedMemory<U> shared;
    U* buf = shared.getPointer();
    U mu, sigma2;
    cuWelfordMuSigma2<T, U, simplified>(vals, n1, n2, i1, mu, sigma2, buf, skip, bias);
    const int offset = i1 * n2;
    const T* lvals = vals + offset;
    const T* skip_vals = (skip != nullptr) ? skip + offset : nullptr;

    V* ovals = output_vals + offset;
    T* skip_input_bias_add_ovals = (skip_input_bias_add_output != nullptr) ? skip_input_bias_add_output + offset : nullptr;
    U c_inv_std_dev = rsqrt(sigma2 + epsilon);

    // Compute the offset of gamma and beta to support broadcasting.
    int gamma_beta_offset = LAYER_NORM_SCALE_BIAS_OFFSET(broadcast_param, i1, n2);

    const int numx = blockDim.x * blockDim.y;
    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
    for (int i = thrx; i < n2; i += numx) {
      U curr = static_cast<U>(lvals[i]);

      if (bias != nullptr) {
        curr += static_cast<U>(bias[i]);
      }

      if (skip_vals != nullptr) {
        curr += static_cast<U>(skip_vals[i]);
      }

      int index = gamma_beta_offset + i;
      U gamma_i = (gamma != nullptr) ? (U)gamma[index] : (U)1;
      U beta_i = (beta != nullptr) ? (U)beta[index] : (U)0;

      if (simplified) {
        ovals[i] = static_cast<V>(gamma_i * c_inv_std_dev * curr);
      } else {
        ovals[i] = static_cast<V>(gamma_i * c_inv_std_dev * (curr - mu) + beta_i);
      }

      if (skip_input_bias_add_ovals != nullptr) {
        skip_input_bias_add_ovals[i] = static_cast<T>(curr);
      }
    }
    if (threadIdx.x == 0 && threadIdx.y == 0) {
      if (mean != nullptr) mean[i1] = mu;
      if (inv_std_dev != nullptr) inv_std_dev[i1] = c_inv_std_dev;
    }
  }
}

int32_t round_up_power_of_2(int32_t v) {
  v--;
  v |= v >> 1;
  v |= v >> 2;
  v |= v >> 4;
  v |= v >> 8;
  v |= v >> 16;
  v++;
  return v;
}

template <typename T, typename U, typename V, bool simplified>
void HostApplyLayerNorm(
    const cudaDeviceProp& prop,
    cudaStream_t stream,
    V* output,
    U* mean,
    U* inv_std_dev,
    const T* input,
    int n1,
    int n2,
    double epsilon,
    const V* gamma,
    const V* beta,
    int broadcast_param,
    const T* skip,
    const T* bias,
    T* skip_input_bias_add_output) {
  const int maxGridY = prop.maxGridSize[1];
  const int warp_size = prop.warpSize;
  ORT_ENFORCE(warp_size == GPU_WARP_SIZE_HOST);

  // Be careful for the logic on treads_y calc:
  //     * 4 is current implementation related, yet it does not using vectorized load
  //     * Do not using maxTreads as it will cause resource issue.
  int threads_y = std::min(round_up_power_of_2((n2 + 4 * warp_size - 1) / (4 * warp_size)),
                           prop.maxThreadsPerBlock / (warp_size * 2));
  // threads_y could be 16, 8, 4, 2, 1. Calculate total number of rows that could be run in parallel.
  int parallel_rows = prop.multiProcessorCount * ((prop.maxThreadsPerBlock / 2) / warp_size / threads_y);
  while (threads_y > 2 && parallel_rows < n1) {
    threads_y /= 2;
    parallel_rows *= 2;
  }
  dim3 threads(warp_size, threads_y, 1);
#ifdef __HIP_PLATFORM_HCC__
  // Optimization for ROCm MI100
  threads.y = 1;
#endif
  const dim3 blocks(1, std::min<unsigned int>(n1, maxGridY), 1);
  int nshared =
      threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;
  cuApplyLayerNorm<T, U, V, simplified><<<blocks, threads, nshared, stream>>>(
      output,
      mean,
      inv_std_dev,
      input,
      n1, n2,
      U(epsilon),
      gamma, beta, broadcast_param,
      skip, bias, skip_input_bias_add_output);
}

#define LAYERNORM_LINEAR_IMPL(T, U, V, simplified)                                                                  \
  template void HostApplyLayerNorm<T, U, V, simplified>(const cudaDeviceProp& prop, cudaStream_t stream, V* output, \
                                                        U* mean, U* inv_std_dev, const T* input, int n1, int n2,    \
                                                        double epsilon, const V* gamma, const V* beta,              \
                                                        int broadcast_param,                                        \
                                                        const T* skip, const T* bias, T* skip_input_bias_add_output);

LAYERNORM_LINEAR_IMPL(float, float, float, true)
LAYERNORM_LINEAR_IMPL(half, float, half, true)
LAYERNORM_LINEAR_IMPL(double, double, double, true)
LAYERNORM_LINEAR_IMPL(float, float, half, true)
LAYERNORM_LINEAR_IMPL(half, float, float, true)
LAYERNORM_LINEAR_IMPL(float, float, float, false)
LAYERNORM_LINEAR_IMPL(half, float, half, false)
LAYERNORM_LINEAR_IMPL(double, double, double, false)
LAYERNORM_LINEAR_IMPL(double, float, double, false)
LAYERNORM_LINEAR_IMPL(float, float, half, false)
LAYERNORM_LINEAR_IMPL(half, float, float, false)
LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, true)
LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, false)

}  // namespace cuda
}  // namespace onnxruntime
