#include "LayerNormExecution.hpp"
namespace MNN {
namespace CUDA {

#define CUDA_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x)

template <typename T>
__global__ 
void input_layernorm(T* out, const T* input, const float* gamma, const float* beta, int m, int n, const float epsilon, int sumPerKnl, bool RMSNorm)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  s_mean = 0;

  if(!RMSNorm){
    for(int idx=0; idx<sumPerKnl && idx*256 + tid < n; idx++) {
     local_out += (float)(input[blockIdx.x * n + idx*256 + tid]);
    }

    mean = blockReduceSum<float>(local_out);
    if(threadIdx.x == 0)
      s_mean = mean / n;
    __syncthreads();
  }
  mean = s_mean;

  float var_tmp = 0.0f;
  for(int idx=0; idx<sumPerKnl && idx*256 + tid < n; idx++) {
    var_tmp += (((float)input[blockIdx.x * n + idx*256 + tid] - mean) * ((float)input[blockIdx.x * n + idx*256 + tid] - mean));
  }
  variance = blockReduceSum<float>(var_tmp);
  if(threadIdx.x == 0)
    s_variance = variance / n + epsilon;
  __syncthreads();

  for(int idx=0; idx<sumPerKnl && idx*256 + tid < n; idx++) {
    float res = (((float)input[blockIdx.x * n + idx*256 + tid] - mean) * rsqrtf(s_variance));
    if(gamma != nullptr && beta != nullptr) {
        res = res * (float)(__ldg(&gamma[idx*256 + tid])) + (float)(__ldg(&beta[idx*256 + tid]));
    }
    out[blockIdx.x * n + idx*256+tid] = (T)res;
  }
}

template <typename T>
__global__ 
void input_layernorm_320(T* out, const T* input, const float* gamma, const float* beta, int m, int n, const float epsilon, bool RMSNorm)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  s_mean = 0;

  float value_tmp[5];
  value_tmp[0] = input[blockIdx.x * n + 0*64 + tid];
  value_tmp[1] = input[blockIdx.x * n + 1*64 + tid];
  value_tmp[2] = input[blockIdx.x * n + 2*64 + tid];
  value_tmp[3] = input[blockIdx.x * n + 3*64 + tid];
  value_tmp[4] = input[blockIdx.x * n + 4*64 + tid];

  if(!RMSNorm){
    for(int idx=0; idx<5; idx++) {
     local_out += value_tmp[idx];
    }

    mean = blockReduceSum<float>(local_out);
    if(threadIdx.x == 0)
      s_mean = mean / n;
    __syncthreads();
  }
  mean = s_mean;

  float var_tmp = 0.0f;
  for(int idx=0; idx<5; idx++) {
    var_tmp += ((value_tmp[idx] - mean) * (value_tmp[idx] - mean));
  }
  variance = blockReduceSum<float>(var_tmp);
  if(threadIdx.x == 0)
    s_variance = variance / n + epsilon;
  __syncthreads();

  for(int idx=0; idx<5; idx++) {
    float res = ((value_tmp[idx] - mean) * rsqrtf(s_variance));
    if(gamma != nullptr && beta != nullptr) {
        res = res * (float)(__ldg(&gamma[idx*64 + tid])) + (float)(__ldg(&beta[idx*64 + tid]));
    }
    out[blockIdx.x * n + idx*64+tid] = (T)res;
  }
}


template <typename T>
__global__ 
void input_layernorm_2048(T* out, const T* input, const float* gamma, const float* beta, int m, int n, const float epsilon, bool RMSNorm)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  s_mean = 0;

  float value_tmp[8];
  value_tmp[0] = input[blockIdx.x * 2048 + 0*256 + tid];
  value_tmp[1] = input[blockIdx.x * 2048 + 1*256 + tid];
  value_tmp[2] = input[blockIdx.x * 2048 + 2*256 + tid];
  value_tmp[3] = input[blockIdx.x * 2048 + 3*256 + tid];
  value_tmp[4] = input[blockIdx.x * 2048 + 4*256 + tid];
  value_tmp[5] = input[blockIdx.x * 2048 + 5*256 + tid];
  value_tmp[6] = input[blockIdx.x * 2048 + 6*256 + tid];
  value_tmp[7] = input[blockIdx.x * 2048 + 7*256 + tid];

  if(!RMSNorm){
    #pragma unroll(8)
    for(int idx=0; idx<8; idx++) {
      local_out += (float)value_tmp[idx];
    }

    mean = blockReduceSum<float>(local_out);
    if(threadIdx.x == 0)
     s_mean = mean / n;
    __syncthreads();
  }
  mean = s_mean;

  float var_tmp = 0.0f;

  #pragma unroll(8)
  for(int idx=0; idx<8; idx++) {
    var_tmp += ((value_tmp[idx] - mean) * (value_tmp[idx] - mean));
  }
  variance = blockReduceSum<float>(var_tmp);
  if(threadIdx.x == 0)
    s_variance = variance / n + epsilon;
  __syncthreads();

  #pragma unroll(8)
  for(int idx=0; idx<8; idx++) {
    float res = ((value_tmp[idx] - mean) * rsqrtf(s_variance));
    if(gamma != nullptr && beta != nullptr) {
        res = res * (float)(__ldg(&gamma[idx*256 + tid])) + (float)(__ldg(&beta[idx*256 + tid]));
    }
    out[blockIdx.x * 2048 + idx*256+tid] = (T)res;
  }
}


