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

//! GPU-oracle tests for `rl_signal_variance_update` — Welford online
//! variance kernel for adaptive controller floors (spec
//! `2026-05-30-adaptive-controller-floor-design`).
//!
//! Three analytical invariants, zero CPU reference implementations:
//!
//!   1. `welford_constant_input_zero_variance` — feeding a constant `k`
//!      for N steps must produce `mean == k` and `m2 == 0` (sample
//!      variance is zero by definition).
//!   2. `welford_known_sequence_variance` — sequence [1, 2, 3, 4, 5]
//!      has known sample variance 2.5; Welford output must match.
//!   3. `welford_sentinel_zero_is_skipped` — input == 0.0 sentinel
//!      must produce no update (per
//!      `pearl_first_observation_bootstrap`).
//!
//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical
//! identities of Welford's algorithm, not CPU reimplementations.
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: ISV bus uses
//! mapped-pinned `MappedF32Buffer` (same pattern as production trainer).
//!
//! Run with:
//!   `cargo test -p ml-alpha --test signal_variance_kernel -- --ignored --nocapture`

use cudarc::driver::{LaunchConfig, PushKernelArg};
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_core::device::MlDevice;

// ISV slot layout for this isolated test: input at 0, Welford triple at 1,2,3.
const INPUT_SLOT: i32 = 0;
const COUNT_SLOT: i32 = 1;
const MEAN_SLOT: i32 = 2;
const M2_SLOT: i32 = 3;
const N_SLOTS: usize = 4;

const SIGNAL_VARIANCE_UPDATE_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/rl_signal_variance_update.cubin"));

fn launch_once(
    dev: &MlDevice,
    isv: &MappedF32Buffer,
) {
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");
    let module = ctx
        .load_cubin(SIGNAL_VARIANCE_UPDATE_CUBIN.to_vec())
        .expect("load signal_variance cubin");
    let func = module
        .load_function("rl_signal_variance_update")
        .expect("load rl_signal_variance_update fn");
    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (1, 1, 1),
        shared_mem_bytes: 0,
    };
    let isv_dev_ptr = isv.dev_ptr;
    unsafe {
        stream
            .launch_builder(&func)
            .arg(&isv_dev_ptr)
            .arg(&INPUT_SLOT)
            .arg(&COUNT_SLOT)
            .arg(&MEAN_SLOT)
            .arg(&M2_SLOT)
            .launch(cfg)
            .expect("rl_signal_variance_update launch");
    }
    stream.synchronize().expect("sync");
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn welford_constant_input_zero_variance() {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return;
        }
    };
    let isv = unsafe { MappedF32Buffer::new(N_SLOTS) }.expect("isv alloc");

    // Welford of constant k for N steps: mean = k, m2 = 0, count = N.
    let k = 0.5_f32;
    let n = 50_usize;

    for _ in 0..n {
        isv.write_record(INPUT_SLOT as usize, k);
        launch_once(&dev, &isv);
    }

    let count = isv.read_record(COUNT_SLOT as usize);
    let mean = isv.read_record(MEAN_SLOT as usize);
    let m2 = isv.read_record(M2_SLOT as usize);

    assert_eq!(count, n as f32, "count must be {n}; got {count}");
    assert!(
        (mean - k).abs() < 1e-6,
        "mean of constant {k} must equal {k}; got {mean}"
    );
    assert!(
        m2.abs() < 1e-5,
        "m2 of constant input must be exactly 0 (or fp-noise); got {m2}"
    );

    eprintln!("G1 OK — constant k={k} for N={n}: mean={mean} m2={m2}");
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn welford_known_sequence_variance() {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return;
        }
    };
    let isv = unsafe { MappedF32Buffer::new(N_SLOTS) }.expect("isv alloc");

    // Sequence [1, 2, 3, 4, 5]: arithmetic mean = 3.
    // Sample variance = Σ(x_i − μ)² / (n − 1) = (4+1+0+1+4)/4 = 2.5.
    let seq = [1.0_f32, 2.0, 3.0, 4.0, 5.0];

    for &x in seq.iter() {
        isv.write_record(INPUT_SLOT as usize, x);
        launch_once(&dev, &isv);
    }

    let count = isv.read_record(COUNT_SLOT as usize);
    let mean = isv.read_record(MEAN_SLOT as usize);
    let m2 = isv.read_record(M2_SLOT as usize);
    let sample_var = m2 / (count - 1.0);

    assert_eq!(count, seq.len() as f32, "count must be {}", seq.len());
    assert!(
        (mean - 3.0).abs() < 1e-5,
        "mean of [1..5] must be 3.0; got {mean}"
    );
    assert!(
        (sample_var - 2.5).abs() < 1e-5,
        "sample variance of [1..5] must be 2.5; got {sample_var} (m2={m2})"
    );

    eprintln!("G2 OK — sequence {seq:?}: mean={mean} sample_var={sample_var}");
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn welford_sentinel_zero_is_skipped() {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return;
        }
    };
    let isv = unsafe { MappedF32Buffer::new(N_SLOTS) }.expect("isv alloc");

    // INPUT_SLOT remains at sentinel 0.0 (MappedF32Buffer::new zero-inits).
    // Kernel must NOT count this as an observation per
    // pearl_first_observation_bootstrap.
    for _ in 0..10 {
        launch_once(&dev, &isv);
    }
    assert_eq!(
        isv.read_record(COUNT_SLOT as usize),
        0.0,
        "sentinel zero must not be counted as a real observation"
    );

    // Now feed a real non-zero observation; counter must advance.
    isv.write_record(INPUT_SLOT as usize, 7.0);
    launch_once(&dev, &isv);
    assert_eq!(
        isv.read_record(COUNT_SLOT as usize),
        1.0,
        "first real observation must be counted"
    );
    assert!(
        (isv.read_record(MEAN_SLOT as usize) - 7.0).abs() < 1e-6,
        "first observation must initialize mean directly per Welford recurrence"
    );

    eprintln!("G3 OK — sentinel skip honoured; first real obs initialized mean");
}