// crates/ml/src/cuda_pipeline/aux_pred_to_isv_tanh_kernel.cu
//
// SP13 Phase 0a P0a.T4 (2026-05-04): aux-head per-bar prediction →
// ISV[AUX_DIR_PREDICTION_INDEX=375] bounded scalar producer.
// SP13 B1.1a (2026-05-05): rewrite — reads the K=2 softmax tile and
// computes `mean(softmax[:, 1] - softmax[:, 0])` ∈ [-1, +1] (structural
// bound, no `tanh()`). The kernel filename + ISV slot name are
// preserved across the rewrite — every consumer of slot 375 reads the
// same `[-1, +1]` bound it always did. The "tanh" in the filename is
// historical (regression-mode wrapped a raw scalar; B1.1a's softmax
// classifier emits structurally-bounded probabilities directly so no
// runtime squash is needed).
//
// Architectural rationale (SP13 B1.1a per
// `pearl_bounded_modifier_outputs_require_structural_activation.md`):
// the previous regression-mode head emitted an unbounded scalar that
// needed a runtime `tanh()` to bound the ISV[375] consumer. The B1.1a
// head emits softmax probabilities directly: `softmax[1] - softmax[0]`
// is naturally bounded in `[-1, +1]` because the components are
// non-negative and sum to 1. The bound is now STRUCTURAL (a property
// of the model output's activation) rather than a runtime clamp on a
// raw scalar. Removing the tanh eliminates one nonlinearity from the
// ISV[375] producer chain.
//
// Why a batch mean instead of per-bar broadcast: ISV is a single
// `[ISV_TOTAL_DIM]` shared array (see `gpu_dqn_trainer.rs::isv_signals_pinned`
// allocation) — there is NO per-batch tile. Every per-step ISV producer
// reduces a per-bar tile to a single scalar before writing (cf.
// `h_s2_rms_ema_kernel.cu` reducing `save_h_s2 [B, SH2]` to a single
// `producer_step_scratch_buf` slot). The batch mean of the directional
// margin is the natural batch-aggregate signal for a shared scalar — it
// sits at zero when the head has no consensus direction (50/50 split),
// drifts toward +1 when the batch consensus is "up" (softmax[1]
// dominates), drifts toward -1 when the batch consensus is "down".
//
// Sentinel-friendly: per `pearl_first_observation_bootstrap`,
// `(prev == 0.0)` is the cold-start sentinel for slot 375 (constructor
// zero-init; no fold-reset entry — the slot is per-step overwritten so
// any non-zero produces a clean replacement on the first launch via
// the existing per-step write semantics). No additional bootstrap
// branch is needed here because the kernel ALWAYS overwrites the
// slot with the current step's batch mean — it's a per-step state,
// not an EMA, and the registry entry skips slot 375 explicitly per
// the SP13 P0a registry comments. B1.1b will swap this slot's
// per-step write semantics to be aware of mask rows once the producer
// kernel for `aux_sign_labels` lands.
//
// Single block, BLOCK_SIZE=256, one shared-memory float array — no
// atomicAdd per `feedback_no_atomicadd.md`. Pure GPU compute per
// `feedback_no_cpu_compute_strict.md`. Stream-ordered with the
// producer that wrote the softmax tile; same-stream barrier suffices
// (cf. `aux_heads_loss_ema_update`'s same-stream contract with
// `aux_next_bar_loss_reduce`).
//
// Cost per launch: 256 threads × ceil(B/256) strided sum passes (two
// loads per row, no transcendentals — this is a bare additive reduce
// since the squash was retired), plus log2(256)=8 reduce steps + 1
// global write. B is O(1k-10k) in production training; total ~1-2µs
// per step, well under the per-step ISV-producer budget. Cheaper than
// the regression-mode kernel by exactly the cost of the tanh transcend.

#include <cuda_runtime.h>

#define BLOCK_SIZE 256

extern "C" __global__ void aux_pred_to_isv_tanh_kernel(
    /* Per-bar K=2 softmax tile [B, K] — produced inside the captured
     * forward graph by `aux_next_bar_forward` (SP13 B1.1a). Valid after
     * the launcher's same-stream sync against the producer. Layout
     * row-major; the kernel reads `softmax[b * K + 0]` and
     * `softmax[b * K + 1]` per-bar. */
    const float* __restrict__ aux_softmax,
    /* Number of bars in the batch. The strided loop covers any
     * `batch_size > 0`; empty batch falls through to the sentinel
     * write (see thread-0 finaliser). */
    int batch_size,
    /* Softmax K (= AUX_NEXT_BAR_K = 2 in B1.1a). Passed runtime so the
     * kernel signature stays stable across head-dim flips; the only
     * indices read are `b*K + 0` and `b*K + 1`. */
    int K,
    /* ISV slot index for the SHARED scalar — `AUX_DIR_PREDICTION_INDEX = 375`
     * in the SP13 layout. The kernel writes
     *   isv[isv_slot_offset] = mean(softmax[:, 1] - softmax[:, 0])
     * after the tree-reduce, with `__threadfence_system()` to make the
     * write visible to subsequent same-stream consumers. */
    int isv_slot_offset,
    /* ISV array (whole bus). The kernel only touches slot
     * `isv_slot_offset`; no other slots are read or written. */
    float* __restrict__ isv)
{
    /* Single-block reducer — guard against accidental multi-block launch. */
    if (blockIdx.x != 0) return;

    extern __shared__ float sh_sum[];
    const int tid  = threadIdx.x;
    const int bdim = blockDim.x;

    /* Per-thread strided accumulation: each thread folds
     * `(softmax[b, 1] - softmax[b, 0])` over its slice of `[0, batch_size)`.
     * The pairwise difference is structurally bounded in [-1, +1] because
     * each softmax component is in [0, 1] and they sum to 1; the batch
     * mean is therefore also in [-1, +1] (the bound used to come from the
     * tanh squash, but B1.1a's softmax classifier provides it
     * structurally per `pearl_bounded_modifier_outputs_require_structural_activation`). */
    float local_sum = 0.0f;
    for (int i = tid; i < batch_size; i += bdim) {
        const float p0 = aux_softmax[(size_t)i * K + 0];
        const float p1 = aux_softmax[(size_t)i * K + 1];
        local_sum += (p1 - p0);
    }
    sh_sum[tid] = local_sum;
    __syncthreads();

    /* Standard log2(BLOCK_SIZE) tree reduction. */
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh_sum[tid] += sh_sum[tid + s];
        }
        __syncthreads();
    }

    /* Thread 0 writes the batch mean. Empty-batch path writes 0.0
     * (matches the constructor-zero sentinel for slot 375 — the
     * downstream Q-head consumer treats 0.0 as "no aux signal yet"
     * via the same logic that handles the cold-start step before
     * the first forward pass). */
    if (tid == 0) {
        const float mean = (batch_size > 0)
            ? (sh_sum[0] / (float)batch_size)
            : 0.0f;
        isv[isv_slot_offset] = mean;
        __threadfence_system();
    }
}
