// crates/ml/src/cuda_pipeline/aux_trade_outcome_loss_reduce_kernel.cu
//
// SP22 H6 vNext Phase A4 (2026-05-13) — sparse cross-entropy reduce over
// the K=3 trade-outcome softmax tile produced by
// `aux_trade_outcome_forward`.
//
// Mirrors `aux_next_bar_loss_reduce`'s single-block shmem-tree reduce
// pattern (mean over valid-row count, fmaxf(p_tgt, 1e-30) numerical floor,
// saved valid_count for backward). Kept as a SEPARATE kernel from the K=2
// sibling because:
//
//   1. Diagnostic isolation: the trade-outcome head's loss feeds a
//      DIFFERENT HEALTH_DIAG slot (`aux_outcome_ce_ema`, allocated in
//      Phase F) than the next-bar head's `aux_next_bar_mse_ema_idx`.
//      Two distinct kernels = two distinct cubins = clean
//      per-loss-source attribution in profiles and reward decomposition.
//
//   2. Sparse-label semantics: the trade-outcome producer emits a `-1`
//      mask on EVERY bar without a trade-close event, which is the
//      common case (~95-99% of bars). The next-bar head sees ~50-100%
//      valid labels (only the trailing H bars near series-end get
//      masked). The aggressive sparsity changes the gradient-scale
//      arithmetic: an active-batch fraction of ~1-5% means each valid
//      row's gradient gets divided by a much smaller B_valid → higher
//      per-trade-close gradient magnitude. The kernel logic handles
//      this correctly via `fmaxf(valid, 1.0f)` divisor (matches the
//      backward kernel's `inv_B = 1 / B_valid` semantics) but the
//      independent kernel lets Phase F's Adam-LR controller tune the
//      trade-outcome head separately from the next-bar head if
//      magnitude calibration becomes necessary.
//
//   3. Future per-class weighting: Profit/Stop/Timeout class imbalance
//      can be ~3:1 to ~10:1 depending on the regime. Phase D+ may need
//      class-weighted CE; an isolated kernel makes the modification
//      surgical without touching the K=2 head's contract.
//
// Inputs / outputs (mirror the K=2 sibling contract)
// ──────────────────────────────────────────────────
// softmax           [B, K=3] f32 row-major — saved by
//                                  `aux_trade_outcome_forward`.
// labels            [B] i32 in {-1, 0, 1, 2} — from
//                                  `trade_outcome_label_kernel` at the
//                                  same step (sparse: -1 = no trade
//                                  close at this bar).
// B                 batch size (= N × L in the replay-batch use case;
//                                  = n_envs in the rollout-per-step use
//                                  case). K-runtime arg even though this
//                                  kernel is K=3-semantic — keeps the
//                                  launcher signature symmetric with
//                                  the K=2 sibling for testing.
// loss_out          [1] mean CE over the B_valid valid rows.
// valid_count_out   [1] B_valid (saved for backward — matches the K=2
//                                  sibling's contract).
//
// Discipline
// ──────────
// - No atomicAdd per `feedback_no_atomicadd.md` — single-block shmem-
//   tree reduction.
// - GPU-only per `feedback_cpu_is_read_only.md`.
// - Numerical floor at 1e-30 prevents -log(0) propagation through the
//   EMA when an untrained softmax emits near-zero probabilities.
//
// Block: AUX_BLOCK threads. Grid: (1, 1, 1).
// Shared memory: 2 * AUX_BLOCK floats (loss numerator + valid-count
// partials, lockstep-reduced).

#include <cuda_runtime.h>
#include <math_constants.h>

#ifndef AUX_BLOCK
#define AUX_BLOCK 256
#endif

extern "C" __global__ void aux_trade_outcome_loss_reduce(
    /* [B, K=3] f32 row-major softmax tile from
     * `aux_trade_outcome_forward`. */
    const float* __restrict__ softmax,
    /* [B] i32 labels in `{-1, 0=Profit, 1=Stop, 2=Timeout}`. Producer is
     * `trade_outcome_label_kernel.cu` — emits -1 (mask) on bars with no
     * trade-close event (the common case at ~95-99% of bars). */
    const int* __restrict__ labels,
    int B,
    int K,
    /* [1] mean CE over the B_valid rows whose label != -1. An all-skip
     * batch produces 0 (no NaN; fmaxf(valid, 1) divisor). */
    float* __restrict__ loss_out,
    /* [1] B_valid saved for backward. Matches the K=2 sibling's contract
     * — the backward kernel uses this as `inv_B = 1 / B_valid` for the
     * d_logits scaling. An all-skip batch saves 0; backward then emits
     * zero gradients across the board (no NaN). */
    float* __restrict__ valid_count_out
) {
    if (blockIdx.x != 0) return;
    extern __shared__ float smem[];
    float* sh_loss  = smem;
    float* sh_valid = smem + blockDim.x;
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    /* Per-thread strided accumulation: each thread folds its slice of
     * `[0, B)` into private (loss_numer, valid_count) accumulators.
     * Mask-encoded rows (label == -1) contribute 0 to BOTH accumulators —
     * the common case for trade-outcome labels where most bars have no
     * trade-close event. The valid-row branch reads `softmax[i, label]`
     * directly (no max-shift needed; the forward kernel already
     * materialised stable softmax probs). */
    float local_loss  = 0.0f;
    float local_valid = 0.0f;
    for (int i = tid; i < B; i += block) {
        const int tgt = labels[i];
        if (tgt == -1) continue;
        /* Defensive bounds — out-of-range labels degrade to skip
         * rather than read past the end of the [B, K] softmax row. The
         * label producer (Phase A2 `trade_outcome_label_kernel`) emits
         * only -1 / {0, 1, 2} so this branch is dead under correct
         * producer behavior; kept for kernel robustness against
         * producer bugs. */
        if (tgt < 0 || tgt >= K) continue;
        const float p_tgt = softmax[(size_t)i * K + tgt];
        /* Numerical floor at 1e-30 — matches the K=2 sibling's clip.
         * An untrained `aux_trade_outcome_forward` with extreme W2 can
         * emit p ≈ 0 for some class; -log(0) = +inf would propagate
         * to the EMA. log(1e-30) ≈ -69 is well outside the trained
         * regime — the floor is robustness, not a substitute for
         * training stability. */
        const float p_clipped = fmaxf(p_tgt, 1e-30f);
        local_loss  += -logf(p_clipped);
        local_valid += 1.0f;
    }
    sh_loss[tid]  = local_loss;
    sh_valid[tid] = local_valid;
    __syncthreads();

    /* Standard log2(BLOCK) tree reduction over both arrays in lockstep
     * — identical pattern to the K=2 sibling so the proven correctness
     * carries over directly. */
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh_loss[tid]  += sh_loss[tid + s];
            sh_valid[tid] += sh_valid[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float valid = sh_valid[0];
        loss_out[0]        = sh_loss[0] / fmaxf(valid, 1.0f);
        valid_count_out[0] = valid;
    }
}
