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

//! Invariant tests for the ISV-driven λ controller.
//!
//! Per `feedback_no_cpu_test_fallbacks`: GPU oracle only. Per
//! `pearl_tests_must_prove_not_lock_observations`: assertions check
//! invariants (sentinel bootstrap, steady-state at target, excess
//! proportionality, floor activation under zero amplitude) — not
//! arbitrary numeric snapshots.
//!
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: all CPU↔GPU
//! transfers go via `MappedF32Buffer` + `memcpy_dtod_async`. The i32
//! sentinel is shuttled by reinterpreting the same mapped-pinned
//! buffer as i32 (canonical pattern, see `loss.rs::download_i32`).

use std::sync::Arc;

use anyhow::{Context, Result};
use approx::assert_relative_eq;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_alpha::heads::N_HORIZONS;
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_core::device::MlDevice;

const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/smoothness_lambda_controller.cubin"));
const LAMBDA_FLOOR: f32 = 1.0e-4;

struct ControllerOut {
    jitter_ema: Vec<f32>,
    first_obs: i32,
    lambda: Vec<f32>,
}

fn run_controller(
    dev: &MlDevice,
    raw_per_h: &[f32; N_HORIZONS],
    jitter_ema_init: &[f32; N_HORIZONS],
    first_obs_init: i32,
    base_lambda: f32,
) -> Result<ControllerOut> {
    let stream: &Arc<CudaStream> = dev.cuda_stream().context("controller stream")?;
    let ctx = dev.cuda_context().context("controller ctx")?;
    let module = ctx.load_cubin(CUBIN.to_vec()).context("load controller cubin")?;
    let func = module.load_function("smoothness_lambda_controller").context("load fn")?;

    let raw_d = upload_f32(stream, raw_per_h.as_slice())?;
    let mut jitter_d = upload_f32(stream, jitter_ema_init.as_slice())?;
    let mut first_obs_d = upload_i32(stream, &[first_obs_init])?;
    let mut lambda_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("lambda alloc")?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (N_HORIZONS as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    // The kernel now accepts two trailing pointer args for the GPU log
    // ring (`LogRing*` + `int* g_step_counter`). Passing null pointers
    // makes `log_record()` short-circuit at the `g_log_ring == nullptr`
    // guard — behaviorally identical to the pre-log signature. We must
    // still push the args because launch_builder validates the count.
    let null_ring: cudarc::driver::sys::CUdeviceptr = 0;
    let null_step: cudarc::driver::sys::CUdeviceptr = 0;
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&raw_d)
        .arg(&mut jitter_d)
        .arg(&mut first_obs_d)
        .arg(&base_lambda)
        .arg(&mut lambda_d)
        .arg(&null_ring)
        .arg(&null_step);
    unsafe {
        launch.launch(cfg).context("controller launch")?;
    }
    stream.synchronize().context("controller sync")?;

    Ok(ControllerOut {
        jitter_ema: download_f32(stream, &jitter_d)?,
        first_obs: download_i32(stream, &first_obs_d)?[0],
        lambda: download_f32(stream, &lambda_d)?,
    })
}

fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr,
                staging.dev_ptr,
                nbytes,
                stream.cu_stream(),
            )
            .context("upload DtoD")?;
        }
    }
    Ok(dst)
}

/// Upload an `i32` slice via the canonical mapped-pinned path by
/// reinterpreting `MappedF32Buffer`'s host pointer as `*mut i32`
/// (same pattern as `trainer/loss.rs::download_i32`). One stream sync
/// barrier inside the kernel launch makes the host write visible
/// to the device through the same mapped page.
fn upload_i32(stream: &Arc<CudaStream>, host: &[i32]) -> Result<CudaSlice<i32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("i32 upload staging: {e}"))?;
    // Write through the host pointer reinterpreted as i32.
    unsafe {
        let host_ptr = staging.host_ptr.cast::<i32>();
        for (i, &v) in host.iter().enumerate() {
            std::ptr::write_volatile(host_ptr.add(i), v);
        }
    }
    let mut dst = stream.alloc_zeros::<i32>(n).context("i32 upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<i32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr,
                staging.dev_ptr,
                nbytes,
                stream.cu_stream(),
            )
            .context("i32 upload DtoD")?;
        }
    }
    // Keep staging alive until the DtoD has been issued onto the stream.
    // `dst` carries an internal record; the subsequent kernel launch and
    // sync barrier in `run_controller` flushes the copy before staging
    // drops here.
    stream.synchronize().context("i32 upload sync")?;
    drop(staging);
    Ok(dst)
}

fn download_f32(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(
            staging.dev_ptr,
            src_ptr,
            nbytes,
            stream.cu_stream(),
        )
        .context("download DtoD")?;
    }
    stream.synchronize().context("download sync")?;
    Ok(staging.read_all())
}