template <typename T>
__global__ 
void input_layernorm_1024(T* out, const T* input, const float* gamma, const float* beta, int m, int n, const float epsilon, bool RMSNorm)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  s_mean = 0;

  float value_tmp[4];
  value_tmp[0] = input[blockIdx.x * 1024 + 0*256 + tid];
  value_tmp[1] = input[blockIdx.x * 1024 + 1*256 + tid];
  value_tmp[2] = input[blockIdx.x * 1024 + 2*256 + tid];
  value_tmp[3] = input[blockIdx.x * 1024 + 3*256 + tid];

  if(!RMSNorm){
    #pragma unroll(4)
    for(int idx=0; idx<4; idx++) {
      local_out += (float)value_tmp[idx];
   }

   mean = blockReduceSum<float>(local_out);
   if(threadIdx.x == 0)
      s_mean = mean / n;
    __syncthreads();
  }
  mean = s_mean;

  float var_tmp = 0.0f;

  #pragma unroll(4)
  for(int idx=0; idx<4; idx++) {
    var_tmp += ((value_tmp[idx] - mean) * (value_tmp[idx] - mean));
  }
  variance = blockReduceSum<float>(var_tmp);
  if(threadIdx.x == 0)
    s_variance = variance / n + epsilon;
  __syncthreads();

  #pragma unroll(4)
  for(int idx=0; idx<4; idx++) {
    float res = ((value_tmp[idx] - mean) * rsqrtf(s_variance));
    if(gamma != nullptr && beta != nullptr) {
        res = res * (float)(__ldg(&gamma[idx*256 + tid])) + (float)(__ldg(&beta[idx*256 + tid]));
    }
    out[blockIdx.x * 1024 + idx*256+tid] = (T)res;
  }
}


template <typename T>
__global__ 
void input_layernorm_512(T* out, const T* input, const float* gamma, const float* beta, int m, int n, const float epsilon, bool RMSNorm)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  s_mean = 0;

  float value_tmp[2];
  value_tmp[0] = input[blockIdx.x * 512 + 0*256 + tid];
  value_tmp[1] = input[blockIdx.x * 512 + 1*256 + tid];

  if(!RMSNorm){
    local_out += (float)value_tmp[0];
    local_out += (float)value_tmp[1];

    mean = blockReduceSum<float>(local_out);
    if(threadIdx.x == 0)
      s_mean = mean / n;
    __syncthreads();
  }
  mean = s_mean;

  float var_tmp = 0.0f;
  var_tmp += ((value_tmp[0] - mean) * (value_tmp[0] - mean));
  var_tmp += ((value_tmp[1] - mean) * (value_tmp[1] - mean));

  variance = blockReduceSum<float>(var_tmp);
  if(threadIdx.x == 0)
    s_variance = variance / n + epsilon;
  __syncthreads();

  float res0 = ((value_tmp[0] - mean) * rsqrtf(s_variance));
  float res1 = ((value_tmp[1] - mean) * rsqrtf(s_variance));

  if(gamma != nullptr && beta != nullptr) {
      res0 = res0 * (float)(__ldg(&gamma[0*256 + tid])) + (float)(__ldg(&beta[0*256 + tid]));
      res1 = res1 * (float)(__ldg(&gamma[1*256 + tid])) + (float)(__ldg(&beta[1*256 + tid]));
  }

  out[blockIdx.x * 512 + 0*256+tid] = (T)res0;
  out[blockIdx.x * 512 + 1*256+tid] = (T)res1;
}


