test(crt-train): invariant tests for smoothness_lambda_controller

crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs

@@ -0,0 +1,254 @@
+//! 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]
+        );
+    }
+}