/// Download an `i32` slice via the canonical mapped-pinned path
/// (matches `trainer/loss.rs::download_i32`).
fn download_i32(stream: &Arc<CudaStream>, src: &CudaSlice<i32>) -> Result<Vec<i32>> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("i32 download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<i32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(
            staging.dev_ptr,
            src_ptr,
            nbytes,
            stream.cu_stream(),
        )
        .context("i32 download DtoD")?;
    }
    stream.synchronize().context("i32 download sync")?;
    let host_ptr = staging.host_ptr.cast::<i32>();
    let mut out = Vec::with_capacity(n);
    unsafe {
        for i in 0..n {
            out.push(std::ptr::read_volatile(host_ptr.add(i)));
        }
    }
    Ok(out)
}

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

// ───── Tests ─────

#[test]
fn first_observation_bootstraps_ema_and_sets_sentinel() {
    let dev = test_device();
    // HORIZONS=[10, 100, 1000]. Geometric decay (factor 1/10) keeps every
    // raw[h] AT OR BELOW jitter_ema[0] * TARGET_K_RATIO[h] for any
    // monotonically-decreasing TARGET_K_RATIO, so excess_ratio = 0 across
    // the board on the bootstrap step. Invariant under test: sentinel
    // advances 0→1, EMA is replaced by raw, and λ relaxes to base.
    let raw = [0.5_f32, 0.05, 0.005];
    let zero_ema = [0.0_f32; N_HORIZONS];
    let base = 0.01_f32;
    let out = run_controller(&dev, &raw, &zero_ema, 0, base).unwrap();
    assert_eq!(
        out.first_obs, 1,
        "sentinel should advance to 1 on first observation"
    );
    for h in 0..N_HORIZONS {
        assert_relative_eq!(out.jitter_ema[h], raw[h], epsilon = 1e-6);
        // excess_ratio = 0 → lambda = base; floor doesn't engage (base > floor).
        assert_relative_eq!(out.lambda[h], base, epsilon = 1e-6);
    }
}

#[test]
fn steady_state_at_target_yields_base_lambda() {
    let dev = test_device();
    // Geometric decay (factor 1/10) for HORIZONS=[10, 100, 1000].
    // All jitter values at-or-below their derived target → excess_ratio = 0
    // → λ relaxes to base. EMA update is a no-op when raw == ema.
    let ema = [0.5_f32, 0.05, 0.005];
    let raw = ema; // identical → EMA update is a no-op (0.5*old + 0.5*old = old)
    let base = 0.01_f32;
    let out = run_controller(&dev, &raw, &ema, 1, base).unwrap();
    for h in 0..N_HORIZONS {
        assert_relative_eq!(out.jitter_ema[h], ema[h], epsilon = 1e-6);
        assert_relative_eq!(out.lambda[h], base, epsilon = 1e-5);
    }
}

#[test]
fn excess_at_h1000_lifts_lambda_proportionally() {
    let dev = test_device();
    // HORIZONS=[10, 100, 1000]. Drive only h1000 (index 2) above its
    // controller-derived target: jitter[2] = 10 × target[2] should yield
    // excess_ratio = 9 → λ[2] = base × 10. The other horizons stay
    // at-or-below target so λ[h<2] relaxes to base.
    //
    // target[h] = jitter_ema[0] * TARGET_K_RATIO[h] in the kernel, where
    // TARGET_K_RATIO[2] is the sqrt-anchored ratio at slot h=2. With
    // ema[0] = 0.5 the kernel's slot-2 target ≈ 0.5 * TARGET_K_RATIO[2].
    // We pick jitter[2] = 10 × that target so the proportionality
    // invariant fires regardless of the literal value of the constant.
    const TARGET_K_RATIO_H2: f32 = 0.3162278; // must mirror cuda/smoothness_lambda_controller.cu
    let ema = [0.5_f32, 0.05, 10.0 * 0.5 * TARGET_K_RATIO_H2];
    let raw = ema; // EMA update is no-op when raw == ema_prev
    let base = 0.01_f32;
    let out = run_controller(&dev, &raw, &ema, 1, base).unwrap();
    // h10: at-target (self-ratio = 1.0) → lambda = base
    assert_relative_eq!(out.lambda[0], base, epsilon = 1e-5);
    // h100: below-target → lambda = base
    assert_relative_eq!(out.lambda[1], base, epsilon = 1e-5);
    // h1000: excess_ratio = 10 - 1 = 9 → lambda = base * (1 + 9) = 0.10
    let expected_lambda_h1000 = base * 10.0;
    assert_relative_eq!(out.lambda[2], expected_lambda_h1000, max_relative = 1e-3);
}

#[test]
fn lambda_floor_when_base_lambda_zero() {
    let dev = test_device();
    let ema = [0.5_f32; N_HORIZONS];
    let raw = ema;
    let base = 0.0_f32;
    let out = run_controller(&dev, &raw, &ema, 1, base).unwrap();
    // base = 0 → lambda_new = 0 (regardless of excess_ratio) → floor activates
    for h in 0..N_HORIZONS {
        assert!(
            (out.lambda[h] - LAMBDA_FLOOR).abs() < 1e-9,
            "lambda[{h}] should equal LAMBDA_FLOOR (={LAMBDA_FLOOR:.1e}); got {}",
            out.lambda[h]
        );
    }
}