/* ══════════════════════════════════════════════════════════════════════════
 * SP20 Phase 1.1 (2026-05-09) — sp20_stats_compute kernel.
 *
 * Component 5 / Kernel 3 of the SP20 design (the *stats* producer of the
 * three-kernel SP20 fused-producer chain; the other two — `sp20_emas_compute`
 * and `sp20_controllers_compute` — land in Phase 1.2 and 1.3 and consume
 * this kernel's outputs as their inputs).
 *
 * Reads a `[B, K=2]` row-major `aux_logits` tensor (the SP14-C aux head's
 * direction logits — `n_classes = AUX_NEXT_BAR_K = 2` for {down, up}, per
 * `gpu_aux_heads.rs:61`, established by SP13 B1.1a) and produces two
 * scalar f32 outputs:
 *
 *   out[0] = aux_conf_p50 — median over the batch of
 *            aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/2
 *   out[1] = aux_conf_std — population std over the batch of aux_conf[i]
 *
 * `aux_conf` is the per-row peak-confidence above the K=2 uniform
 * baseline (uniform softmax at K=2 gives 1/2 per class, so subtracting
 * 1/2 makes the signal zero when the head is maximally uncertain and
 * approaches 1/2 when the head is fully concentrated on one class).
 *
 * K=2 fixup (2026-05-09, post Phase 1.3): the kernel originally assumed
 * K=3 `{short, hold, long}` with baseline `1/3`. The production aux head
 * emits K=2 per `AUX_NEXT_BAR_K = 2` (SP13 B1.1a). Wiring K=2 production
 * aux into a K=3 kernel = OOB reads + corrupt stats. This kernel was
 * retargeted to K=2 atomically with the launcher, tests, and audit doc.
 *
 * Producer cadence: hot-path (per-training-step). Subsequent kernels
 * (sp20_emas_compute, sp20_controllers_compute) Wiener-blend these into
 * ISV slots [510..520) — this kernel itself does NOT write to ISV.
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd` — block tree-reduce only; no atomicAdd.
 *    p50 uses `sp4_histogram_p99<256>` with target_quantile = 0.50
 *    (per-warp tile binning + cumulative-from-top, no atomicAdd).
 *    std uses two block tree-reductions (sum, sum-of-squares) sharing
 *    one shmem tile sequentially.
 *  - `pearl_fused_per_group_statistics_oracle` — one fused kernel
 *    reading `aux_logits` once for both stats; no separate p50 / std
 *    producers redundantly re-streaming the per-row aux_conf.
 *  - `pearl_no_host_branches_in_captured_graph` — single-block kernel,
 *    no host branches; safe to capture in the per-step CUDA Graph.
 *  - `feedback_no_htod_htoh_only_mapped_pinned` — outputs land in a
 *    `MappedF32Buffer<2>` written via `__threadfence_system()` for
 *    PCIe-visible coherence; host reads after stream sync.
 *  - `pearl_symmetric_clamp_audit` — the `aux_conf` formula is itself
 *    bounded by softmax composition (max ∈ [1/2, 1] ⇒ aux_conf ∈
 *    [0, 1/2]); no runtime clamp needed on the per-row signal.
 *  - `feedback_no_cpu_compute_strict` — every operation lives in the
 *    kernel; no CPU fallback for the softmax / max / median / std math.
 *
 * Algorithm (single block, 256 threads):
 *
 *   Pass A: compute `aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/2`
 *           for i in [0, B), strided across threads, into a flat
 *           per-row scratch `aux_conf_buf [B]` in dynamic shmem.
 *           Numerically-stable softmax: `m = max(l[i,0], l[i,1]);
 *           e_c = exp(l[i,c] - m); s = e_0 + e_1;
 *           sm_c = e_c / s; max_sm = max_c sm_c`. Then
 *           aux_conf[i] = max_sm - 0.5f.
 *           __syncthreads() after the pass.
 *
 *   Pass B: block tree-reduce sum(aux_conf) and sum(aux_conf^2) for
 *           the std computation. Two reductions share the second half
 *           of the shmem tile sequentially. Thread 0 captures both.
 *
 *   Pass C: single-block call to `sp4_histogram_p99<256>(aux_conf_buf, B)`
 *           with the percentile target overridden to 0.50 via the
 *           `cumul ≥ ceil(B/2)` cumulative-from-bottom variant
 *           (the `sp4_histogram_p99` cumulative-from-top variant
 *           returns the bin upper edge for the top 1% — for p50 we
 *           re-run pass 3 of the same histogram with a different
 *           cumulative target). To keep the shared `sp4_histogram_p99`
 *           contract intact, we inline pass 1+2 (which fill `s_bins`)
 *           and substitute pass 3 with a from-bottom cumulative scan.
 *           See per-pass code below for the detailed adaptation.
 *
 *   Pass D: thread 0 writes:
 *             out[0] = p50              (median of aux_conf)
 *             out[1] = sqrt(var)        (population std of aux_conf)
 *           via `__threadfence_system()` for mapped-pinned coherence.
 *
 * Launch contract:
 *   grid_dim  = (1, 1, 1)
 *   block_dim = (SP20_STATS_BLOCK, 1, 1)   — 256 threads
 *   shared mem (dynamic):
 *     [SP4_HIST_BINS warp tiles]   8 warps × 256 ints × 4B = 8192 bytes
 *     +  B × sizeof(float)         per-row aux_conf_buf
 *     +  block tile reuses s_bins[0..256] (already 1024 bytes static)
 *   With B=128, shmem ≈ 8192 + 512 = 8704 bytes (well under L40S 48 KB).
 *
 * Args:
 *   aux_logits — `[B, K=2]` row-major aux head logits. f32 device ptr.
 *                Row stride is 2 floats; contiguous per row.
 *   out        — `MappedF32Buffer<2>` device ptr for [p50, std].
 *   B          — batch size (rows of aux_logits).
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>
#include "sp4_histogram_p99.cuh"

#define SP20_STATS_BLOCK 256
#define SP20_K_CLASSES   2
#define SP20_UNIFORM_K2  0.5f            /* 1.0f / 2.0f exactly representable */

