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

//! Phase 3D-C kernel-level invariants for `rl_pi_grad_blend`.
//!
//! Three GPU-oracle invariants on the PPO surrogate × Q-distill blend
//! operator (`crates/ml-alpha/cuda/rl_pi_grad_blend.cu`):
//!
//!   1. `blend_disabled_zeros_buffer` — at `RL_PPO_SURROGATE_ENABLED_INDEX`
//!      ≤ 0.5, the kernel zeros the buffer regardless of weight (legacy
//!      zero-then-Q-distill path, bit-equal to pre-3D behaviour).
//!   2. `blend_enabled_scales_by_weight` — at `ENABLED > 0.5`, the kernel
//!      multiplies each entry by `RL_PPO_SURROGATE_WEIGHT_INDEX`. Linear
//!      relationship: `out[i] = w · in[i]`.
//!   3. `blend_weight_zero_is_equivalent_to_disabled` — with ENABLED=1 and
//!      WEIGHT=0, every element becomes 0, matching the disabled path
//!      (relevant for ablation runs that want PPO=off via weight rather
//!      than the gate).
//!
//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical
//! (zero output, linear scaling, weight=0 → zero).
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: ISV bus uses
//! `MappedF32Buffer`; the grad buffer is a regular `CudaSlice` written
//! via the public `write_slice_f32_d_pub` helper.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test phase_3d_blend_kernel_invariants \
//!       --release -- --ignored --nocapture`

use anyhow::{Context, Result};
use cudarc::driver::PushKernelArg;
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_alpha::rl::common::N_ACTIONS;
use ml_alpha::rl::isv_slots::{
    RL_PPO_SURROGATE_ENABLED_INDEX,
    RL_PPO_SURROGATE_WEIGHT_INDEX,
    RL_SLOTS_END,
};
use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub};
use ml_core::device::MlDevice;

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

const B_SIZE: usize = 8;

fn build_device() -> Option<MlDevice> {
    match MlDevice::cuda(0) {
        Ok(d) => Some(d),
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            None
        }
    }
}

