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,
    };
    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);
    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();
    // raw_per_h satisfies the target ratio exactly (raw[h]/raw[0] = K_h30/K_h),
    // so on the bootstrap step jitter_ema == raw → target == jitter → excess_ratio = 0.
    let raw = [0.5_f32, 0.15, 0.05, 0.015, 0.0025];
    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();
    let ema = [0.5_f32, 0.15, 0.05, 0.015, 0.0025];
    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_h6000_lifts_lambda_proportionally() {
    let dev = test_device();
    // h6000 EMA at 0.025 = 10× target (= jitter_ema[h30] * 30/6000 = 0.0025)
    let ema = [0.5_f32, 0.15, 0.05, 0.015, 0.025];
    let raw = ema; // EMA update is no-op
    let base = 0.01_f32;
    let out = run_controller(&dev, &raw, &ema, 1, base).unwrap();
    // h30: at-target → lambda = base
    assert_relative_eq!(out.lambda[0], base, epsilon = 1e-5);
    // h100, h300, h1000: at-target → lambda = base
    assert_relative_eq!(out.lambda[1], base, epsilon = 1e-5);
    assert_relative_eq!(out.lambda[2], base, epsilon = 1e-5);
    assert_relative_eq!(out.lambda[3], base, epsilon = 1e-5);
    // h6000: excess_ratio = 0.025/0.0025 - 1 = 9 → lambda = base * (1 + 9) = 0.10
    let expected_lambda_h6000 = base * 10.0;
    assert_relative_eq!(out.lambda[4], expected_lambda_h6000, 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]
        );
    }
}