#include "../common.cuh"

// LayerNorm: y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
// Each block processes one row of shape [hidden_size].

__global__ void layernorm_f32(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ out,
    int hidden_size, float eps
) {
    int row = blockIdx.x;
    const float* x_row = x + row * hidden_size;
    float* out_row = out + row * hidden_size;

    // Welford online: compute mean and variance in one pass
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float v = x_row[i];
        local_sum += v;
        local_sum_sq += v * v;
    }
    local_sum = block_reduce_sum(local_sum);
    local_sum_sq = block_reduce_sum(local_sum_sq);

    __shared__ float s_mean, s_inv_std;
    if (threadIdx.x == 0) {
        float mean = local_sum / hidden_size;
        float var = local_sum_sq / hidden_size - mean * mean;
        s_mean = mean;
        s_inv_std = rsqrtf(var + eps);
    }
    __syncthreads();

    float mean = s_mean;
    float inv_std = s_inv_std;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        out_row[i] = gamma[i] * (x_row[i] - mean) * inv_std + beta[i];
    }
}

__global__ void layernorm_bf16(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ gamma,
    const __nv_bfloat16* __restrict__ beta,
    __nv_bfloat16* __restrict__ out,
    int hidden_size, float eps
) {
    int row = blockIdx.x;
    const __nv_bfloat16* x_row = x + row * hidden_size;
    __nv_bfloat16* out_row = out + row * hidden_size;

    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float v = __bfloat162float(x_row[i]);
        local_sum += v;
        local_sum_sq += v * v;
    }
    local_sum = block_reduce_sum(local_sum);
    local_sum_sq = block_reduce_sum(local_sum_sq);

    __shared__ float s_mean, s_inv_std;
    if (threadIdx.x == 0) {
        float mean = local_sum / hidden_size;
        float var = local_sum_sq / hidden_size - mean * mean;
        s_mean = mean;
        s_inv_std = rsqrtf(var + eps);
    }
    __syncthreads();

    float mean = s_mean;
    float inv_std = s_inv_std;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float v = __bfloat162float(x_row[i]);
        float g = __bfloat162float(gamma[i]);
        float b = __bfloat162float(beta[i]);
        out_row[i] = __float2bfloat16(g * (v - mean) * inv_std + b);
    }
}

extern "C" {

void launch_layernorm_f32(const void* x, const void* gamma, const void* beta,
                          void* out, int rows, int hidden_size, float eps, void* stream) {
    int block = (hidden_size < 1024) ? hidden_size : 1024;
    layernorm_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
        (const float*)x, (const float*)gamma, (const float*)beta,
        (float*)out, hidden_size, eps);
}

void launch_layernorm_bf16(const void* x, const void* gamma, const void* beta,
                           void* out, int rows, int hidden_size, float eps, void* stream) {
    int block = (hidden_size < 1024) ? hidden_size : 1024;
    layernorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma, (const __nv_bfloat16*)beta,
        (__nv_bfloat16*)out, hidden_size, eps);
}

}