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

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

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

    // Thread-local accumulation (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();

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

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