// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Fused online softmax over the last dim. One threadgroup per row; each
// thread keeps a running (max, sum) over a strided slice of the row, lanes
// combine with simd reductions, simdgroups combine through a small
// threadgroup scratch (<= 32 simdgroups, so one simd-reduce finishes it).
// Row layout: X[row * C + j].
//
// Masked/padded values are the caller's problem; -FLT_MAX sentinels are safe
// here (never -INFINITY: fast-math exp(-inf) is undefined — RESEARCH.md §8).
#include <metal_stdlib>
using namespace metal;

kernel void softmax_f32(device const float* X   [[buffer(0)]],
                        device float*       OUT [[buffer(1)]],
                        constant uint&      C   [[buffer(2)]],
                        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) * C;
    device float* out = OUT + ulong(row) * C;

    // per-thread online (m, l)
    float m = -FLT_MAX;
    float l = 0.0f;
    for (uint j = lid; j < C; j += tg_size) {
        const float v = x[j];
        const float m_new = max(m, v);
        l = l * exp(m - m_new) + exp(v - m_new);
        m = m_new;
    }

    // combine across the simdgroup
    float m_simd = simd_max(m);
    float l_simd = simd_sum(l * exp(m - m_simd));

    // combine across simdgroups
    threadgroup float tg_m[32];
    threadgroup float tg_l[32];
    if (simd_is_first()) {
        tg_m[simd_idx] = m_simd;
        tg_l[simd_idx] = l_simd;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    float m_row, l_row;
    {
        // Indexed by lane-within-simdgroup so EVERY simdgroup runs the same
        // reduction over the same scratch and lands on identical row stats.
        const uint i = min(lane, n_simds - 1); // lanes >= n_simds mirror the last entry
        const float mi = tg_m[i];
        const float li = tg_l[i];
        m_row = simd_max(mi);
        // lanes beyond n_simds would double-count: zero their contribution
        const float contrib = (lane < n_simds) ? li * exp(mi - m_row) : 0.0f;
        l_row = simd_sum(contrib);
    }

    const float inv_l = 1.0f / l_row;
    for (uint j = lid; j < C; j += tg_size) {
        out[j] = exp(x[j] - m_row) * inv_l;
    }
}