// projection.cu — h_128 -> z_8 with affine + layer-norm.
//
// pre[j]   = sum_i w[j, i] * h[i] + b[j]                  for j in 0..8
// mean     = (1/8) * sum_j pre[j]
// var      = (1/8) * sum_j (pre[j] - mean)^2
// out[j]   = ln_gain[j] * (pre[j] - mean) * rsqrt(var + 1e-6) + ln_bias[j]
//
// Single block, 8 threads (one per output unit). Mean and variance
// computed by a single thread (block-coordinated) to keep the kernel
// simple — the cost is negligible for 8 outputs. No atomicAdd anywhere.

extern "C" __global__ void projection_kernel(
    const float* __restrict__ w,        // [8, 128]
    const float* __restrict__ b,        // [8]
    const float* __restrict__ ln_gain,  // [8]
    const float* __restrict__ ln_bias,  // [8]
    const float* __restrict__ h,        // [128]
    float* __restrict__ out             // [8]
) {
    __shared__ float pre[8];
    __shared__ float mean_var[2]; // [mean, var]

    int j = threadIdx.x;
    if (j < 8) {
        float v = b[j];
        for (int i = 0; i < 128; ++i) {
            v += w[j * 128 + i] * h[i];
        }
        pre[j] = v;
    }
    __syncthreads();

    if (j == 0) {
        float s = 0.0f, ss = 0.0f;
        for (int k = 0; k < 8; ++k) {
            s += pre[k];
            ss += pre[k] * pre[k];
        }
        const float mean = s * 0.125f; // /8
        const float var = ss * 0.125f - mean * mean;
        mean_var[0] = mean;
        mean_var[1] = fmaxf(var, 1e-6f);
    }
    __syncthreads();

    if (j < 8) {
        const float mean = mean_var[0];
        const float invstd = rsqrtf(mean_var[1]);
        out[j] = ln_gain[j] * (pre[j] - mean) * invstd + ln_bias[j];
    }
}
