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

//! Invariant tests for the GPU log ring: allocator, magic validation,
//! ring wrap, drainer fall-behind detection.
//!
//! Per `feedback_no_cpu_test_fallbacks`: unit tests #1-#6 exercise the
//! HOST-SIDE invariants only (struct layout, slot arithmetic, magic
//! validation, drain dispatch). They write LogRecord structs directly
//! into the mapped-pinned buffer — `#[repr(C)]` makes the host write
//! binary-identical to what a kernel would write.
//!
//! Test #7 is the GPU oracle path: launches the tick kernel via cudarc,
//! confirms the device-side step counter increments.

#![cfg(feature = "kernel-step-trace")]

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use ml_alpha::gpu_log::{
    spawn_drain_task_to_file, LogHeader, LogRecord, LogRing, KID_SMOOTHNESS_CONTROLLER, MAGIC,
    N_RECORDS_PER_STEP, N_RT_PER_KERNEL, N_SLOTS, PAYLOAD_F32, RT_INPUT, RT_OUTPUT, RT_STATE,
};
use ml_alpha::pinned_mem::MappedRecordBuffer;
use ml_core::device::MlDevice;
use tempfile::NamedTempFile;

// ── helpers ──────────────────────────────────────────────────────────

/// Acquire a CUDA device — required even for "unit" tests because
/// `LogRing::alloc()` calls `cuMemHostAlloc` which needs an active
/// context on the calling thread. Per `feedback_no_cpu_test_fallbacks`,
/// the only valid CPU↔GPU memory channel is mapped-pinned (see
/// `feedback_no_htod_htoh_only_mapped_pinned`), so a CUDA-less unit
/// test is not possible for the ring allocator.
fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for gpu_log_ring tests")
}

/// Compute the slot index for a (step, kernel_id, record_type) tuple.
/// Mirrors the kernel-side arithmetic in `gpu_log_ring.cu`.
fn slot_of(step: u32, kernel_id: u8, record_type: u8) -> usize {
    ((step as usize) * N_RECORDS_PER_STEP
        + (kernel_id as usize) * N_RT_PER_KERNEL
        + record_type as usize)
        & (N_SLOTS - 1)
}

/// Build a `LogRecord` with the given fields.
fn make_record(step: u32, kid: u8, rt: u8, payload: &[f32]) -> LogRecord {
    let mut rec = LogRecord {
        header: LogHeader {
            magic: MAGIC,
            step,
            kernel_id: kid,
            record_type: rt,
            payload_words: payload.len().min(PAYLOAD_F32) as u8,
            reserved: 0,
            _padding: 0,
        },
        payload: [0.0; PAYLOAD_F32],
    };
    for (i, &v) in payload.iter().take(PAYLOAD_F32).enumerate() {
        rec.payload[i] = v;
    }
    rec
}

/// Allocate a host-visible step-counter shadow (1 i32).
fn alloc_host_counter() -> Arc<MappedRecordBuffer<i32>> {
    let buf = unsafe { MappedRecordBuffer::<i32>::new(1) }
        .expect("alloc step counter shadow");
    buf.write_record(0, 0);
    Arc::new(buf)
}

// ── unit tests (require CUDA context for mapped-pinned alloc only) ───

#[test]
fn ring_alloc_zeros_magic() {
    let _dev = test_device();
    let ring = LogRing::alloc().expect("alloc ring");
    // Every slot's magic must be 0 (zero-initialised — must not match MAGIC).
    for i in 0..N_SLOTS {
        let rec = ring.buffer.read_record(i);
        assert_eq!(
            rec.header.magic, 0,
            "slot {i} magic should be 0 on fresh alloc, got 0x{:x}",
            rec.header.magic
        );
    }
}

#[test]
fn host_write_round_trip() {
    let _dev = test_device();
    let ring = LogRing::alloc().expect("alloc ring");
    let payload = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
    let rec = make_record(42, KID_SMOOTHNESS_CONTROLLER, RT_INPUT, &payload);

    // Write at the canonical slot for (step=42, KID_SMOOTHNESS_CONTROLLER, RT_INPUT).
    let slot = slot_of(42, KID_SMOOTHNESS_CONTROLLER, RT_INPUT);
    ring.buffer.write_record(slot, rec);

    // Read back — all fields must match.
    let read = ring.buffer.read_record(slot);
    assert_eq!(read.header.magic, MAGIC);
    assert_eq!(read.header.step, 42);
    assert_eq!(read.header.kernel_id, KID_SMOOTHNESS_CONTROLLER);
    assert_eq!(read.header.record_type, RT_INPUT);
    assert_eq!(read.header.payload_words, 10);
    for (i, &expected) in payload.iter().enumerate() {
        assert_eq!(read.payload[i], expected, "payload[{i}] mismatch");
    }
    // Untouched payload slots remain zero.
    for i in payload.len()..PAYLOAD_F32 {
        assert_eq!(read.payload[i], 0.0, "untouched payload[{i}] should be 0");
    }
}

