// rl_frd_softmax_ce_grad.cu — FRD head per-horizon softmax + CE loss +
// gradient w.r.t. logits (SP20 P3 backward stage 1).
//
// Per batch (b) and horizon (h), the 21 atom logits feed a softmax:
//   p[b, h, a] = exp(logits[b, h, a] - max) / Σ_a' exp(logits[b, h, a'] - max)
//
// Cross-entropy against one-hot label:
//   CE[b, h] = -log p[b, h, label[b, h]]
//
// Gradient (standard softmax-CE derivative, mean-reduced over batch
// per the v_head_bwd convention so downstream reduce_axis0 + Adam
// step land mean-grad without an explicit /B):
//   grad_logits[b, h, a] = (1/B) × (p[b, h, a] - 1{a == label[b, h]})
//
// Label sentinel: a negative label index (e.g. -1) means "missing
// horizon" (future return not yet realized at the rightmost edge of
// the snapshot stream); the kernel zeros grad_logits for that row
// and writes loss = 0 so the caller's sum is masked-out.
//
// Block layout: 1 block per (batch, horizon), 21 threads per block.
//   * Cooperative max-reduction across 21 atoms (thread 0 serial in
//     shared mem — 21 atoms is small enough that the warp-shuffle
//     pattern's overhead isn't worth it)
//   * Cooperative sum-exp across 21 atoms (same pattern)
//   * Each thread writes one grad_logits slot
//   * Thread 0 writes loss_per_batch_h[b, h]
//
// Per `feedback_no_atomicadd`: no atomics — per-(batch, horizon, atom)
// sole-writer pattern.
// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs.

#include <stdint.h>

#define FRD_N_HORIZONS  3
#define FRD_N_ATOMS    21
#define FRD_OUT_DIM   (FRD_N_HORIZONS * FRD_N_ATOMS)

extern "C" __global__ void rl_frd_softmax_ce_grad(
    const float* __restrict__ logits,             // [B, FRD_OUT_DIM]
    const int*   __restrict__ labels,             // [B, FRD_N_HORIZONS]
    int                       b_size,
    float*       __restrict__ grad_logits,        // [B, FRD_OUT_DIM]   OUT
    float*       __restrict__ loss_per_batch_h    // [B, FRD_N_HORIZONS] OUT
) {
    const int b = blockIdx.x;
    const int h = blockIdx.y;
    const int a = threadIdx.x;
    if (b >= b_size || h >= FRD_N_HORIZONS || a >= FRD_N_ATOMS) return;

    const int row_off = b * FRD_OUT_DIM + h * FRD_N_ATOMS;

    __shared__ float exp_buf[FRD_N_ATOMS];
    __shared__ float row_max;
    __shared__ float row_sum;
    __shared__ int   s_label;

    // Phase 0: stage label (thread 0) — same value broadcast to all threads.
    if (a == 0) {
        s_label = labels[b * FRD_N_HORIZONS + h];
    }
    __syncthreads();

    const int label = s_label;

    // Sentinel label → emit zero gradient + zero loss for this row.
    // (Caller masks out missing-horizon contributions from the total
    // before backward propagation continues.)
    if (label < 0 || label >= FRD_N_ATOMS) {
        grad_logits[row_off + a] = 0.0f;
        if (a == 0) {
            loss_per_batch_h[b * FRD_N_HORIZONS + h] = 0.0f;
        }
        return;
    }

    // Phase 1: stage logits + compute max (thread 0 serial).
    const float my_logit = logits[row_off + a];
    exp_buf[a] = my_logit;
    __syncthreads();
    if (a == 0) {
        float m = exp_buf[0];
        #pragma unroll
        for (int i = 1; i < FRD_N_ATOMS; ++i) {
            m = fmaxf(m, exp_buf[i]);
        }
        row_max = m;
    }
    __syncthreads();

    // Phase 2: each thread computes its exp(logit - max), stages.
    const float my_exp = expf(my_logit - row_max);
    exp_buf[a] = my_exp;
    __syncthreads();
    if (a == 0) {
        float s = 0.0f;
        #pragma unroll
        for (int i = 0; i < FRD_N_ATOMS; ++i) {
            s += exp_buf[i];
        }
        row_sum = s;
    }
    __syncthreads();

    // Phase 3: per-atom softmax probability + gradient.
    const float p_a = my_exp / row_sum;
    const float target = (a == label) ? 1.0f : 0.0f;
    const float inv_B = 1.0f / (float) b_size;
    grad_logits[row_off + a] = (p_a - target) * inv_B;

    // Phase 4: loss (thread 0 only). p_label = exp_buf[label] / row_sum.
    if (a == 0) {
        const float p_label = exp_buf[label] / row_sum;
        loss_per_batch_h[b * FRD_N_HORIZONS + h] = -logf(fmaxf(p_label, 1e-30f));
    }
}
