// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Fused AdamW update (llm.c formulation): one thread per element; the host
// precomputes the bias corrections and folds the gradient-clip scale into
// grad_scale, so the kernel reads the raw gradient once and writes w/m/v.
// eps sits OUTSIDE the sqrt; weight decay is decoupled (never through m/v)
// and the host sets wd = 0 for dim<2 tensors.
//
// sumsq_f32: grid-stride sum of squares of one tensor into partials[tg] —
// used for the global grad-norm (host sums the partials at the sync point).
#include <metal_stdlib>
using namespace metal;

struct AdamWParams {
    float lr, beta1, beta2, bc1, bc2, eps, wd, grad_scale;
};

kernel void adamw_f32(device float*       w [[buffer(0)]],
                      device const float* g [[buffer(1)]],
                      device float*       m [[buffer(2)]],
                      device float*       v [[buffer(3)]],
                      constant AdamWParams& p [[buffer(4)]],
                      uint gid [[thread_position_in_grid]]) {
    const float grad = g[gid] * p.grad_scale;
    const float mi = p.beta1 * m[gid] + (1.0f - p.beta1) * grad;
    const float vi = p.beta2 * v[gid] + (1.0f - p.beta2) * grad * grad;
    m[gid] = mi;
    v[gid] = vi;
    const float mhat = mi / p.bc1;
    const float vhat = vi / p.bc2;
    w[gid] -= p.lr * (mhat / (sqrt(vhat) + p.eps) + p.wd * w[gid]);
}

kernel void sumsq_f32(device const float* x        [[buffer(0)]],
                      device float*       partials [[buffer(1)]],
                      constant uint&      n        [[buffer(2)]],
                      uint gid      [[thread_position_in_grid]],
                      uint grid_sz  [[threads_per_grid]],
                      uint tg_id    [[threadgroup_position_in_grid]],
                      uint lane     [[thread_index_in_simdgroup]],
                      uint simd_idx [[simdgroup_index_in_threadgroup]],
                      uint n_simds  [[simdgroups_per_threadgroup]]) {
    float acc = 0.0f;
    for (uint i = gid; i < n; i += grid_sz) acc = fma(x[i], x[i], acc);

    threadgroup float scratch[32];
    const float s = simd_sum(acc);
    if (simd_is_first()) scratch[simd_idx] = s;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;
    const float total = simd_sum(mine);
    if (simd_idx == 0 && lane == 0) partials[tg_id] = total;
}

// Single-threadgroup final sum (also reduces CE per-row losses).
kernel void sum_f32(device const float* x   [[buffer(0)]],
                    device float*       out [[buffer(1)]],
                    constant uint&      n   [[buffer(2)]],
                    constant float&     mul [[buffer(3)]],
                    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]]) {
    float acc = 0.0f;
    for (uint i = lid; i < n; i += tg_size) acc += x[i];

    threadgroup float scratch[32];
    const float s = simd_sum(acc);
    if (simd_is_first()) scratch[simd_idx] = s;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;
    const float total = simd_sum(mine);
    if (simd_idx == 0 && lane == 0) out[0] = total * mul;
}