// crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu
//
// SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up extractor.
//
// Role
// ────
// At rollout step t, after the aux head's K=2 softmax forward produces the
// per-env tile `aux_softmax [n_envs, K=2]`, this kernel copies the second
// column (p_up = softmax[1]) into the per-env cache buffer `prev_aux_dir_prob
// [n_envs]`. The cache is read by `experience_state_gather` at step t+1 to
// populate `state[AUX_DIR_PROB_INDEX = SL_PADDING_START + 0 = 121]`.
//
// Why a separate copy kernel (not direct softmax → state)
// ───────────────────────────────────────────────────────
// Step t's state was already assembled BEFORE step t's aux forward ran (state
// is the network input, aux output is a forward result). So the state at step
// t+1 reads the aux output from step t — a one-step lag. The cache buffer is
// the lag carrier. H=200 label horizon makes the aux signal slow-moving, so
// one-bar lag is fine (per `pearl_first_observation_bootstrap`'s tolerance for
// slow signals on shared scalars).
//
// Why one slot per env (not batch mean)
// ─────────────────────────────────────
// Unlike `aux_pred_to_isv_tanh_kernel` (which reduces the batch to a single
// shared scalar in ISV[375]), this kernel preserves per-env granularity
// because each env's state row needs its own aux signal. The aux softmax is
// already per-env (`[n_envs, K]`), so this is just a strided gather into
// `[n_envs]`.
//
// Sentinel / cold-start / FoldReset
// ─────────────────────────────────
// SP22 H6 Phase 2 (2026-05-12): buffer is filled with 0.0 (neutral
// under the recentered `2*p - 1` encoding; aux contributes nothing) at
// construction and at every FoldReset by the GPU `fill_f32` kernel
// (epsilon_greedy_kernel.cu). Pre-Phase-2 the sentinel was 0.5, which
// violated `pearl_first_observation_bootstrap` (every other state slot
// uses 0 as the "no signal" baseline). This kernel ALWAYS overwrites
// — no sentinel branch needed here. The first rollout step's state
// assembly reads 0.0 (no aux signal yet), then this kernel writes the
// real recentered value for step 2's state, and so on.
//
// Discipline
// ──────────
// - No atomicAdd per `feedback_no_atomicadd.md` (no reduction, just gather).
// - No HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md` (GPU↔GPU only).
// - Stream-ordered with the producer that wrote the softmax tile; same-stream
//   barrier suffices (cf. `aux_pred_to_isv_tanh_kernel`'s same-stream contract
//   with `aux_next_bar_loss_reduce`).
//
// Launch
// ──────
// grid = ((n_envs + 255) / 256), block = 256. K is passed runtime so the
// kernel ABI stays stable if the aux head's K flips (currently 2).
//
// Cost
// ────
// One coalesced 4-byte load + one 4-byte store per env. n_envs is O(1k); ~1µs.

#include <cuda_runtime.h>

extern "C" __global__ void aux_softmax_to_per_env_kernel(
    /* Per-env K-way softmax tile [n_envs, K]. Written by
     * `aux_next_bar_forward` during the captured rollout-step graph; valid
     * after same-stream sync. Layout row-major; this kernel reads
     * `aux_softmax[env * K + 1]` per env. */
    const float* __restrict__ aux_softmax,
    /* Per-env cache buffer [n_envs]. Written here, read by
     * `experience_state_gather` and `backtest_state_gather*` at the next
     * step's state assembly. */
    float*       __restrict__ prev_aux_dir_prob,
    /* Number of envs (rollout-time batch width = `alloc_episodes`). */
    int n_envs,
    /* Softmax K (= AUX_NEXT_BAR_K = 2 in B1.1a). Kernel reads slot K=1
     * (p_up). Passed runtime so the kernel signature stays stable across
     * head-dim flips — the only index read is `env * K + 1`. */
    int K
) {
    int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= n_envs) return;

    /* SP22 H6 Phase 2 (2026-05-12): write RECENTERED p_up = `2*p - 1`
     * so slot 121 uses the same "no signal = 0" baseline as every
     * other state slot, per `pearl_first_observation_bootstrap`.
     * Range: [-1, +1] (softmax components non-negative and sum to 1,
     * so 2*p - 1 is structurally bounded). Sentinel 0 = neutral
     * (aux contributes nothing); +1 = up with full conviction;
     * -1 = down with full conviction. */
    prev_aux_dir_prob[env] =
        2.0f * aux_softmax[(size_t)env * (size_t)K + 1] - 1.0f;
}