// P3: Adaptive LayerNorm/RMSNorm kernel using configurable thread count (up to 1024)
// Each block handles one row; threads cooperatively process all elements
template <typename T>
__global__
void input_layernorm_adaptive(T* out, const T* input, const float* gamma, const float* beta,
                               int m, int n, const float epsilon, bool RMSNorm)
{
    const int tid = threadIdx.x;
    const int num_threads = blockDim.x;
    const int row = blockIdx.x;

    __shared__ float s_mean;
    __shared__ float s_variance;

    float mean = 0.0f;
    s_mean = 0.0f;

    // First pass: compute mean (for LayerNorm) or skip (for RMSNorm)
    if (!RMSNorm) {
        float local_sum = 0.0f;
        for (int i = tid; i < n; i += num_threads) {
            local_sum += (float)input[row * n + i];
        }
        mean = blockReduceSum<float>(local_sum);
        if (tid == 0) s_mean = mean / n;
        __syncthreads();
    }
    mean = s_mean;

    // Second pass: compute variance
    float var_sum = 0.0f;
    for (int i = tid; i < n; i += num_threads) {
        float val = (float)input[row * n + i] - mean;
        var_sum += val * val;
    }
    float variance = blockReduceSum<float>(var_sum);
    if (tid == 0) s_variance = variance / n + epsilon;
    __syncthreads();

    float inv_std = rsqrtf(s_variance);

    // Third pass: normalize and write output
    for (int i = tid; i < n; i += num_threads) {
        float res = ((float)input[row * n + i] - mean) * inv_std;
        if (gamma != nullptr && beta != nullptr) {
            res = res * (float)__ldg(&gamma[i]) + (float)__ldg(&beta[i]);
        }
        out[row * n + i] = (T)res;
    }
}

template<typename T>
__global__ void LAYERNORM(const int count, const int outside, const int inside, const float epsilon,
                          const T* in, T* out, const float* gamma_data, const float* beta_data, bool RMSNorm) {
    CUDA_KERNEL_LOOP(i, count) {
        const int o = i / inside;
        const int index = i % inside;
        const T* inner_input = in + o * inside;
        T* inner_output = out + o * inside;
        float mean = 0.0f;
        if(!RMSNorm){
          float sum = 0.f;
          for (int j = 0; j < inside; ++j) {
            sum += (float)inner_input[j];
          }
          mean = sum / inside;
        }
        float square_sum = 0.f;
        for (int j = 0; j < inside; ++j) {
            square_sum += ((float)inner_input[j] - mean) * ((float)inner_input[j] - mean);
        }
        float variable = square_sum / inside;
        variable = 1.f / sqrt(variable + epsilon);

        float res = ((float)inner_input[index] - mean) * variable;
        if(gamma_data != nullptr && beta_data != nullptr) {
            res = res * gamma_data[index] + beta_data[index];
        }
        inner_output[index] = (T)res;
    }
}

LayerNormExecution::LayerNormExecution(const LayerNorm* layer_norm_param, Backend *backend) : Execution(backend) {
    if (nullptr != layer_norm_param->axis()) {
        mAxises = layer_norm_param->axis()->size();
    }

    mEps = layer_norm_param->epsilon();
    mGroup = layer_norm_param->group();
    RMSNorm = layer_norm_param->useRMSNorm();

    if (layer_norm_param->gamma() && layer_norm_param->beta()) {
        int size = layer_norm_param->gamma()->size();
        mGammaTensor.reset(Tensor::createDevice<int32_t>({size}));
        auto status = backend->onAcquireBuffer(mGammaTensor.get(), Backend::STATIC);
        if (!status) {
            MNN_ERROR("Out of memory when gamma is acquired in CudaLayerNorm.\n");
        }

        mDeviceGamma = (void *)mGammaTensor.get()->buffer().device;
        const float* gamma_data = layer_norm_param->gamma()->data();
        cudaMemcpy(mDeviceGamma, gamma_data, size * sizeof(float), cudaMemcpyHostToDevice);

        if (layer_norm_param->beta()->size() != size) {
            MNN_ERROR("Size of gamma and beta are not match in CudaLayerNorm.\n");
        }
        mBetaTensor.reset(Tensor::createDevice<int32_t>({size}));
        status = backend->onAcquireBuffer(mBetaTensor.get(), Backend::STATIC);
        if (!status) {
            MNN_ERROR("Out of memory when beta is acquired in CudaLayerNorm.\n");
        }

        mDeviceBeta = (void *)mBetaTensor.get()->buffer().device;
        const float* beta_data = layer_norm_param->beta()->data();
        cudaMemcpy(mDeviceBeta, beta_data, size * sizeof(float), cudaMemcpyHostToDevice);
    }
}
LayerNormExecution::~LayerNormExecution() {
    // Do nothing
}

