// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Row-wise RMSNorm / LayerNorm. One threadgroup per row; simd reductions
// first (register shuffles beat threadgroup memory 2:1 on Apple GPUs),
// tiny threadgroup scratch only to cross simdgroups. LayerNorm is two-pass
// (mean, then variance) — the E[x²]−E[x]² shortcut loses precision and
// these kernels are held to 1e-4 parity vs the two-pass CPU reference.
#include <metal_stdlib>
using namespace metal;

namespace {

// Threadgroup-wide sum of one per-thread value; returns the same total to
// every thread. Scratch must hold 32 floats.
inline float tg_sum(float v, uint lane, uint simd_idx, uint n_simds,
                    threadgroup float* scratch) {
    const float s = simd_sum(v);
    if (simd_is_first()) scratch[simd_idx] = s;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;
    return simd_sum(mine);
}

} // namespace

struct NormParams {
    uint C;
    float eps;
};

kernel void rmsnorm_f32(device const float* X   [[buffer(0)]],
                        device const float* W   [[buffer(1)]],
                        device float*       OUT [[buffer(2)]],
                        constant NormParams& p  [[buffer(3)]],
                        uint row      [[threadgroup_position_in_grid]],
                        uint lid      [[thread_index_in_threadgroup]],
                        uint tg_size  [[threads_per_threadgroup]],
                        uint lane     [[thread_index_in_simdgroup]],
                        uint simd_idx [[simdgroup_index_in_threadgroup]],
                        uint n_simds  [[simdgroups_per_threadgroup]]) {
    device const float* x = X + ulong(row) * p.C;
    device float* out = OUT + ulong(row) * p.C;
    threadgroup float scratch[32];

    float ss = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) ss = fma(x[j], x[j], ss);
    ss = tg_sum(ss, lane, simd_idx, n_simds, scratch);

    const float inv_rms = rsqrt(ss / float(p.C) + p.eps);
    for (uint j = lid; j < p.C; j += tg_size) out[j] = W[j] * x[j] * inv_rms;
}

kernel void layernorm_f32(device const float* X   [[buffer(0)]],
                          device const float* W   [[buffer(1)]],
                          device const float* B   [[buffer(2)]],
                          device float*       OUT [[buffer(3)]],
                          constant NormParams& p  [[buffer(4)]],
                          uint row      [[threadgroup_position_in_grid]],
                          uint lid      [[thread_index_in_threadgroup]],
                          uint tg_size  [[threads_per_threadgroup]],
                          uint lane     [[thread_index_in_simdgroup]],
                          uint simd_idx [[simdgroup_index_in_threadgroup]],
                          uint n_simds  [[simdgroups_per_threadgroup]]) {
    device const float* x = X + ulong(row) * p.C;
    device float* out = OUT + ulong(row) * p.C;
    threadgroup float scratch[32];

    float s = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) s += x[j];
    const float mean = tg_sum(s, lane, simd_idx, n_simds, scratch) / float(p.C);
    threadgroup_barrier(mem_flags::mem_threadgroup); // scratch reuse

    float var = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) {
        const float d = x[j] - mean;
        var = fma(d, d, var);
    }
    var = tg_sum(var, lane, simd_idx, n_simds, scratch) / float(p.C);

    const float inv_std = rsqrt(var + p.eps);
    for (uint j = lid; j < p.C; j += tg_size)
        out[j] = fma(W[j], (x[j] - mean) * inv_std, B[j]);
}

// ---- backwards -------------------------------------------------------------
// dx is row-parallel (one threadgroup per row) and saves the row statistics
// (inv_rms, or mean+inv_std) it computes anyway; dw/db are column-parallel
// (one thread per column) reading those saved stats — deterministic, no
// atomics, no per-column recomputation.

kernel void rmsnorm_bwd_dx_f32(device const float* X       [[buffer(0)]],
                               device const float* W       [[buffer(1)]],
                               device const float* dOUT    [[buffer(2)]],
                               device float*       dX      [[buffer(3)]],
                               device float*       inv_rms [[buffer(4)]], // [rows]
                               constant NormParams& p      [[buffer(5)]],
                               uint row      [[threadgroup_position_in_grid]],
                               uint lid      [[thread_index_in_threadgroup]],
                               uint tg_size  [[threads_per_threadgroup]],
                               uint lane     [[thread_index_in_simdgroup]],
                               uint simd_idx [[simdgroup_index_in_threadgroup]],
                               uint n_simds  [[simdgroups_per_threadgroup]]) {
    device const float* x = X + ulong(row) * p.C;
    device const float* d = dOUT + ulong(row) * p.C;
    device float* dx = dX + ulong(row) * p.C;
    threadgroup float scratch[32];

    float ss = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) ss = fma(x[j], x[j], ss);
    ss = tg_sum(ss, lane, simd_idx, n_simds, scratch);
    const float ir = rsqrt(ss / float(p.C) + p.eps);
    if (lid == 0) inv_rms[row] = ir;
    threadgroup_barrier(mem_flags::mem_threadgroup);

    float dot = 0.0f; // sum_j g_j w_j x_j
    for (uint j = lid; j < p.C; j += tg_size) dot = fma(d[j] * W[j], x[j], dot);
    dot = tg_sum(dot, lane, simd_idx, n_simds, scratch);

    const float coef = dot * ir * ir * ir / float(p.C);
    for (uint j = lid; j < p.C; j += tg_size)
        dx[j] += d[j] * W[j] * ir - x[j] * coef;
}

