// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Fused softmax cross-entropy (llm.c fused_classifier pattern): one
// threadgroup per row computes online (max, sumexp), the loss in
// logsumexp form, and — in the same kernel — the logit gradient
// (softmax − onehot) · inv_n, accumulated into dlogits. The full softmax
// is never materialized. targets use ignore_index = -1 (loss 0, grad 0).
//
// losses[row] receives the per-row loss; the host reduces (sum_f32) and
// scales by 1/n_valid.
#include <metal_stdlib>
using namespace metal;

struct CEParams {
    uint V;
    float inv_n;    // 1 / n_valid
    uint want_grad; // 0: loss only
};

kernel void cross_entropy_f32(device const float* logits  [[buffer(0)]],
                              device const int*   targets [[buffer(1)]],
                              device float*       losses  [[buffer(2)]],
                              device float*       dlogits [[buffer(3)]],
                              constant CEParams&  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]]) {
    const int tgt = targets[row];
    device const float* x = logits + ulong(row) * p.V;

    if (tgt < 0) {
        if (lid == 0) losses[row] = 0.0f;
        return;
    }

    // online (m, l) as in softmax.metal
    float m = -FLT_MAX;
    float l = 0.0f;
    for (uint j = lid; j < p.V; 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;
    }
    float m_simd = simd_max(m);
    float l_simd = simd_sum(l * exp(m - m_simd));

    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;
    {
        const uint i = min(lane, n_simds - 1);
        const float mi = tg_m[i];
        const float li = tg_l[i];
        m_row = simd_max(mi);
        const float contrib = (lane < n_simds) ? li * exp(mi - m_row) : 0.0f;
        l_row = simd_sum(contrib);
    }

    if (lid == 0) {
        // loss = logsumexp − logit[target]
        losses[row] = m_row + log(l_row) - x[uint(tgt)];
    }

    if (p.want_grad) {
        const float inv_sum = 1.0f / l_row;
        device float* drow = dlogits + ulong(row) * p.V;
        for (uint j = lid; j < p.V; j += tg_size) {
            const float prob = exp(x[j] - m_row) * inv_sum;
            const float ind = (j == uint(tgt)) ? 1.0f : 0.0f;
            drow[j] += (prob - ind) * p.inv_n;
        }
    }
}