// RMSNorm + LayerNorm forward kernels -- BF16-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), 0.0f, 1.0f, etc.

// ── RMSNorm forward ────────────────────────────────────────────────────
// Each block processes one sample (row) of the [batch, features] input.
// Computes: y = x / sqrt(mean(x^2) + eps) * weight
//
// Two-pass approach:
//   Pass 1: compute mean(x^2) via parallel reduction in shared memory (f32 accum).
//   Pass 2: normalize and scale each element.

extern "C" __global__
void rmsnorm_forward(const float* __restrict__ x,
                     float* __restrict__ y,
                     const float* __restrict__ weight,
                     int batch,
                     int features,
                     float eps) {
    int row = blockIdx.x;
    if (row >= batch) return;

    const float* x_row = x + row * features;
    float* y_row = y + row * features;

    extern __shared__ float smem[];

    // Thread-local accumulation of x^2 over features (f32 for precision)
    float local_sq = 0.0f;
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = (x_row[j]);
        local_sq += val * val;
    }

    smem[threadIdx.x] = local_sq;
    __syncthreads();

    // Parallel reduction for sum of squares
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            smem[threadIdx.x] += smem[threadIdx.x + stride];
        }
        __syncthreads();
    }

    // rms = sqrt(mean(x^2) + eps)
    float inv_rms = rsqrtf(smem[0] / (float)features + eps);

    // Normalize and scale
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float xf = (x_row[j]);
        float wf = (weight[j]);
        y_row[j] = (xf * inv_rms * wf);
    }
}

// ── LayerNorm forward ──────────────────────────────────────────────────
// Each block processes one sample (row) of the [batch, features] input.
// Two-pass approach:
//   Pass 1: compute mean and variance via parallel reduction in shared memory.
//   Pass 2: normalize, scale, and shift each element.

extern "C" __global__
void layernorm_forward(const float* __restrict__ x,
                       float* __restrict__ y,
                       const float* __restrict__ gamma,
                       const float* __restrict__ beta,
                       int batch,
                       int features,
                       float eps) {
    int row = blockIdx.x;
    if (row >= batch) return;

    const float* x_row = x + row * features;
    float* y_row = y + row * features;

    // Shared memory: first half for sum, second half for sum of squares
    extern __shared__ float smem[];
    float* s_sum = smem;
    float* s_sq  = smem + blockDim.x;

    // Thread-local accumulation over features (f32 for precision)
    float local_sum = 0.0f;
    float local_sq  = 0.0f;
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = (x_row[j]);
        local_sum += val;
        local_sq  += val * val;
    }

    s_sum[threadIdx.x] = local_sum;
    s_sq[threadIdx.x]  = local_sq;
    __syncthreads();

    // Parallel reduction
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
            s_sq[threadIdx.x]  += s_sq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    float mean = s_sum[0] / (float)features;
    float var  = s_sq[0] / (float)features - mean * mean;
    float inv_std = rsqrtf(var + eps);

    // Normalize, scale, shift
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float xf = (x_row[j]);
        float gf = (gamma[j]);
        float bf = (beta[j]);
        float normalized = (xf - mean) * inv_std;
        y_row[j] = (normalized * gf + bf);
    }
}
