// crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
//
// SP22 H6 Phase 3 α Step 11 → SP22 H6 vNext Phase D Step 11 — Adam
// optimizer update for w_aux_to_q_dir [b0_size=4 × K=3] = 12 weights.
//
// Standard Adam (Kingma & Ba 2015) with bias correction:
//   m_t = β1 × m_{t-1} + (1 - β1) × g_t
//   v_t = β2 × v_{t-1} + (1 - β2) × g_t²
//   m_hat = m_t / (1 - β1^t)
//   v_hat = v_t / (1 - β2^t)
//   W_t = W_{t-1} - lr × m_hat / (√v_hat + ε)
//
// Discipline
// ──────────
// - 12 threads, single block — trivial cost (<1µs). Launched once per
//   training step OUTSIDE the captured forward/backward graph (mirrors
//   the existing main-params Adam launch). Phase D bumped from 4 to 12
//   threads to cover the K=3 per-outcome weights per direction action.
// - β1, β2, ε, lr passed by value from host — host-stable across captured
//   launches per pearl_no_host_branches_in_captured_graph (this kernel is
//   not captured anyway).
// - step is i32 (1-indexed) per the trainer's existing Adam step counter
//   semantics. `__powf` for the bias-correction denominator (sufficient
//   precision for the few-thousand-step horizon).
// - Per pearl_first_observation_bootstrap.md: m and v default to 0 at
//   alloc_zeros init; the first step's bias-correction divides by
//   (1 - β^1) = (1 - β), not (1 - β^0) = 0 → no NaN.
// - No HtoD inside kernel per feedback_no_htod_htoh_only_mapped_pinned.md.

#include <cuda_runtime.h>

/* SP22 H6 vNext Phase D (2026-05-14): W_AUX_DIM bumped 4 → 12. Layout:
 * 4 direction actions × 3 outcome classes = 12 weights, indexed
 * as `w[a * K + k]` where K=3 = AUX_OUTCOME_K. */
#define W_AUX_DIM 12

/* Graph-capture-safe signature: `step` and `lr` are POINTERS to
 * device-mapped pinned scratch the host updates each training step
 * (mirrors the main Adam launcher's `t_buf` + `lr_dev_ptr` pattern).
 * The captured launch records the pointers; the kernel reads the
 * current values at replay time. Beta/eps are constants — host-stable
 * across all replays per pearl_no_host_branches_in_captured_graph. */
extern "C" __global__ void adam_w_aux_update(
    float* __restrict__ w,           /* [W_AUX_DIM=12] in-place update */
    const float* __restrict__ dw,    /* [W_AUX_DIM=12] gradient from c51_aux_dw_kernel */
    float* __restrict__ m,           /* [W_AUX_DIM=12] first moment in-place */
    float* __restrict__ v,           /* [W_AUX_DIM=12] second moment in-place */
    const float* __restrict__ lr_ptr,   /* [1] device-mapped pinned f32 */
    float beta1,
    float beta2,
    float eps,
    const int*   __restrict__ step_ptr  /* [1] device-mapped pinned i32, 1-indexed */
) {
    int a = threadIdx.x;
    if (a >= W_AUX_DIM) return;

    float lr   = lr_ptr[0];
    int   step = step_ptr[0];

    float g = dw[a];
    /* SP22 H6 Phase 3 α Step 11 DEFENSIVE NaN guard (2026-05-13):
     * skip update entirely if any input is non-finite. Belt-and-
     * suspenders with c51_aux_dw_kernel's own guard — prevents NaN
     * propagation into W through any pathway. */
    if (!isfinite(g) || !isfinite(lr) || !isfinite(m[a]) || !isfinite(v[a])) {
        return;
    }
    float m_new = beta1 * m[a] + (1.0f - beta1) * g;
    float v_new = beta2 * v[a] + (1.0f - beta2) * g * g;
    m[a] = m_new;
    v[a] = v_new;

    /* Bias-corrected moments. step is at least 1 on the first invocation
     * (trainer increments before submit_adam_ops), so (1 - β^step) > 0. */
    float bc1 = 1.0f - __powf(beta1, (float)step);
    float bc2 = 1.0f - __powf(beta2, (float)step);
    /* Floors at fp32 epsilon to avoid /0 in degenerate first-step
     * configurations (β1 ≈ 0 or β2 ≈ 0). */
    bc1 = fmaxf(bc1, 1e-30f);
    bc2 = fmaxf(bc2, 1e-30f);
    float m_hat = m_new / bc1;
    float v_hat = v_new / bc2;
    w[a] = w[a] - lr * m_hat / (sqrtf(v_hat) + eps);
}