extern "C" __global__ void sp20_stats_compute_kernel(
    const float* __restrict__ aux_logits,   /* [B, K=2] row-major        */
    float*       __restrict__ out,          /* MappedF32Buffer<2>        */
    int          B)
{
    /* Single-block launch contract. */
    if (blockIdx.x != 0) return;

    const int tid  = (int)threadIdx.x;
    const int bdim = (int)blockDim.x;

    /* Degenerate guard: empty batch — write zeros and return. The
     * caller's downstream EMA producer (Phase 1.2) treats zero stats
     * as a no-op via Pearl-A bootstrap (sentinel = 0). */
    if (B <= 0) {
        if (tid == 0) {
            out[0] = 0.0f;
            out[1] = 0.0f;
            __threadfence_system();
        }
        return;
    }

    /* Static shared: 256 ints for s_bins (1024 B) + 1 float s_step_max
     * (4 B). Used by sp4_histogram_p99-style passes.                  */
    __shared__ int   s_bins[SP4_HIST_BINS];
    __shared__ float s_step_max;
    __shared__ float s_sum;
    __shared__ float s_sumsq;

    /* Dynamic shared layout (caller MUST size accordingly):
     *   [0 .. warps × SP4_HIST_BINS × sizeof(int))      per-warp tiles
     *   [warps × SP4_HIST_BINS × sizeof(int) ..)        aux_conf_buf [B]
     *
     * Both regions are accessed as `int*` then `float*` views; the
     * caller passes one contiguous dynamic-shmem allocation sized
     * (warps × SP4_HIST_BINS × sizeof(int)) + (B × sizeof(float)).
     */
    extern __shared__ int s_dyn_int[];
    int*   s_warp_tiles = s_dyn_int;
    const int warps     = bdim >> 5;
    float* s_aux_conf   = reinterpret_cast<float*>(
                              s_warp_tiles + warps * SP4_HIST_BINS);

    /* ── Pass A: per-row K=2 softmax-max → aux_conf[i] ──────────────── */
    for (int i = tid; i < B; i += bdim) {
        const int row_off = i * SP20_K_CLASSES;
        const float l0 = aux_logits[row_off + 0];
        const float l1 = aux_logits[row_off + 1];

        /* Numerically-stable softmax: subtract row-max before exp. */
        const float m  = fmaxf(l0, l1);
        const float e0 = __expf(l0 - m);
        const float e1 = __expf(l1 - m);
        const float denom = e0 + e1;
        const float sm0 = e0 / denom;
        const float sm1 = e1 / denom;
        const float max_sm = fmaxf(sm0, sm1);

        s_aux_conf[i] = max_sm - SP20_UNIFORM_K2;
    }
    __syncthreads();

    /* ── Pass B1: block tree-reduce sum(aux_conf) ───────────────────── */
    float local_sum = 0.0f;
    for (int i = tid; i < B; i += bdim) {
        local_sum += s_aux_conf[i];
    }
    /* Reuse s_bins[0..bdim) as a float-via-int reinterpret reduce tile.
     * BLOCK_SIZE=256 ≤ SP4_HIST_BINS=256 so the alias fits exactly. */
    s_bins[tid] = __float_as_int(local_sum);
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            const float a = __int_as_float(s_bins[tid]);
            const float b = __int_as_float(s_bins[tid + s]);
            s_bins[tid] = __float_as_int(a + b);
        }
        __syncthreads();
    }
    if (tid == 0) s_sum = __int_as_float(s_bins[0]);
    __syncthreads();

    /* ── Pass B2: block tree-reduce sum(aux_conf^2) ─────────────────── */
    float local_sumsq = 0.0f;
    for (int i = tid; i < B; i += bdim) {
        const float x = s_aux_conf[i];
        local_sumsq += x * x;
    }
    s_bins[tid] = __float_as_int(local_sumsq);
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            const float a = __int_as_float(s_bins[tid]);
            const float b = __int_as_float(s_bins[tid + s]);
            s_bins[tid] = __float_as_int(a + b);
        }
        __syncthreads();
    }
    if (tid == 0) s_sumsq = __int_as_float(s_bins[0]);
    __syncthreads();

    /* ── Pass C: histogram over aux_conf for p50 (cumulative-from-bottom) ──
     * Mirrors sp4_histogram_p99's pass 1 (max-reduce of |aux_conf|) and
     * pass 2 (per-warp tile binning) verbatim, then substitutes a
     * cumulative-from-bottom scan over s_bins[] to recover the median's
     * bin upper-edge. We can't call sp4_histogram_p99<256>() directly
     * because its pass 3 hardcodes the top-1% (p99) cumulative target;
     * the rest of the algorithm is identical, so the inlining cost is
     * one extra `for (b)` block (we save a kernel boundary).
     */

    /* Pass C.1: max-reduce of |aux_conf| via the same tile alias.
     * aux_conf is non-negative by construction (max_sm ≥ 1/2 ⇒
     * aux_conf ≥ 0), so |x| = x; we still use fabsf to guarantee the
     * sign discipline that sp4_histogram_p99 demands.                 */
    float local_max = 0.0f;
    for (int i = tid; i < B; i += bdim) {
        local_max = fmaxf(local_max, fabsf(s_aux_conf[i]));
    }
    s_bins[tid] = __float_as_int(local_max);
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            const float a = __int_as_float(s_bins[tid]);
            const float b = __int_as_float(s_bins[tid + s]);
            s_bins[tid] = __float_as_int(fmaxf(a, b));
        }
        __syncthreads();
    }
    if (tid == 0) s_step_max = __int_as_float(s_bins[0]);
    __syncthreads();

    /* Degenerate-distribution branch: every aux_conf is 0 (e.g.,
     * uniform-logit input). p50 = 0; std also = 0 by construction. */
    if (s_step_max == 0.0f) {
        if (tid == 0) {
            out[0] = 0.0f;
            out[1] = 0.0f;
            __threadfence_system();
        }
        return;
    }

    const float step_max  = s_step_max;
    const float bin_width = step_max / (float)SP4_HIST_BINS;

    /* Pass C.2: per-warp tile binning (no atomicAdd). Verbatim from
     * sp4_histogram_p99 pass 2 — the caller-provided dynamic shmem
     * already reserves `warps × SP4_HIST_BINS` ints at the front. */
    const int warp_id = tid >> 5;
    const int lane    = tid & 31;

    /* Zero this warp's tile (32 lanes × 8 ints = 256). */
    for (int b = lane; b < SP4_HIST_BINS; b += 32) {
        s_warp_tiles[warp_id * SP4_HIST_BINS + b] = 0;
    }
    __syncwarp();

    for (int i = tid; i < B; i += bdim) {
        int bin_idx = (int)floorf(fabsf(s_aux_conf[i]) / bin_width);
        if (bin_idx >= SP4_HIST_BINS) bin_idx = SP4_HIST_BINS - 1;
        s_warp_tiles[warp_id * SP4_HIST_BINS + bin_idx] += 1;
    }
    __syncthreads();

    /* Tree-reduce warp tiles into s_bins[]. */
    for (int b = tid; b < SP4_HIST_BINS; b += bdim) {
        int sum = 0;
        for (int w = 0; w < warps; ++w) {
            sum += s_warp_tiles[w * SP4_HIST_BINS + b];
        }
        s_bins[b] = sum;
    }
    __syncthreads();

    /* Pass C.3: cumulative-from-BOTTOM scan → p50 = first bin where
     * cumul ≥ ⌈B/2⌉. Mirrors sp4_histogram_p99 pass 3 mechanically
     * but inverts the scan direction (bottom-up for the median, top-
     * down for the 99th percentile).                                  */
    if (tid == 0) {
        const int target = (B + 1) / 2;   /* ceil(B / 2) */
        int cumul   = 0;
        int p50_bin = 0;
        for (int b = 0; b < SP4_HIST_BINS; ++b) {
            cumul += s_bins[b];
            if (cumul >= target) { p50_bin = b; break; }
        }
        /* p50 = upper edge of p50_bin = (p50_bin + 1) × bin_width. */
        const float p50 = (float)(p50_bin + 1) * bin_width;

        /* Population std from the captured sum / sumsq. */
        const float n_f  = (float)B;
        const float mean = s_sum / n_f;
        const float var  = fmaxf(0.0f, s_sumsq / n_f - mean * mean);
        const float std  = sqrtf(var);

        out[0] = p50;
        out[1] = std;
        __threadfence_system();
    }
}
