/* =====================================================================
 * rl_outcome_fused.cu -- Fused fwd + CE + bwd for K=3 outcome head.
 *
 * Eliminates two global-memory round-trips (logits, grad_logits) by
 * keeping all intermediates in shared memory.
 *
 * Grid  = (B, 1, 1)   -- one block per batch element.
 * Block = (128, 1, 1)  -- matches backward's HIDDEN_DIM parallelism.
 *
 * Three phases separated by __syncthreads():
 *
 *   Phase 1 (threads 0..2):  linear forward  logits = h_t × W + b
 *                             → s_logits[3] in shared memory
 *
 *   Phase 2 (thread 0):      softmax CE + gradient
 *                             → loss_pb[batch], s_grad_logits[3]
 *
 *   Phase 3 (all 128 threads): backward through linear layer
 *                             → grad_w_per_batch, grad_b_per_batch, grad_h_t
 *
 * No atomicAdd -- sole-writer per output element.
 * No nvrtc     -- precompiled cubin.
 * ===================================================================== */

#include <stdint.h>

#define HIDDEN_DIM 128
#define K_CLASSES  3

extern "C" __global__ void rl_outcome_fused(
    const float* __restrict__ h_t,                /* [B, HIDDEN_DIM=128]            */
    const float* __restrict__ w,                  /* [HIDDEN_DIM, K_CLASSES=3] col-major */
    const float* __restrict__ b,                  /* [K_CLASSES=3]                  */
    const int*   __restrict__ labels,             /* [B]  (0..2 or -1=skip)         */
    int                       b_size,
    float*       __restrict__ loss_pb,            /* [B]                            */
    float*       __restrict__ grad_w_per_batch,   /* [B, HIDDEN_DIM, K_CLASSES]     */
    float*       __restrict__ grad_b_per_batch,   /* [B, K_CLASSES]                 */
    float*       __restrict__ grad_h_t            /* [B, HIDDEN_DIM]                */
) {
    const int batch = blockIdx.x;
    const int tid   = threadIdx.x;

    if (batch >= b_size) return;

    /* Shared intermediates -- eliminates two global round-trips. */
    __shared__ float s_logits[K_CLASSES];
    __shared__ float s_grad_logits[K_CLASSES];

    const float* h_row = h_t + batch * HIDDEN_DIM;

    /* ── Phase 1: linear forward (threads 0..2) ────────────────────── */
    if (tid < K_CLASSES) {
        float sum = b[tid];
        #pragma unroll 16
        for (int i = 0; i < HIDDEN_DIM; ++i) {
            sum += h_row[i] * w[i * K_CLASSES + tid];
        }
        s_logits[tid] = sum;
    }
    __syncthreads();

    /* ── Phase 2: softmax CE + gradient (thread 0 only) ──────────── */
    if (tid == 0) {
        int label = labels[batch];

        if (label < 0 || label >= K_CLASSES) {
            /* Masked sample: no label available. */
            loss_pb[batch] = 0.0f;
            s_grad_logits[0] = 0.0f;
            s_grad_logits[1] = 0.0f;
            s_grad_logits[2] = 0.0f;
        } else {
            float l0 = s_logits[0];
            float l1 = s_logits[1];
            float l2 = s_logits[2];

            /* Numerically stable softmax: subtract max. */
            float mx = fmaxf(l0, fmaxf(l1, l2));
            float e0 = expf(l0 - mx);
            float e1 = expf(l1 - mx);
            float e2 = expf(l2 - mx);
            float sum_exp = e0 + e1 + e2;

            float p0 = e0 / sum_exp;
            float p1 = e1 / sum_exp;
            float p2 = e2 / sum_exp;

            /* CE loss = -log(p[label]). */
            float p_label = (label == 0) ? p0 : ((label == 1) ? p1 : p2);
            loss_pb[batch] = -logf(fmaxf(p_label, 1e-12f));

            /* Gradient: p[k] - 1(k == label). */
            s_grad_logits[0] = p0 - ((label == 0) ? 1.0f : 0.0f);
            s_grad_logits[1] = p1 - ((label == 1) ? 1.0f : 0.0f);
            s_grad_logits[2] = p2 - ((label == 2) ? 1.0f : 0.0f);
        }
    }
    __syncthreads();

    /* ── Phase 3: backward through linear layer (all 128 threads) ── */
    if (tid >= HIDDEN_DIM) return;

    const float h_bi = h_row[tid];

    /* grad_w_per_batch[batch, tid, k] = h_bi × s_grad_logits[k] */
    const int row_off = (batch * HIDDEN_DIM + tid) * K_CLASSES;
    #pragma unroll
    for (int k = 0; k < K_CLASSES; ++k) {
        grad_w_per_batch[row_off + k] = h_bi * s_grad_logits[k];
    }

    /* grad_h_t[batch, tid] = Σ_k W[tid, k] × s_grad_logits[k] */
    float acc = 0.0f;
    #pragma unroll
    for (int k = 0; k < K_CLASSES; ++k) {
        acc += w[tid * K_CLASSES + k] * s_grad_logits[k];
    }
    grad_h_t[batch * HIDDEN_DIM + tid] = acc;

    /* grad_b_per_batch[batch, k] = s_grad_logits[k] — threads 0..2 only. */
    if (tid < K_CLASSES) {
        grad_b_per_batch[batch * K_CLASSES + tid] = s_grad_logits[tid];
    }
}
