// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Mixture-of-experts kernels. All operate on router-sized rows (E experts,
// E <= 64 in practice), so every kernel is one thread per row with a plain
// loop — no threadgroup reductions needed. row_scale kernels are flat
// elementwise over [N, C] activations.
#include <metal_stdlib>
using namespace metal;

// Generic row softmax backward: dx += p ∘ (dout − dot(dout, p)), row = last
// dim. Thread-per-row; meant for small C (the router's E columns).
kernel void softmax_bwd_f32(device const float* P    [[buffer(0)]],
                            device const float* DOUT [[buffer(1)]],
                            device float*       DX   [[buffer(2)]],
                            constant uint&      C    [[buffer(3)]],
                            uint row [[thread_position_in_grid]]) {
    device const float* prob = P + ulong(row) * C;
    device const float* dout = DOUT + ulong(row) * C;
    device float* dx = DX + ulong(row) * C;
    float dot = 0.0f;
    for (uint j = 0; j < C; ++j) dot += dout[j] * prob[j];
    for (uint j = 0; j < C; ++j) dx[j] += prob[j] * (dout[j] - dot);
}

// Top-k gating with a selection bias (DeepSeek-V3 "noaux"): the k experts
// are chosen by score+BIAS, but gate VALUES come from the biasless score —
// the bias only steers routing. norm=1 renormalizes kept gates to sum 1
// (classic top-k softmax); norm=0 keeps raw scores (sigmoid routing, V3).
// Ties broken by lower index, which keeps forward and backward selections
// identical. p = (E, k, norm)
struct TopkParams { uint E, K, norm; };

static inline void topk_select_biased(device const float* in,
                                      device const float* bias, uint E, uint K,
                                      thread bool* kept, thread float* S) {
    for (uint j = 0; j < E; ++j) kept[j] = false;
    float acc = 0.0f;
    for (uint sel = 0; sel < K; ++sel) {
        float best = -FLT_MAX;
        uint arg = 0;
        for (uint j = 0; j < E; ++j) {
            const float v = in[j] + bias[j];
            if (!kept[j] && v > best) { best = v; arg = j; }
        }
        kept[arg] = true;
        acc += in[arg]; // gate mass is biasless
    }
    *S = acc;
}

kernel void topk_renorm_f32(device const float* P    [[buffer(0)]],
                            device const float* BIAS [[buffer(1)]],
                            device float*       OUT  [[buffer(2)]],
                            constant TopkParams& p   [[buffer(3)]],
                            uint row [[thread_position_in_grid]]) {
    device const float* in = P + ulong(row) * p.E;
    device float* out = OUT + ulong(row) * p.E;

    bool kept[64];
    float S;
    topk_select_biased(in, BIAS, p.E, p.K, kept, &S);
    const float inv = p.norm ? 1.0f / max(S, 1e-12f) : 1.0f;
    for (uint j = 0; j < p.E; ++j) out[j] = kept[j] ? in[j] * inv : 0.0f;
}

// Backward: with norm, g_i = p_i / S over kept entries so
// dp_i += (dg_i − Σ_j dg_j g_j) / S; without norm, g_i = p_i so dp_i += dg_i.
// The kept set (incl. bias) is recomputed with the same tie-breaking.
kernel void topk_renorm_bwd_f32(device const float* P    [[buffer(0)]],
                                device const float* BIAS [[buffer(1)]],
                                device const float* DOUT [[buffer(2)]],
                                device float*       DP   [[buffer(3)]],
                                constant TopkParams& p   [[buffer(4)]],
                                uint row [[thread_position_in_grid]]) {
    device const float* in = P + ulong(row) * p.E;
    device const float* dout = DOUT + ulong(row) * p.E;
    device float* dp = DP + ulong(row) * p.E;

    bool kept[64];
    float S;
    topk_select_biased(in, BIAS, p.E, p.K, kept, &S);
    if (!p.norm) {
        for (uint j = 0; j < p.E; ++j)
            if (kept[j]) dp[j] += dout[j];
        return;
    }
    const float inv = 1.0f / max(S, 1e-12f);
    float dot = 0.0f;
    for (uint j = 0; j < p.E; ++j)
        if (kept[j]) dot += dout[j] * in[j] * inv;
    for (uint j = 0; j < p.E; ++j)
        if (kept[j]) dp[j] += (dout[j] - dot) * inv;
}

// counts[e] += number of rows whose gate for expert e is nonzero — the load
// statistic driving the noaux bias update. One thread per expert.
kernel void expert_counts_f32(device const float* G      [[buffer(0)]],
                              device float*       COUNTS [[buffer(1)]],
                              constant uint2&     p      [[buffer(2)]], // (E, N)
                              uint e [[thread_position_in_grid]]) {
    float acc = 0.0f;
    for (uint i = 0; i < p.y; ++i)
        if (G[ulong(i) * p.x + e] != 0.0f) acc += 1.0f;
    COUNTS[e] += acc;
}

// 12-byte layout matching the host-side struct (uint3 would pad to 16).
struct RowScaleParams { uint C, E, e; };

// out[i, c] = x[i, c] * G[i, e] — scale each row of X by one gate column.
// Flat over N*C.
kernel void row_scale_f32(device const float*       X   [[buffer(0)]],
                          device const float*       G   [[buffer(1)]],
                          device float*             OUT [[buffer(2)]],
                          constant RowScaleParams&  p   [[buffer(3)]],
                          uint gid [[thread_position_in_grid]]) {
    const uint i = gid / p.C;
    OUT[gid] = X[gid] * G[ulong(i) * p.E + p.e];
}

// dst[i, c] += x[i, c] * G[i, e] — accumulating variant (dx of row_scale).
kernel void row_scale_acc_f32(device const float*       X   [[buffer(0)]],
                              device const float*       G   [[buffer(1)]],
                              device float*             DST [[buffer(2)]],
                              constant RowScaleParams&  p   [[buffer(3)]],
                              uint gid [[thread_position_in_grid]]) {
    const uint i = gid / p.C;
    DST[gid] = fma(X[gid], G[ulong(i) * p.E + p.e], DST[gid]);
}

// dG[i, e] += dot(dout[i, :], x[i, :]) — gate gradient, one thread per row.
kernel void row_scale_gate_bwd_f32(device const float*       DOUT [[buffer(0)]],
                                   device const float*       X    [[buffer(1)]],
                                   device float*             DG   [[buffer(2)]],
                                   constant RowScaleParams&  p    [[buffer(3)]],
                                   uint row [[thread_position_in_grid]]) {
    device const float* dout = DOUT + ulong(row) * p.C;
    device const float* x = X + ulong(row) * p.C;
    float acc = 0.0f;
    for (uint j = 0; j < p.C; ++j) acc += dout[j] * x[j];
    DG[ulong(row) * p.E + p.e] += acc;
}