foxhunt/crates/ml/tests/sp20_stats_compute_test.rs

//! SP20 Phase 1.1 (2026-05-09) — `sp20_stats_compute_kernel` GPU oracle tests.
//!
//! Verifies the fused stats producer emits
//! `[aux_conf_p50, aux_conf_std]` consistent with the closed-form CPU
//! oracle for two characteristic distributions of `aux_logits [B, K=2]`:
//!
//!   1. **Uniform**: `aux_logits = 0` everywhere → softmax = [0.5, 0.5]
//!      → `aux_conf = max - 0.5 = 0` → p50 ≈ 0, std ≈ 0.
//!   2. **Concentrated**: `aux_logits[*, 0] = 10`, rest = 0 → softmax
//!      ≈ [1, 0] (one-hot) → `aux_conf ≈ 0.5` → p50 ≈ 0.5, std ≈ 0.
//!
//! Tests also cover a heterogeneous distribution with non-trivial p50
//! and std so the kernel is exercised on both the histogram path and
//! the running-sum/sum-of-squares paths together.
//!
//! K=2 fixup (2026-05-09, post Phase 1.3): the kernel originally assumed
//! K=3 `{down, flat, up}` but the production aux head emits K=2
//! `{down, up}` per `gpu_aux_heads.rs:61` (`AUX_NEXT_BAR_K = 2`,
//! established by SP13 B1.1a). This test file was rewritten to exercise
//! the K=2 contract; the CPU oracle, expected p50/std ranges, and the
//! degenerate-uniform branch all reflect the K=2 baseline `1/2 = 0.5`
//! and the resulting `aux_conf ∈ [0, 0.5]` range.
//!
//! Per `feedback_no_atomicadd` the kernel uses block tree-reduce + per-
//! warp tile binning. 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 sp20_stats_compute_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};
    use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;
    use ml::cuda_pipeline::sp20_stats_compute::{
        launch_sp20_stats_compute, AUX_K_CLASSES,
    };

    const SP20_STATS_COMPUTE_CUBIN: &[u8] =
        include_bytes!(concat!(env!("OUT_DIR"), "/sp20_stats_compute_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_sp20_stats_compute(stream: &Arc<CudaStream>) -> CudaFunction {
        let module = stream
            .context()
            .load_cubin(SP20_STATS_COMPUTE_CUBIN.to_vec())
            .expect("load sp20_stats_compute_kernel cubin");
        module
            .load_function("sp20_stats_compute_kernel")
            .expect("load sp20_stats_compute_kernel function")
    }

    /// CPU oracle: per-row K=2 softmax → max → minus 1/2, then median + std.
    /// Mirrors the kernel's math exactly so the tolerances below pin
    /// fp32 rounding only, not algorithmic deltas.
    fn cpu_oracle(aux_logits: &[f32], b: usize) -> (f32, f32) {
        assert_eq!(aux_logits.len(), b * AUX_K_CLASSES);
        let one_half = 0.5_f32;

        let mut conf: Vec<f32> = (0..b)
            .map(|i| {
                let off = i * AUX_K_CLASSES;
                let l0 = aux_logits[off];
                let l1 = aux_logits[off + 1];
                let m = l0.max(l1);
                let e0 = (l0 - m).exp();
                let e1 = (l1 - m).exp();
                let s = e0 + e1;
                let max_sm = (e0 / s).max(e1 / s);
                max_sm - one_half
            })
            .collect();

        // Sort for the median: matches the histogram-bin upper-edge
        // semantics of the kernel (the sorted-array median is what the
        // histogram converges to as bin_width → 0; for our two
        // characteristic tests below the values are concentrated at
        // 0 or 0.5 which the histogram resolves exactly).
        conf.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median_idx = (b + 1) / 2 - 1;            // 0-based ⌈B/2⌉ - 1
        let median = conf[median_idx];

        let mean: f32 = conf.iter().sum::<f32>() / (b as f32);
        let var: f32 = conf
            .iter()
            .map(|x| (x - mean).powi(2))
            .sum::<f32>() / (b as f32);
        let std = var.max(0.0).sqrt();

        (median, std)
    }

    fn run_kernel(b: usize, aux_logits: &[f32]) -> (f32, f32) {
        let stream = make_test_stream();
        let kernel = load_sp20_stats_compute(&stream);

        let aux_buf = unsafe { MappedF32Buffer::new(aux_logits.len()) }
            .expect("alloc aux_logits buf");
        aux_buf.write_from_slice(aux_logits);
        let out_buf = unsafe { MappedF32Buffer::new(2) }
            .expect("alloc out buf");
        out_buf.write_from_slice(&[0.0, 0.0]);

        unsafe {
            launch_sp20_stats_compute(
                &stream,
                &kernel,
                aux_buf.dev_ptr,
                out_buf.dev_ptr,
                b as i32,
            )
            .expect("launch sp20_stats_compute_kernel");
        }
        stream.synchronize().expect("sync after sp20_stats_compute_kernel");

        let out = out_buf.read_all();
        (out[0], out[1])
    }

    /// Test 1 — uniform-logit distribution.
    ///
    /// `aux_logits = 0` everywhere ⇒ softmax = [0.5, 0.5] for every
    /// row ⇒ `aux_conf = 0.5 - 0.5 = 0` for every row.
    /// Expected: p50 ≈ 0, std ≈ 0. The kernel's degenerate-distribution
    /// branch (step_max == 0) writes both outputs to 0 directly.
    #[test]
    #[ignore = "requires GPU"]
    fn uniform_logits_emit_zero_stats() {
        const B: usize = 128;
        let aux_logits = vec![0.0_f32; B * AUX_K_CLASSES];

        let (p50, std) = run_kernel(B, &aux_logits);

        // Both outputs MUST be exactly zero — the kernel's
        // step_max == 0 branch writes 0 directly.
        assert!(
            p50.abs() < 1e-6,
            "uniform: p50 expected 0, got {p50}",
        );
        assert!(
            std.abs() < 1e-6,
            "uniform: std expected 0, got {std}",
        );
    }

    /// Test 2 — concentrated-logit distribution (varied confidence).
    ///
    /// Per-row class-0 logit varies linearly from 0 to 3 across the
    /// batch, with class 1 held at 0. softmax(class 0) varies from 0.5
    /// (logit 0, uniform) to ≈ 0.953 (logit 3), so `aux_conf` varies
    /// from 0 to ≈ 0.453 — a 0.45 range spread across well over half of
    /// the 256 histogram bins (per
    /// `pearl_sp4_histogram_warp_tile_undercount`: lockstep-uniform
    /// inputs cause warp-tile collisions in the non-atomic per-warp
    /// histogram tiles, so test data MUST vary per row across many
    /// bins, not just within one bin's saturating-softmax tail).
    ///
    /// Logits are kept under the saturating-softmax regime (max < 4)
    /// because the softmax tail compresses 0.95 → 1.0 confidences
    /// into a range smaller than one bin_width (≈ 0.4% of step_max),
    /// re-creating the warp-tile collision pattern this test
    /// explicitly avoids.
    ///
    /// Expected: p50 is the median of a monotonically-increasing
    /// aux_conf vector — i.e., the value at the lower midpoint
    /// (`(B+1)/2 - 1` 0-based). The CPU oracle computes it exactly;
    /// the GPU kernel matches to within a couple of bin_widths.
    #[test]
    #[ignore = "requires GPU"]
    fn varied_confidence_logits_match_cpu_oracle() {
        const B: usize = 128;
        let mut aux_logits = vec![0.0_f32; B * AUX_K_CLASSES];
        for i in 0..B {
            // Linear ramp 0 → 3 on class 0 — confidence ramps from
            // uniform (0.5) to a non-saturating peak (≈ 0.953).
            let frac = (i as f32) / ((B - 1) as f32);
            aux_logits[i * AUX_K_CLASSES] = 3.0 * frac;
        }

        let (p50_gpu, std_gpu) = run_kernel(B, &aux_logits);
        let (p50_cpu, std_cpu) = cpu_oracle(&aux_logits, B);

        // The GPU histogram bin_width ≈ step_max / 256 ≈ 0.45 / 256 ≈
        // 0.0018; agreement to within a few bin_widths is the design
        // tolerance for the histogram path.
        let p50_tol = 0.02_f32;
        assert!(
            (p50_gpu - p50_cpu).abs() < p50_tol,
            "concentrated: p50_gpu={p50_gpu} vs cpu_oracle={p50_cpu} (diff={:.3e}, tol={p50_tol})",
            (p50_gpu - p50_cpu).abs(),
        );
        // CPU oracle sanity: with logits ramping 0 → 3, the median
        // aux_conf is the row at logit ≈ 1.5 → softmax ≈ 0.818 →
        // aux_conf ≈ 0.318. We pin it to (0.10, 0.40) to catch
        // contract drift while tolerating fp32 / softmax math.
        assert!(
            p50_cpu > 0.10_f32 && p50_cpu < 0.40_f32,
            "CPU oracle: concentrated p50 should land in (0.10, 0.40), got {p50_cpu}",
        );
        // std non-trivial; CPU and GPU within fp32 noise.
        let std_tol = 5e-3_f32;
        assert!(
            (std_gpu - std_cpu).abs() < std_tol,
            "concentrated: std_gpu={std_gpu} vs cpu_oracle={std_cpu} (diff={:.3e}, tol={std_tol})",
            (std_gpu - std_cpu).abs(),
        );
        // CPU oracle std: the aux_conf vector spans ≈ 0.45 with a
        // monotonic-but-saturating ramp, so std lands roughly in the
        // 0.05-0.20 range for a sigmoid-shaped spread.
        assert!(
            std_cpu > 0.05_f32 && std_cpu < 0.20_f32,
            "CPU oracle: concentrated std should land in (0.05, 0.20), got {std_cpu}",
        );
    }

    /// Test 3 — heterogeneous distribution exercising both p50 and std.
    ///
    /// Two clusters with ramped (not lockstep) aux_conf within each
    /// cluster:
    ///   - Lower half (rows 0..64): logit-0 ramps 0.0 → 1.0, logit-1
    ///     held at 0 ⇒ aux_conf ramps 0 → sigmoid(1.0)-0.5 ≈ 0.231
    ///   - Upper half (rows 64..128): logit-0 ramps 2.5 → 4.0, logit-1
    ///     held at 0 ⇒ aux_conf ramps sigmoid(2.5)-0.5 ≈ 0.424 to
    ///     sigmoid(4.0)-0.5 ≈ 0.491
    ///
    /// The two ramps are well-separated so the median of the sorted
    /// aux_conf vector lands at the top of the lower cluster
    /// (≈ 0.231). Std is dominated by the inter-cluster gap and lands
    /// near 0.15.
    ///
    /// Per-row ramping (rather than per-row jitter on a constant
    /// cluster center) ensures each warp's 32 lanes compute distinct
    /// bin indices, avoiding the
    /// `pearl_sp4_histogram_warp_tile_undercount` lockstep-uniform
    /// trap where 32 lanes racing on the same non-atomic tile slot
    /// undercount the histogram. Production aux logits are naturally
    /// distributed across bins, so spread test data IS the realistic-
    /// input model.
    #[test]
    #[ignore = "requires GPU"]
    fn heterogeneous_logits_match_cpu_oracle() {
        const B: usize = 128;
        let mut aux_logits = vec![0.0_f32; B * AUX_K_CLASSES];
        // Lower cluster: l0 ramps 0 → 1.0; aux_conf 0 → 0.231.
        for i in 0..(B / 2) {
            let frac = (i as f32) / ((B / 2 - 1) as f32);
            aux_logits[i * AUX_K_CLASSES + 0] = 1.0 * frac;
            // l1 = 0
        }
        // Upper cluster: l0 ramps 2.5 → 4.0; aux_conf 0.424 → 0.491.
        for i in (B / 2)..B {
            let k = i - B / 2;
            let frac = (k as f32) / ((B / 2 - 1) as f32);
            aux_logits[i * AUX_K_CLASSES + 0] = 2.5 + 1.5 * frac;
            // l1 = 0
        }

        let (p50_gpu, std_gpu) = run_kernel(B, &aux_logits);
        let (p50_cpu, std_cpu) = cpu_oracle(&aux_logits, B);

        // p50 sanity: the sorted aux_conf vector has 64 lower-cluster
        // values (0 → 0.231) then 64 upper-cluster values (0.424 →
        // 0.491). ⌈B/2⌉-th element (1-based) = 64th = the largest
        // lower-cluster value ≈ 0.231. Both CPU and GPU should agree
        // within a few bin_widths (bin_width = step_max/256 ≈
        // 0.491/256 ≈ 0.002).
        let p50_tol = 0.02_f32;
        assert!(
            (p50_gpu - p50_cpu).abs() < p50_tol,
            "het: p50_gpu={p50_gpu} vs cpu_oracle={p50_cpu} (diff={:.3e}, tol={p50_tol})",
            (p50_gpu - p50_cpu).abs(),
        );
        // CPU oracle sanity: median should be at the top of the lower
        // cluster, ≈ 0.231. Pin to (0.18, 0.27) to catch contract
        // drift while tolerating fp32 / sigmoid math.
        assert!(
            p50_cpu > 0.18_f32 && p50_cpu < 0.27_f32,
            "CPU oracle: het p50 should land in (0.18, 0.27), got {p50_cpu}",
        );
        // std sanity: with two well-separated ramped clusters (lower
        // 0..0.231, upper 0.424..0.491), the variance is dominated
        // by the inter-cluster gap. Jitter perturbs this slightly;
        // allow ±5e-3 between the GPU result and the CPU oracle.
        let std_tol = 5e-3_f32;
        assert!(
            (std_gpu - std_cpu).abs() < std_tol,
            "het: std_gpu={std_gpu} vs cpu_oracle={std_cpu} (diff={:.3e}, tol={std_tol})",
            (std_gpu - std_cpu).abs(),
        );
        // CPU oracle std lands roughly in the 0.12–0.20 band given
        // the bimodal ramped layout.
        assert!(
            std_cpu > 0.12_f32 && std_cpu < 0.20_f32,
            "CPU oracle: het std should land in (0.12, 0.20), got {std_cpu}",
        );
    }

    /// Test 4 — empty-batch degenerate guard.
    ///
    /// `B = 0` ⇒ the kernel's degenerate branch writes `[0, 0]` and
    /// returns without launching the full pipeline. This is the
    /// trainer's pre-warmup path (no aux logits yet).
    #[test]
    #[ignore = "requires GPU"]
    fn empty_batch_writes_zero_stats() {
        // We still need a non-zero allocation for the device pointer to
        // be valid; the kernel will not read it because B == 0.
        let aux_logits = vec![0.0_f32; AUX_K_CLASSES];
        let (p50, std) = run_kernel(0, &aux_logits);
        assert!(p50.abs() < 1e-6, "empty: p50 expected 0, got {p50}");
        assert!(std.abs() < 1e-6, "empty: std expected 0, got {std}");
    }
}