#[test]
fn magic_validation_skips_torn() {
    let _dev = test_device();
    let ring = Arc::new(LogRing::alloc().expect("alloc ring"));
    let counter = alloc_host_counter();
    let stop = Arc::new(AtomicBool::new(false));

    // Slot for step=1, KID_SMOOTHNESS_CONTROLLER, RT_INPUT — valid record.
    let valid_slot = slot_of(1, KID_SMOOTHNESS_CONTROLLER, RT_INPUT);
    // Slot for step=1, KID_SMOOTHNESS_CONTROLLER, RT_STATE — torn (magic=0).
    let torn_slot = slot_of(1, KID_SMOOTHNESS_CONTROLLER, RT_STATE);
    // Slot for step=1, KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT — header.step mismatched.
    let mismatch_slot = slot_of(1, KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT);

    let valid_payload = [100.0_f32; 10];
    ring.buffer.write_record(
        valid_slot,
        make_record(1, KID_SMOOTHNESS_CONTROLLER, RT_INPUT, &valid_payload),
    );

    // Torn record: write with magic=0 so drainer must skip.
    let mut torn = make_record(1, KID_SMOOTHNESS_CONTROLLER, RT_STATE, &[1.0; 10]);
    torn.header.magic = 0;
    ring.buffer.write_record(torn_slot, torn);

    // Mismatch record: claims step=99 but lives in slot for step=1 — drainer must skip.
    let mismatch = make_record(99, KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT, &[2.0; 10]);
    ring.buffer.write_record(mismatch_slot, mismatch);

    counter.write_record(0, 1);

    let tmp = NamedTempFile::new().expect("tempfile");
    let path = tmp.path().to_path_buf();
    let h = spawn_drain_task_to_file(ring.clone(), counter.clone(), stop.clone(), path.clone())
        .expect("spawn drain");
    // Give drainer one tick (>500 ms cadence + a bit of slack).
    std::thread::sleep(Duration::from_millis(700));
    stop.store(true, Ordering::Relaxed);
    // Wait for the drain thread to observe the stop flag and flush.
    let _ = h.join();

    // The valid record should produce exactly one JSONL line; the torn
    // (magic=0) and mismatched-step records must NOT appear.
    let contents = std::fs::read_to_string(&path).expect("read trace file");
    let lines: Vec<&str> = contents.lines().filter(|l| !l.is_empty()).collect();
    assert_eq!(lines.len(), 1, "expected exactly 1 valid record, got {}", lines.len());
    let v: serde_json::Value = serde_json::from_str(lines[0]).expect("parse JSONL");
    assert_eq!(v["step"], 1);
    assert_eq!(v["kid"], KID_SMOOTHNESS_CONTROLLER);
    assert_eq!(v["rt"], RT_INPUT);
    assert_eq!(v["kname"], "smoothness_controller");
    assert_eq!(v["rt_name"], "input");
    // Payload field comes from the smoothness_input schema.
    assert!((v["payload"]["raw_h10"].as_f64().unwrap() - 100.0).abs() < 1e-6);
}

#[test]
fn ring_wrap_modulo_arithmetic() {
    // For step values past N_SLOTS/N_RECORDS_PER_STEP, slot indices wrap.
    let wrap_step = (N_SLOTS / N_RECORDS_PER_STEP) as u32; // = 1024
    // step 0 and step `wrap_step` for (KID=0, RT=0) hash to the same slot.
    assert_eq!(
        slot_of(0, 0, 0),
        slot_of(wrap_step, 0, 0),
        "wrap_step={wrap_step} should collide with step=0 at same (kid, rt)"
    );
    // step 0 and step (wrap_step * 2) collide on the next cycle, etc.
    assert_eq!(slot_of(0, 0, 0), slot_of(wrap_step * 2, 0, 0));
    // Different (kid, rt) at same step never collide (within one cycle).
    assert_ne!(slot_of(0, 0, 0), slot_of(0, 0, 1));
    assert_ne!(slot_of(0, 0, 0), slot_of(0, 1, 0));
    // All slots stay in [0, N_SLOTS).
    for step in 0..(wrap_step * 3) {
        for kid in 0..8u8 {
            for rt in 0..4u8 {
                let s = slot_of(step, kid, rt);
                assert!(
                    s < N_SLOTS,
                    "slot {s} out of bounds for (step={step}, kid={kid}, rt={rt})"
                );
            }
        }
    }
}