// dw[j] += sum_i g[i,j] * x[i,j] * inv_rms[i]
kernel void rmsnorm_bwd_dw_f32(device const float* X       [[buffer(0)]],
                               device const float* dOUT    [[buffer(1)]],
                               device const float* inv_rms [[buffer(2)]],
                               device float*       dW      [[buffer(3)]],
                               constant NormParams& p      [[buffer(4)]],
                               constant uint&      rows    [[buffer(5)]],
                               uint j [[thread_position_in_grid]]) {
    float acc = 0.0f;
    for (uint i = 0; i < rows; ++i)
        acc = fma(dOUT[ulong(i) * p.C + j] * X[ulong(i) * p.C + j], inv_rms[i], acc);
    dW[j] += acc;
}

kernel void layernorm_bwd_dx_f32(device const float* X    [[buffer(0)]],
                                 device const float* W    [[buffer(1)]],
                                 device const float* dOUT [[buffer(2)]],
                                 device float*       dX   [[buffer(3)]],
                                 device float*       mean_out [[buffer(4)]], // [rows]
                                 device float*       istd_out [[buffer(5)]], // [rows]
                                 constant NormParams& p   [[buffer(6)]],
                                 uint row      [[threadgroup_position_in_grid]],
                                 uint lid      [[thread_index_in_threadgroup]],
                                 uint tg_size  [[threads_per_threadgroup]],
                                 uint lane     [[thread_index_in_simdgroup]],
                                 uint simd_idx [[simdgroup_index_in_threadgroup]],
                                 uint n_simds  [[simdgroups_per_threadgroup]]) {
    device const float* x = X + ulong(row) * p.C;
    device const float* d = dOUT + ulong(row) * p.C;
    device float* dx = dX + ulong(row) * p.C;
    threadgroup float scratch[32];

    float s = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) s += x[j];
    const float mean = tg_sum(s, lane, simd_idx, n_simds, scratch) / float(p.C);
    threadgroup_barrier(mem_flags::mem_threadgroup);

    float var = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) {
        const float dv = x[j] - mean;
        var = fma(dv, dv, var);
    }
    var = tg_sum(var, lane, simd_idx, n_simds, scratch) / float(p.C);
    const float inv_std = rsqrt(var + p.eps);
    if (lid == 0) {
        mean_out[row] = mean;
        istd_out[row] = inv_std;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    float m_dxhat = 0.0f, m_dxhat_xhat = 0.0f;
    for (uint j = lid; j < p.C; j += tg_size) {
        const float xhat = (x[j] - mean) * inv_std;
        const float dxhat = d[j] * W[j];
        m_dxhat += dxhat;
        m_dxhat_xhat = fma(dxhat, xhat, m_dxhat_xhat);
    }
    m_dxhat = tg_sum(m_dxhat, lane, simd_idx, n_simds, scratch) / float(p.C);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    m_dxhat_xhat = tg_sum(m_dxhat_xhat, lane, simd_idx, n_simds, scratch) / float(p.C);

    for (uint j = lid; j < p.C; j += tg_size) {
        const float xhat = (x[j] - mean) * inv_std;
        dx[j] += inv_std * (d[j] * W[j] - m_dxhat - xhat * m_dxhat_xhat);
    }
}

// dw[j] += sum_i g[i,j] * xhat[i,j] ; db[j] += sum_i g[i,j]
kernel void layernorm_bwd_dwdb_f32(device const float* X    [[buffer(0)]],
                                   device const float* dOUT [[buffer(1)]],
                                   device const float* mean [[buffer(2)]],
                                   device const float* istd [[buffer(3)]],
                                   device float*       dW   [[buffer(4)]],
                                   device float*       dB   [[buffer(5)]],
                                   constant NormParams& p   [[buffer(6)]],
                                   constant uint&      rows [[buffer(7)]],
                                   uint j [[thread_position_in_grid]]) {
    float aw = 0.0f, ab = 0.0f;
    for (uint i = 0; i < rows; ++i) {
        const float g = dOUT[ulong(i) * p.C + j];
        const float xhat = (X[ulong(i) * p.C + j] - mean[i]) * istd[i];
        aw = fma(g, xhat, aw);
        ab += g;
    }
    dW[j] += aw;
    dB[j] += ab;
}