// LayerNorm forward kernel — f32-native.
// Common header prepended by build.rs.

extern "C" __global__
void layer_norm_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;

    extern __shared__ float smem[];
    float* s_sum = smem;
    float* s_sq  = smem + blockDim.x;

    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();

    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 n = ((float)features);
    float mean = s_sum[0] / n;
    float var  = s_sq[0] / n - mean * mean;
    float inv_std = 1.0f / sqrtf(var + eps);

    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = x_row[j];
        float normalized = (val - mean) * inv_std;
        y_row[j] = normalized * gamma[j] + beta[j];
    }
}