#[test]
fn step_gap_detection_emits_warn() {
    let _dev = test_device();
    let ring = Arc::new(LogRing::alloc().expect("alloc ring"));
    let counter = alloc_host_counter();
    let stop = Arc::new(AtomicBool::new(false));

    // Drainer ring window = N_SLOTS / N_RECORDS_PER_STEP = 1024 steps.
    // Jump the counter past the window so drainer must clamp.
    let big_jump = ((N_SLOTS / N_RECORDS_PER_STEP) + 10) as i32;
    counter.write_record(0, big_jump);

    let tmp = NamedTempFile::new().expect("tempfile");
    let h = spawn_drain_task_to_file(
        ring.clone(),
        counter.clone(),
        stop.clone(),
        tmp.path().to_path_buf(),
    )
    .expect("spawn drain");
    std::thread::sleep(Duration::from_millis(700));
    stop.store(true, Ordering::Relaxed);
    let _ = h.join();

    // Pass if no panic — drainer clamped `last_drained_step` to `oldest_safe_step`
    // and emitted a warn. We don't tap tracing here; the absence of panic +
    // the drainer's known structure (warns then clamps) is the pass condition.
}

#[test]
fn decoder_dispatch_unknown_falls_through() {
    let _dev = test_device();
    let ring = Arc::new(LogRing::alloc().expect("alloc ring"));
    let counter = alloc_host_counter();
    let stop = Arc::new(AtomicBool::new(false));

    // Write a valid record with kernel_id=99 (out of declared range — no decoder registered).
    let unknown_kid: u8 = 99;
    let rec = make_record(1, unknown_kid, RT_INPUT, &[7.0; 10]);
    let slot = ((1usize) * N_RECORDS_PER_STEP
        + (unknown_kid as usize) * N_RT_PER_KERNEL
        + RT_INPUT as usize)
        & (N_SLOTS - 1);
    ring.buffer.write_record(slot, rec);
    counter.write_record(0, 1);

    let tmp = NamedTempFile::new().expect("tempfile");
    let h = spawn_drain_task_to_file(
        ring.clone(),
        counter.clone(),
        stop.clone(),
        tmp.path().to_path_buf(),
    )
    .expect("spawn drain");
    std::thread::sleep(Duration::from_millis(700));
    stop.store(true, Ordering::Relaxed);
    let _ = h.join();

    // The drainer's iteration is `for kid in 0..N_KERNELS` — so kid=99 will
    // never be visited by the drainer's slot walk. But the `decode_to_json()`
    // function itself MUST handle unknown kernel_id without panic. Exercise
    // it directly — unknown (kid, rt) pairs fall through to a "values" array.
    let v1 = ml_alpha::gpu_log::decode_to_json(unknown_kid, RT_INPUT, 1, &[7.0; 10]);
    assert_eq!(v1["kname"], "unknown");
    assert!(v1["payload"]["values"].is_array());
    let v2 = ml_alpha::gpu_log::decode_to_json(0, 99, 1, &[7.0; 10]); // unknown record_type
    assert_eq!(v2["rt_name"], "unknown");
}

// ── CUDA integration test #7 ─────────────────────────────────────────

/// Fire the `gpu_log_tick` kernel via cudarc twice; confirm that the
/// device-side step counter increments from 0 → 1 → 2. This validates
/// the cubin includes the `gpu_log_tick` symbol and that the host can
/// drive the GPU log ring's step-counter advance path end-to-end.
///
/// The full kernel-writes → drainer-reads loop is covered by the
/// `perception_overfit` smoke once a kernel wires `log_record()` into
/// the ring (Task 7).
#[test]
#[ignore = "requires CUDA — run via `cargo test --features kernel-step-trace -- --ignored`"]
fn kernel_writes_visible_to_host_smoke() {
    use cudarc::driver::{LaunchConfig, PushKernelArg};

    let dev = test_device();
    let stream = dev.cuda_stream().expect("stream");
    let ctx = dev.cuda_context().expect("ctx");

    let cubin = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin"));
    let module = ctx.load_cubin(cubin.to_vec()).expect("load cubin");
    let tick_fn = module.load_function("gpu_log_tick").expect("load gpu_log_tick");

    // Allocate the step counter, fire tick, confirm device increments to 1.
    let mut step_counter = stream.alloc_zeros::<i32>(1).expect("alloc counter");
    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (1, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&tick_fn);
    launch.arg(&mut step_counter);
    unsafe {
        launch.launch(cfg).expect("launch tick");
    }
    stream.synchronize().expect("sync");

    let host: Vec<i32> = stream.clone_dtoh(&step_counter).expect("dtoh");
    assert_eq!(host[0], 1, "after one tick, step_counter should be 1");

    // Fire again — should increment to 2.
    let mut launch = stream.launch_builder(&tick_fn);
    launch.arg(&mut step_counter);
    unsafe {
        launch.launch(cfg).expect("launch tick 2");
    }
    stream.synchronize().expect("sync");
    let host: Vec<i32> = stream.clone_dtoh(&step_counter).expect("dtoh 2");
    assert_eq!(host[0], 2, "after two ticks, step_counter should be 2");
}