/// Launches `rl_pi_grad_blend` on the given (already-uploaded) buffer
/// using the supplied (enabled, weight) ISV values. Returns the host
/// copy of the buffer post-launch.
fn run_blend(
    dev: &MlDevice,
    input: &[f32],
    enabled: f32,
    weight: f32,
) -> Result<Vec<f32>> {
    let stream = dev.cuda_stream().context("cuda_stream")?.clone();
    let ctx = dev.cuda_context().context("cuda_context")?;
    let module = ctx
        .load_cubin(RL_PI_GRAD_BLEND_CUBIN.to_vec())
        .context("load rl_pi_grad_blend cubin")?;
    let func = module
        .load_function("rl_pi_grad_blend")
        .context("load rl_pi_grad_blend fn")?;

    // Allocate input buffer and upload (no raw HtoD per
    // `feedback_no_htod_htoh_only_mapped_pinned`).
    let mut buf = stream.alloc_zeros::<f32>(input.len())?;
    write_slice_f32_d_pub(&stream, input, &mut buf)?;

    // Mapped-pinned ISV buffer (matches production layout).
    let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
        .map_err(|e| anyhow::anyhow!("alloc isv: {e}"))?;
    isv.write_record(RL_PPO_SURROGATE_ENABLED_INDEX, enabled);
    isv.write_record(RL_PPO_SURROGATE_WEIGHT_INDEX, weight);

    let b_i = B_SIZE as i32;
    let n_a = N_ACTIONS as i32;
    let total = (B_SIZE * N_ACTIONS) as u32;
    let block: u32 = 128;
    let grid: u32 = (total + block - 1) / block;

    let isv_dev_ptr = isv.dev_ptr;
    unsafe {
        stream
            .launch_builder(&func)
            .arg(&mut buf)
            .arg(&isv_dev_ptr)
            .arg(&b_i)
            .arg(&n_a)
            .launch(cudarc::driver::LaunchConfig {
                grid_dim: (grid, 1, 1),
                block_dim: (block, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| anyhow::anyhow!("rl_pi_grad_blend launch: {:?}", e))?;
    }
    stream.synchronize().context("sync")?;
    read_slice_d_pub(&stream, &buf, input.len())
}

/// Synthetic input: deterministic per-index ramp so each invariant has
/// a non-trivial signal to verify against.
fn ramp_input() -> Vec<f32> {
    (0..B_SIZE * N_ACTIONS)
        .map(|i| (i as f32) * 0.1 - 5.0) // values in [-5, ~3.7]
        .collect()
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn blend_disabled_zeros_buffer() -> Result<()> {
    let Some(dev) = build_device() else { return Ok(()); };
    let input = ramp_input();
    // ENABLED=0.0 → kernel must zero buffer regardless of weight.
    let out = run_blend(&dev, &input, 0.0, 1.0)?;
    assert_eq!(out.len(), input.len());
    for (i, &v) in out.iter().enumerate() {
        assert!(
            v.abs() < 1e-9,
            "ENABLED=0 must zero buffer: out[{i}] = {v} (input was {})",
            input[i]
        );
    }
    // Also: ENABLED=0.5 (boundary, `> 0.5f` is FALSE → disabled).
    let out2 = run_blend(&dev, &input, 0.5, 1.0)?;
    for (i, &v) in out2.iter().enumerate() {
        assert!(
            v.abs() < 1e-9,
            "ENABLED=0.5 must zero buffer (strict > test): out[{i}] = {v}"
        );
    }
    eprintln!("PASS — ENABLED ≤ 0.5 zeros buffer (legacy zero-then-distill path equivalence)");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn blend_enabled_scales_by_weight() -> Result<()> {
    let Some(dev) = build_device() else { return Ok(()); };
    let input = ramp_input();
    // ENABLED=1.0 + weight=0.005 (production bootstrap): each entry
    // scaled by 0.005.
    let out = run_blend(&dev, &input, 1.0, 0.005)?;
    for (i, (&v_in, &v_out)) in input.iter().zip(out.iter()).enumerate() {
        let expected = 0.005_f32 * v_in;
        assert!(
            (v_out - expected).abs() < 1e-6,
            "ENABLED=1 weight=0.005: out[{i}] = {v_out}, expected 0.005 * {v_in} = {expected}"
        );
    }
    // Also: ENABLED=1.0 + weight=1.0 (identity scaling) — verifies
    // pure pass-through (the surrogate grad lands verbatim in
    // ss_pi_grad_logits_d, ready for Q-distill +=).
    let out_id = run_blend(&dev, &input, 1.0, 1.0)?;
    for (i, (&v_in, &v_out)) in input.iter().zip(out_id.iter()).enumerate() {
        assert!(
            (v_out - v_in).abs() < 1e-6,
            "ENABLED=1 weight=1.0 (identity): out[{i}] = {v_out}, expected {v_in}"
        );
    }
    eprintln!(
        "PASS — ENABLED=1 scales buffer linearly by weight (verified at w=0.005 and w=1.0)"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn blend_weight_zero_is_equivalent_to_disabled() -> Result<()> {
    let Some(dev) = build_device() else { return Ok(()); };
    let input = ramp_input();
    // ENABLED=1.0 but weight=0.0 → out = 0 * in = 0 everywhere.
    // Should be bit-equal to the ENABLED=0 (disabled) path.
    let out_w0 = run_blend(&dev, &input, 1.0, 0.0)?;
    let out_disabled = run_blend(&dev, &input, 0.0, 0.0)?;
    assert_eq!(out_w0.len(), out_disabled.len());
    for (i, (&a, &b)) in out_w0.iter().zip(out_disabled.iter()).enumerate() {
        assert!(
            (a - b).abs() < 1e-9 && a.abs() < 1e-9,
            "weight=0 vs disabled at [{i}]: w0={a} disabled={b} (both must be 0)"
        );
    }
    eprintln!(
        "PASS — ENABLED=1 weight=0 is bit-equal to ENABLED=0 (ablation path equivalence)"
    );
    Ok(())
}