foxhunt/crates/ml/tests/sp18_weight_drift_test.rs

//! SP18 v2 weight-drift diagnostic kernel oracle tests (2026-05-09).
//!
//! Verifies the `weight_drift_diag_kernel` produces
//! `[||params - best||₂, ||params - best||₂ / max(||best||₂, EPS)]`
//! to within ε=1e-5 of the CPU oracle on synthetic GPU buffers.
//!
//! Per `feedback_no_atomicadd` the kernel uses block tree-reduce.
//! Per `feedback_no_htod_htoh_only_mapped_pinned` all CPU↔GPU buffers
//! are `MappedF32Buffer`.
//!
//! Run on a GPU host:
//!
//!     SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \
//!       cargo test -p ml --test sp18_weight_drift_test --features cuda \
//!         -- --ignored --nocapture

#![allow(clippy::tests_outside_test_module)]

#[cfg(feature = "cuda")]
#[allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.
mod gpu {
    use std::sync::Arc;

    use cudarc::driver::{CudaContext, CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
    use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;

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

    fn make_test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
        ctx.default_stream()
    }

    fn load_weight_drift_diag(stream: &Arc<CudaStream>) -> CudaFunction {
        let module = stream
            .context()
            .load_cubin(SP18_WEIGHT_DRIFT_DIAG_CUBIN.to_vec())
            .expect("load weight_drift_diag_kernel cubin");
        module
            .load_function("weight_drift_diag_update")
            .expect("load weight_drift_diag_update function")
    }

    /// CPU oracle: `(||params - best||₂, ||params - best||₂ / max(||best||₂, EPS))`.
    /// EPS matches the kernel's `SP18_DRIFT_EPS_F` = 1e-12.
    fn cpu_oracle(params: &[f32], best: &[f32]) -> (f32, f32) {
        assert_eq!(params.len(), best.len());
        let sumsq_diff: f32 = params
            .iter()
            .zip(best.iter())
            .map(|(&p, &b)| (p - b).powi(2))
            .sum();
        let sumsq_best: f32 = best.iter().map(|x| x * x).sum();
        let l2_diff = sumsq_diff.sqrt();
        let norm_best = sumsq_best.sqrt();
        let denom = if norm_best > 1e-12 { norm_best } else { 1e-12 };
        let rel = l2_diff / denom;
        (l2_diff, rel)
    }