ErrorCode LayerNormExecution::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    MNN_ASSERT(inputs.size() == 1);
    MNN_ASSERT(outputs.size() == 1);
    auto input = inputs[0];

    mOutside = 1;
    mInside = 1;
    int rank = input->dimensions();
    if (mGroup > 1) {
        mOutside = input->length(0) * mGroup;
        for (int i = 1; i < rank; i++) {
            mInside *= input->length(i);
        }
        mInside /= mGroup;
        return NO_ERROR;
    }
    for (int i = 0; i < rank - mAxises; ++i) {
        mOutside *= input->length(i);
    }
    for (int i = rank - mAxises; i < rank; ++i) {
        mInside *= input->length(i);
    }
    return NO_ERROR;
}

ErrorCode LayerNormExecution::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    auto runtime = static_cast<CUDABackend*>(backend())->getCUDARuntime();
 
    int block_num = runtime->blocks_num(mOutside*mInside);
    int threads_num = runtime->threads_num();
    auto input_addr = (void*)inputs[0]->deviceId();
    auto output_addr = (void*)outputs[0]->deviceId();

    //printf("ln:%d-%d\n", mOutside, mInside);
    // P3: Choose optimal thread count based on hidden dimension
    // For large hidden dims (common in LLMs: 3584, 4096, 5120, 8192), use more threads
    auto launchLayerNorm = [&](auto* output_typed, const auto* input_typed) {
        using TT = std::remove_pointer_t<decltype(output_typed)>;
        if (mInside < 128) {
            LAYERNORM<<<block_num, threads_num>>>(mOutside*mInside, mOutside, mInside, mEps,
                input_typed, output_typed, (const float *)mDeviceGamma, (const float *)mDeviceBeta, RMSNorm);
        } else if (mInside == 2048) {
            input_layernorm_2048<<<mOutside, 256>>>(output_typed, input_typed, (const float *)mDeviceGamma,
                (const float *)mDeviceBeta, mOutside, mInside, mEps, RMSNorm);
        } else if (mInside == 1024) {
            input_layernorm_1024<<<mOutside, 256>>>(output_typed, input_typed, (const float *)mDeviceGamma,
                (const float *)mDeviceBeta, mOutside, mInside, mEps, RMSNorm);
        } else if (mInside == 512) {
            input_layernorm_512<<<mOutside, 256>>>(output_typed, input_typed, (const float *)mDeviceGamma,
                (const float *)mDeviceBeta, mOutside, mInside, mEps, RMSNorm);
        } else if (mInside == 320) {
            input_layernorm_320<<<mOutside, 64>>>(output_typed, input_typed, (const float *)mDeviceGamma,
                (const float *)mDeviceBeta, mOutside, mInside, mEps, RMSNorm);
        } else if (mInside > 2048) {
            // P3: For large hidden dims (3584, 4096, 5120, 8192, etc.), use adaptive kernel with more threads
            // Use up to 1024 threads to minimize iterations per thread
            int adaptive_threads = 256;
            if (mInside > 2048) adaptive_threads = 512;
            if (mInside > 4096) adaptive_threads = 1024;
            input_layernorm_adaptive<<<mOutside, adaptive_threads>>>(output_typed, input_typed,
                (const float *)mDeviceGamma, (const float *)mDeviceBeta, mOutside, mInside, mEps, RMSNorm);
        } else {
            int sumPerKnl = (mInside+255) / 256;
            input_layernorm<<<mOutside, 256>>>(output_typed, input_typed, (const float *)mDeviceGamma,
                (const float *)mDeviceBeta, mOutside, mInside, mEps, sumPerKnl, RMSNorm);
        }
    };

    if (static_cast<CUDABackend*>(backend())->useFp16()) {
        launchLayerNorm((half *)output_addr, (const half *)input_addr);
        return NO_ERROR;
    }

    launchLayerNorm((float *)output_addr, (const float *)input_addr);
    return NO_ERROR;
}

class LayerNormCreator : public CUDABackend::Creator {
public:
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                                const MNN::Op* op, Backend* backend) const override {
        auto param = op->main_as_LayerNorm();
        return new LayerNormExecution(param, backend);
    }
};

static CUDACreatorRegister<LayerNormCreator> __init(OpType_LayerNorm);

}
}