foxhunt/crates/ml-alpha/tests/frd_head.rs

//! GPU-oracle tests for the Forward-Return-Distribution head (SP20 P3).
//!
//! Forward-pass invariants verified:
//!   1. `frd_forward_zero_input_emits_zero_logits` — when `h_t = 0` and
//!      both biases are zero (the default `FrdHead::new` init), the
//!      output logits must be exactly zero. Provides an unambiguous
//!      analytical oracle for the matmul + ReLU + matmul chain.
//!   2. `frd_forward_shape_matches_spec` — output buffer length is
//!      `b_size × FRD_OUT_DIM` and the per-horizon softmax sums to 1
//!      within fp tolerance. Exercises a random `h_t` to verify the
//!      head produces valid distributions for any input.
//!   3. `frd_forward_caches_hidden_for_bwd` — after a forward pass, the
//!      cached hidden activation buffer matches a re-computed ReLU
//!      sentinel: when `h_t` is set such that `W1 @ h_t + b1` is
//!      strictly negative for some output (driven via biases), those
//!      slots stay at exactly zero.
//!
//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical
//! (zero-input → zero-output; softmax sum invariant), no CPU reference
//! impl.
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: all host↔device
//! transfers use mapped-pinned helpers.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test frd_head -- --ignored --nocapture`

use anyhow::Result;
use cudarc::driver::{CudaSlice, CudaStream};
use ml_alpha::heads::HIDDEN_DIM;
use ml_alpha::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS};
use ml_alpha::rl::frd::{FrdHead, FrdHeadConfig, FRD_OUT_DIM};
use ml_alpha::trainer::integrated::{
    read_slice_d_pub, write_slice_f32_d_pub, write_slice_i32_d_pub,
};
use ml_core::device::MlDevice;
use std::sync::Arc;

fn build_head() -> Option<(MlDevice, FrdHead)> {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return None;
        }
    };
    let head = FrdHead::new(&dev, FrdHeadConfig { seed: 0xF8D_42 }).expect("FrdHead::new");
    Some((dev, head))
}

fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let mut d = stream.alloc_zeros::<f32>(host.len())?;
    write_slice_f32_d_pub(stream, host, &mut d)?;
    Ok(d)
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_forward_zero_input_emits_zero_logits() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4;

    // h_t = 0, b1 = 0 (default init), so hidden_pre = 0 → ReLU(0) = 0.
    // Then hidden @ W2 + b2 = 0 + 0 = 0 regardless of W2's values.
    let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut hidden_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
    let mut logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;

    head.forward(&h_t_d, &mut hidden_d, &mut logits_d, b_size)?;

    let logits = read_slice_d_pub(&stream, &logits_d, b_size * FRD_OUT_DIM)?;
    let hidden = read_slice_d_pub(&stream, &hidden_d, b_size * FRD_HIDDEN_DIM)?;

    for (i, v) in logits.iter().enumerate() {
        assert_eq!(
            *v, 0.0,
            "logit[{i}] should be exactly 0 for h_t=0 with default b1=b2=0; got {v}"
        );
    }
    for (i, v) in hidden.iter().enumerate() {
        assert_eq!(
            *v, 0.0,
            "hidden[{i}] should be exactly 0 for h_t=0 with default b1=0 (ReLU(0)=0); got {v}"
        );
    }

    eprintln!(
        "PASS — zero-input → zero-logits invariant ({} logits, {} hidden slots)",
        logits.len(),
        hidden.len()
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_forward_shape_matches_spec() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 8;

    // Random h_t in [-1, 1] — Xavier-scaled W1 keeps pre-activations
    // O(0.1), but enough nonzero that softmax distributions are
    // meaningfully non-uniform.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xCAFE);
    let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-1.0..1.0))
        .collect();
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut hidden_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
    let mut logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;

    head.forward(&h_t_d, &mut hidden_d, &mut logits_d, b_size)?;
    let logits = read_slice_d_pub(&stream, &logits_d, b_size * FRD_OUT_DIM)?;

    // Per-horizon softmax sums to 1 (standard softmax invariant).
    for b in 0..b_size {
        for h in 0..FRD_N_HORIZONS {
            let off = b * FRD_OUT_DIM + h * FRD_N_ATOMS;
            let row = &logits[off..off + FRD_N_ATOMS];
            // log-sum-exp for numerical stability.
            let max_l = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let denom: f32 = row.iter().map(|x| (x - max_l).exp()).sum();
            let sum_p: f32 = row.iter().map(|x| (x - max_l).exp() / denom).sum();
            assert!(
                (sum_p - 1.0).abs() < 1e-5,
                "softmax sum for batch {b} horizon {h} should be 1.0; got {sum_p}"
            );
        }
    }

    eprintln!(
        "PASS — output shape {} (= {} × {}); per-horizon softmax sums to 1.0",
        logits.len(),
        b_size,
        FRD_OUT_DIM
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_forward_relu_mask_consistent_with_cached_hidden() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4;

    // Use a uniformly negative input (-1) so SOME hidden pre-activations
    // are strictly negative (depending on the random W1 sign pattern).
    // The cached hidden buffer must contain the post-ReLU values:
    // every value MUST be ≥ 0 (ReLU output is non-negative).
    let h_t = vec![-1.0_f32; b_size * HIDDEN_DIM];
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut hidden_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
    let mut logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;

    head.forward(&h_t_d, &mut hidden_d, &mut logits_d, b_size)?;
    let hidden = read_slice_d_pub(&stream, &hidden_d, b_size * FRD_HIDDEN_DIM)?;

    let mut zero_count = 0;
    for (i, v) in hidden.iter().enumerate() {
        assert!(
            *v >= 0.0,
            "cached hidden[{i}] = {v} violates ReLU non-negativity invariant"
        );
        if *v == 0.0 {
            zero_count += 1;
        }
    }
    // With Xavier-symmetric init around 0, ≈half of the pre-activations
    // should be negative → zero in the cached buffer. Assert at least
    // SOME entries got masked (a sign-test that the ReLU actually fires);
    // we don't lock in an exact ratio since it depends on the RNG.
    assert!(
        zero_count > 0,
        "expected at least one ReLU-masked hidden slot under negative input; got {zero_count} zeros out of {}",
        hidden.len()
    );

    eprintln!(
        "PASS — cached hidden is non-negative; {}/{} slots ReLU-masked (negative input)",
        zero_count,
        hidden.len()
    );
    Ok(())
}

