#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;

    // Pass 1: compute mean
    float local_sum = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        local_sum += x_row[i];
    }
    local_sum = block_reduce_sum(local_sum);

    __shared__ float s_mean, s_inv_std;
    if (threadIdx.x == 0) {
        s_mean = local_sum / hidden_size;
    }
    __syncthreads();

    float mean = s_mean;

    // Pass 2: compute variance = sum((x - mean)^2) / N
    float local_var = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float d = x_row[i] - mean;
        local_var += d * d;
    }
    local_var = block_reduce_sum(local_var);

    if (threadIdx.x == 0) {
        s_inv_std = rsqrtf(local_var / hidden_size + eps);
    }
    __syncthreads();

    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;

    // Pass 1: compute mean
    float local_sum = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        local_sum += __bfloat162float(x_row[i]);
    }
    local_sum = block_reduce_sum(local_sum);

    __shared__ float s_mean, s_inv_std;
    if (threadIdx.x == 0) {
        s_mean = local_sum / hidden_size;
    }
    __syncthreads();

    float mean = s_mean;

    // Pass 2: compute variance = sum((x - mean)^2) / N
    float local_var = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float d = __bfloat162float(x_row[i]) - mean;
        local_var += d * d;
    }
    local_var = block_reduce_sum(local_var);

    if (threadIdx.x == 0) {
        s_inv_std = rsqrtf(local_var / hidden_size + eps);
    }
    __syncthreads();

    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;
    if (block < 32) block = 32;
    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;
    if (block < 32) block = 32;
    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);
}

}