// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Elementwise kernels. Flat 1-D dispatch via dispatchThreads (non-uniform
// threadgroups — no bounds-check tail), threadgroup size a multiple of 32.
// Purely bandwidth-bound; float4-vectorized variants come with the M5
// fusion pass.
#include <metal_stdlib>
using namespace metal;

kernel void add_f32(device const float* a   [[buffer(0)]],
                    device const float* b   [[buffer(1)]],
                    device float*       out [[buffer(2)]],
                    uint gid [[thread_position_in_grid]]) {
    out[gid] = a[gid] + b[gid];
}

kernel void mul_f32(device const float* a   [[buffer(0)]],
                    device const float* b   [[buffer(1)]],
                    device float*       out [[buffer(2)]],
                    uint gid [[thread_position_in_grid]]) {
    out[gid] = a[gid] * b[gid];
}

kernel void scale_f32(device const float* a   [[buffer(0)]],
                      device float*       out [[buffer(1)]],
                      constant float&     s   [[buffer(2)]],
                      uint gid [[thread_position_in_grid]]) {
    out[gid] = a[gid] * s;
}

// out[i] = x[i] + bias[i % C] — row-broadcast bias
kernel void add_bias_f32(device const float* x    [[buffer(0)]],
                         device const float* bias [[buffer(1)]],
                         device float*       out  [[buffer(2)]],
                         constant uint&      C    [[buffer(3)]],
                         uint gid [[thread_position_in_grid]]) {
    out[gid] = x[gid] + bias[gid % C];
}

kernel void silu_f32(device const float* x   [[buffer(0)]],
                     device float*       out [[buffer(1)]],
                     uint gid [[thread_position_in_grid]]) {
    const float v = x[gid];
    out[gid] = v / (1.0f + exp(-v));
}

kernel void gelu_f32(device const float* x   [[buffer(0)]],
                     device float*       out [[buffer(1)]],
                     uint gid [[thread_position_in_grid]]) {
    const float v = x[gid];
    const float k = 0.7978845608028654f; // sqrt(2/pi)
    out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));
}

kernel void sigmoid_f32(device const float* x   [[buffer(0)]],
                        device float*       out [[buffer(1)]],
                        uint gid [[thread_position_in_grid]]) {
    out[gid] = 1.0f / (1.0f + exp(-x[gid]));
}

// dx += dout * y * (1 - y), using the forward output y.
kernel void sigmoid_bwd_f32(device const float* y    [[buffer(0)]],
                            device const float* dout [[buffer(1)]],
                            device float*       dx   [[buffer(2)]],
                            uint gid [[thread_position_in_grid]]) {
    const float v = y[gid];
    dx[gid] = fma(dout[gid], v * (1.0f - v), dx[gid]);
}

// ReLU² (nanoGPT-speedrun lineage): out = max(x, 0)^2.
kernel void relu2_f32(device const float* x   [[buffer(0)]],
                      device float*       out [[buffer(1)]],
                      uint gid [[thread_position_in_grid]]) {
    const float v = max(x[gid], 0.0f);
    out[gid] = v * v;
}

kernel void relu2_bwd_f32(device const float* x    [[buffer(0)]],
                          device const float* dout [[buffer(1)]],
                          device float*       dx   [[buffer(2)]],
                          uint gid [[thread_position_in_grid]]) {
    dx[gid] = fma(dout[gid], 2.0f * max(x[gid], 0.0f), dx[gid]);
}

// Gemma-style logit softcap: out = cap * tanh(x / cap).
kernel void softcap_f32(device const float* x   [[buffer(0)]],
                        device float*       out [[buffer(1)]],
                        constant float&     cap [[buffer(2)]],
                        uint gid [[thread_position_in_grid]]) {
    out[gid] = cap * precise::tanh(x[gid] / cap);
}

// dx += dout * (1 - (y/cap)^2), using the forward output y.
kernel void softcap_bwd_f32(device const float* y    [[buffer(0)]],
                            device const float* dout [[buffer(1)]],
                            device float*       dx   [[buffer(2)]],
                            constant float&     cap  [[buffer(3)]],
                            uint gid [[thread_position_in_grid]]) {
    const float t = y[gid] / cap;
    dx[gid] = fma(dout[gid], 1.0f - t * t, dx[gid]);
}