fn upload_i32(stream: &Arc<CudaStream>, host: &[i32]) -> Result<CudaSlice<i32>> {
    let mut d = stream.alloc_zeros::<i32>(host.len())?;
    write_slice_i32_d_pub(stream, host, &mut d)?;
    Ok(d)
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_softmax_ce_grad_uniform_logits_match_log_n_atoms() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4;

    // Uniform logits (all zeros) → softmax = 1/FRD_N_ATOMS everywhere,
    // so CE = -log(1/21) = ln(21) ≈ 3.0445 for ANY label index.
    let logits = vec![0.0_f32; b_size * FRD_OUT_DIM];
    let logits_d = upload_f32(&stream, &logits)?;
    // Set label = 10 (mid-bucket) for every (batch, horizon) row.
    let labels: Vec<i32> = vec![10; b_size * FRD_N_HORIZONS];
    let labels_d = upload_i32(&stream, &labels)?;
    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;

    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;

    let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
    let expected = (FRD_N_ATOMS as f32).ln();
    for (i, v) in loss.iter().enumerate() {
        assert!(
            (v - expected).abs() < 1e-5,
            "uniform-logit CE[{i}] should be ln({}) = {expected}; got {v}",
            FRD_N_ATOMS
        );
    }

    // Softmax-CE gradient invariant: Σ_a (p[a] - 1{a==label}) = 0
    // (mean-reduced over B doesn't change the per-row sum being zero).
    let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
    for b in 0..b_size {
        for h in 0..FRD_N_HORIZONS {
            let off = b * FRD_OUT_DIM + h * FRD_N_ATOMS;
            let row_sum: f32 = grad[off..off + FRD_N_ATOMS].iter().sum();
            assert!(
                row_sum.abs() < 1e-6,
                "Σ grad_logits per (b={b}, h={h}) must be 0 (softmax-CE invariant); got {row_sum}"
            );
        }
    }

    eprintln!(
        "PASS — uniform logits → CE = ln({}) = {:.4}; per-row grad sum = 0",
        FRD_N_ATOMS, expected
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_softmax_ce_grad_sentinel_label_zeros_row() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2;

    // Non-trivial logits so a non-sentinel label would produce non-zero
    // loss and gradient — proving the sentinel really is masking.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0x5EE);
    let logits: Vec<f32> = (0..b_size * FRD_OUT_DIM)
        .map(|_| r.gen_range(-1.0..1.0))
        .collect();
    let logits_d = upload_f32(&stream, &logits)?;
    // Labels: ALL sentinel (-1) → every row should be zeroed.
    let labels: Vec<i32> = vec![-1; b_size * FRD_N_HORIZONS];
    let labels_d = upload_i32(&stream, &labels)?;
    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;

    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;

    let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
    let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
    for (i, v) in loss.iter().enumerate() {
        assert_eq!(*v, 0.0, "sentinel label loss[{i}] must be 0; got {v}");
    }
    for (i, v) in grad.iter().enumerate() {
        assert_eq!(
            *v, 0.0,
            "sentinel label grad_logits[{i}] must be 0; got {v}"
        );
    }

    eprintln!("PASS — sentinel label (-1) zeros loss + grad rows");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> {
    let Some((dev, head)) = build_head() else { return Ok(()) };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 1;

    // Single batch, fixed (h=0, target_a=5). Pick a smooth random logit
    // pattern so probabilities aren't degenerate at 0/1.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xFD1FF);
    let mut logits: Vec<f32> = (0..b_size * FRD_OUT_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    let labels: Vec<i32> = vec![5, 10, 15]; // distinct per horizon
    let labels_d = upload_i32(&stream, &labels)?;

    // Analytical gradient via the kernel.
    let logits_d = upload_f32(&stream, &logits)?;
    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
    let grad_analytical = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;

    // Finite-difference for slot (b=0, h=0, a=3). Note: gradient was
    // mean-reduced by 1/B at the kernel source; for b_size=1 the 1/B
    // factor is 1.0 so finite-diff matches directly.
    let probe_h = 0_usize;
    let probe_a = 3_usize;
    let probe_off = probe_h * FRD_N_ATOMS + probe_a;
    let eps = 1e-3_f32;

    // L(logits + ε · e_j) — perturb only the target slot upward.
    logits[probe_off] += eps;
    let logits_plus_d = upload_f32(&stream, &logits)?;
    head.softmax_ce_grad(&logits_plus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
    let loss_plus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
    let l_plus = loss_plus[probe_h]; // only h=0 affected — h=1,2 share the perturbation only if probe was in their horizon block

    // L(logits - ε · e_j)
    logits[probe_off] -= 2.0 * eps;
    let logits_minus_d = upload_f32(&stream, &logits)?;
    head.softmax_ce_grad(&logits_minus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
    let loss_minus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
    let l_minus = loss_minus[probe_h];

    let numerical = (l_plus - l_minus) / (2.0 * eps);
    let analytical = grad_analytical[probe_off];

    let rel_err = (numerical - analytical).abs() / (analytical.abs().max(1e-6));
    // 5e-3 tolerance: fp32 single-precision finite-difference at ε=1e-3
    // saturates around the per-evaluation rounding error (~1e-7) divided
    // by ε, so 1e-4-ish absolute error is intrinsic to the test method,
    // not the kernel. Tighter threshold would force fp64 finite-diff
    // (not worth the complexity here).
    assert!(
        rel_err < 5e-3,
        "finite-diff gradient mismatch at (h={probe_h}, a={probe_a}): \
         analytical={analytical:.6}, numerical={numerical:.6}, rel_err={rel_err:.6}"
    );

    eprintln!(
        "PASS — finite-diff matches analytical: analytical={:.6} numerical={:.6} rel_err={:.2e}",
        analytical, numerical, rel_err
    );
    Ok(())
}