#include "../common.cuh"

// RMSNorm: y[i] = x[i] * rsqrt(mean(x²) + eps) * gamma[i]
// Each block processes one row of shape [hidden_size].

__global__ void rmsnorm_f32(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    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;

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

    __shared__ float s_rms_inv;
    if (threadIdx.x == 0) {
        s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
    }
    __syncthreads();

    float rms_inv = s_rms_inv;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        out_row[i] = x_row[i] * rms_inv * gamma[i];
    }
}

__global__ void rmsnorm_bf16(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ gamma,
    __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 sum_sq = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float v = __bfloat162float(x_row[i]);
        sum_sq += v * v;
    }
    sum_sq = block_reduce_sum(sum_sq);

    __shared__ float s_rms_inv;
    if (threadIdx.x == 0) {
        s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
    }
    __syncthreads();

    float rms_inv = s_rms_inv;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float v = __bfloat162float(x_row[i]);
        float g = __bfloat162float(gamma[i]);
        out_row[i] = __float2bfloat16(v * rms_inv * g);
    }
}

extern "C" {

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

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

}