    fn launch_drift(
        stream: &Arc<CudaStream>,
        kernel: &CudaFunction,
        params: &MappedF32Buffer,
        best: &MappedF32Buffer,
        out: &MappedF32Buffer,
        n: i32,
    ) {
        const BLOCK_DIM: u32 = 256;
        let shmem_bytes: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(kernel)
                .arg(&params.dev_ptr)
                .arg(&best.dev_ptr)
                .arg(&out.dev_ptr)
                .arg(&n)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (BLOCK_DIM, 1, 1),
                    shared_mem_bytes: shmem_bytes,
                })
                .expect("launch weight_drift_diag_update");
        }
        stream.synchronize().expect("sync after weight_drift_diag");
    }

    /// Test A: synthetic params + best buffers with known L2 distance.
    /// Uses N=4096 (representative of trainer params count) and matches
    /// the CPU oracle to ε=1e-5.
    #[test]
    #[ignore = "requires GPU"]
    fn weight_drift_matches_cpu_oracle_4096() {
        let stream = make_test_stream();
        let kernel = load_weight_drift_diag(&stream);

        const N: usize = 4096;

        // best: smooth pattern (best-Sharpe peak weights).
        let best: Vec<f32> = (0..N).map(|i| (i as f32 * 0.001).sin() * 0.5).collect();
        // params: best + drift (HD5-like decay).
        let params: Vec<f32> = best
            .iter()
            .enumerate()
            .map(|(i, &b)| b + 0.05 * (i as f32 * 0.013).cos())
            .collect();

        let (cpu_l2, cpu_rel) = cpu_oracle(&params, &best);

        let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params");
        params_buf.write_from_slice(&params);
        let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best");
        best_buf.write_from_slice(&best);
        let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out");

        launch_drift(&stream, &kernel, &params_buf, &best_buf, &out_buf, N as i32);

        let out = out_buf.read_all();
        let gpu_l2 = out[0];
        let gpu_rel = out[1];

        let eps = 1e-5_f32;
        // Use relative tolerance for L2 since the magnitude depends on N.
        let l2_tol = (cpu_l2.abs() * 1e-4).max(eps);
        assert!(
            (gpu_l2 - cpu_l2).abs() < l2_tol,
            "L2 mismatch: gpu={} cpu={} diff={:.3e} tol={:.3e}",
            gpu_l2, cpu_l2, (gpu_l2 - cpu_l2).abs(), l2_tol,
        );
        assert!(
            (gpu_rel - cpu_rel).abs() < eps,
            "relative mismatch: gpu={} cpu={} diff={:.3e}",
            gpu_rel, cpu_rel, (gpu_rel - cpu_rel).abs(),
        );
    }

    /// Test B: params == best → drift is exactly zero.
    /// This is the "snapshot just saved, no training step yet" case.
    #[test]
    #[ignore = "requires GPU"]
    fn weight_drift_is_zero_when_params_equals_best() {
        let stream = make_test_stream();
        let kernel = load_weight_drift_diag(&stream);

        const N: usize = 1024;

        let w: Vec<f32> = (0..N).map(|i| (i as f32 * 0.01).sin()).collect();

        let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params");
        params_buf.write_from_slice(&w);
        let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best");
        best_buf.write_from_slice(&w);
        let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out");

        launch_drift(&stream, &kernel, &params_buf, &best_buf, &out_buf, N as i32);

        let out = out_buf.read_all();
        assert!(
            out[0].abs() < 1e-5,
            "L2 should be 0 when params == best, got {}",
            out[0],
        );
        assert!(
            out[1].abs() < 1e-5,
            "relative should be 0 when params == best, got {}",
            out[1],
        );
    }

    /// Test C: best == 0 → relative drift uses EPS floor (not divide by zero).
    /// Catches the cold-start "snapshot uninitialized but kernel still runs"
    /// edge case (e.g. if save_best_params was called with an all-zero
    /// trainer for some test setup).
    #[test]
    #[ignore = "requires GPU"]
    fn weight_drift_eps_floor_when_best_is_zero() {
        let stream = make_test_stream();
        let kernel = load_weight_drift_diag(&stream);

        const N: usize = 256;

        let params: Vec<f32> = vec![0.1; N];
        let best: Vec<f32> = vec![0.0; N];

        let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params");
        params_buf.write_from_slice(&params);
        let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best");
        best_buf.write_from_slice(&best);
        let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out");

        launch_drift(&stream, &kernel, &params_buf, &best_buf, &out_buf, N as i32);

        let out = out_buf.read_all();
        // L2 = sqrt(N × 0.1²) = sqrt(2.56) ≈ 1.6
        let expected_l2 = (N as f32 * 0.01).sqrt();
        assert!(
            (out[0] - expected_l2).abs() < 1e-4,
            "L2 expected ≈ {}, got {}",
            expected_l2, out[0],
        );
        // norm_best = 0, denom = EPS = 1e-12, so rel = L2 / 1e-12 ≈ 1.6e12.
        // Just verify it's finite and very large (not NaN).
        assert!(
            out[1].is_finite(),
            "relative should be finite (got NaN/inf) when best is all-zero: {}",
            out[1],
        );
        assert!(
            out[1] > 1e10,
            "relative should be very large (>1e10) when best is all-zero, got {}",
            out[1],
        );
    }

    /// Test D: n == 0 degenerate guard.
    #[test]
    #[ignore = "requires GPU"]
    fn weight_drift_n_zero_emits_zeros() {
        let stream = make_test_stream();
        let kernel = load_weight_drift_diag(&stream);

        // Allocate something — kernel uses n=0 to early-return.
        let params_buf = unsafe { MappedF32Buffer::new(16) }.expect("alloc params");
        params_buf.write_from_slice(&vec![1.0; 16]);
        let best_buf = unsafe { MappedF32Buffer::new(16) }.expect("alloc best");
        best_buf.write_from_slice(&vec![2.0; 16]);
        let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out");
        // Pre-fill out with sentinel so we can verify it gets written to 0.
        out_buf.write_from_slice(&[99.0_f32, 99.0_f32]);

        launch_drift(&stream, &kernel, &params_buf, &best_buf, &out_buf, 0_i32);

        let out = out_buf.read_all();
        assert_eq!(out[0], 0.0, "L2 should be 0 for n=0, got {}", out[0]);
        assert_eq!(out[1], 0.0, "relative should be 0 for n=0, got {}", out[1]);
    }
}