// ---- backward / accumulation kernels (all ACCUMULATE into their outputs) ----

kernel void accum_f32(device float*       dst [[buffer(0)]],
                      device const float* src [[buffer(1)]],
                      uint gid [[thread_position_in_grid]]) {
    dst[gid] += src[gid];
}

// dst += src * s ; s in a 1-element buffer so it can be produced on-GPU
// (e.g. d(loss) scaling) without a sync.
kernel void axpy_f32(device float*       dst [[buffer(0)]],
                     device const float* src [[buffer(1)]],
                     device const float* s   [[buffer(2)]],
                     uint gid [[thread_position_in_grid]]) {
    dst[gid] = fma(src[gid], s[0], dst[gid]);
}

kernel void silu_bwd_f32(device const float* x    [[buffer(0)]],
                         device const float* dout [[buffer(1)]],
                         device float*       dx   [[buffer(2)]],
                         uint gid [[thread_position_in_grid]]) {
    const float v = x[gid];
    const float sig = 1.0f / (1.0f + exp(-v));
    dx[gid] = fma(dout[gid], sig * (1.0f + v * (1.0f - sig)), dx[gid]);
}

kernel void gelu_bwd_f32(device const float* x    [[buffer(0)]],
                         device const float* dout [[buffer(1)]],
                         device float*       dx   [[buffer(2)]],
                         uint gid [[thread_position_in_grid]]) {
    const float v = x[gid];
    const float k = 0.7978845608028654f;
    const float u = k * (v + 0.044715f * v * v * v);
    const float t = precise::tanh(u);
    const float du = k * (1.0f + 3.0f * 0.044715f * v * v);
    dx[gid] = fma(dout[gid], 0.5f * (1.0f + t) + 0.5f * v * (1.0f - t * t) * du, dx[gid]);
}

// dbias[j] += sum_rows dout[i,j] — one thread per column, strided rows.
// Column-major walk is uncoalesced but this kernel is a tiny fraction of a
// step; revisit in the M5 fusion pass if it ever shows in a trace.
kernel void add_bias_bwd_f32(device const float* dout  [[buffer(0)]],
                             device float*       dbias [[buffer(1)]],
                             constant uint2&     nc    [[buffer(2)]], // (N, C)
                             uint j [[thread_position_in_grid]]) {
    float acc = 0.0f;
    for (uint i = 0; i < nc.x; ++i) acc += dout[ulong(i) * nc.y + j];
    dbias[j] += acc;
}

// ---- RoPE (interleaved pairs; INV=true applies the inverse rotation and
// accumulates — the backward pass) ----

constant bool ROPE_INV [[function_constant(2)]];

struct RopeParams {
    uint T, H, HD;
    uint pos_offset;
};

// freqs[k] is the per-pair inverse frequency (theta^(-2k/HD), possibly
// rescaled — llama3 rope_scaling, per-layer theta). Host-precomputed.
kernel void rope_f32(device const float* x     [[buffer(0)]],
                     device float*       out   [[buffer(1)]],
                     constant RopeParams& p    [[buffer(2)]],
                     device const float* freqs [[buffer(3)]],
                     uint gid [[thread_position_in_grid]]) {
    // gid indexes (bt, h, k) pairs: one thread per rotated pair
    const uint pairs_per_row = p.H * (p.HD / 2);
    const uint bt = gid / pairs_per_row;
    const uint rem = gid % pairs_per_row;
    const uint h = rem / (p.HD / 2);
    const uint k = rem % (p.HD / 2);

    const float pos = float(bt % p.T + p.pos_offset);
    const float angle = pos * freqs[k];
    const float c = cos(angle);
    const float s = ROPE_INV ? -sin(angle) : sin(angle);

    const ulong i0 = ulong(bt) * (p.H * p.HD) + h * p.HD + 2 * k;
    const float x0 = x[i0], x1 = x[i0 + 1];
    if (ROPE_INV) {
        out[i0]     += x0 * c - x1 * s;
        out[i0 + 1] += x0 * s + x1 * c;
    } else {
        out[i0]     = x0 * c - x1 * s;
        out[i0 + 1] = x0 * s + x1 * c;
    }
}