open-standards/wai/wai-rs/tests/neural_npu.rs

//! Neural dispatch in the native sink (extensions/neural-dispatch.md): a neural
//! capability's synthesis on the platform's neural processor through its model runtime,
//! held to the CPU reference at the capability's tier, and the neural processor refused
//! wherever a decode needs `decode-equivalence`.
//!
//! The set and payload come from `tools/wai_bmshj2018_npu_set.py` (CompressAI's
//! bmshj2018-factorized weights, BSD-3-Clause-Clear), named by `WAI_NPU_SET`: a directory
//! holding `decoder.onnx`, `cdfs_q3.json`, `decoder.mlmodel`, `decoder.weight.bin` and
//! `vector.bmshj2018`. Without it the tests that need it print SKIPPED; with
//! `WAI_REQUIRE_NPU_SET=1` they fail instead. `WAI_REQUIRE_NPU_PLACEMENT=1` also fails a
//! host whose model runtime places no operation on the neural processor, where the
//! dispatch reports `dispatch_not_placed` and decodes on the CPU.
//!
//! The integer capability's refusal reads the committed `neural-conformance` corpus and
//! always runs.

#![cfg(feature = "neural_npu")]

use std::path::{Path, PathBuf};

use wai::codecs::DeclaredTier;
use wai::container::{Conditioning, Manifest, ModelRequirement, Prior, Wai};
use wai::neural::npu::{decode_envelope_dispatched, decode_payload_dispatched, npu_plan, Offer};
use wai::neural::{DecodeError, Decoded, ModelRegistry, ParamSource};
use wai::neural_dispatch::{self as nd, Path as UnitPath, Refusal};

const CAP: &str = "wai.neural.bmshj2018";

/// The set's directory, or `None` (SKIPPED) where it is not supplied.
fn set_dir() -> Option<PathBuf> {
    let dir = std::env::var_os("WAI_NPU_SET").map(PathBuf::from);
    let required = std::env::var("WAI_REQUIRE_NPU_SET").as_deref() == Ok("1");
    match dir {
        Some(d) if d.join("decoder.mlmodel").is_file() && d.join("vector.bmshj2018").is_file() => Some(d),
        other => {
            assert!(!required, "WAI_REQUIRE_NPU_SET=1 and WAI_NPU_SET ({other:?}) holds no set");
            eprintln!("SKIPPED: WAI_NPU_SET names no set from tools/wai_bmshj2018_npu_set.py");
            None
        }
    }
}

fn image(d: &Decoded) -> &[u8] {
    match d {
        Decoded::Image(img) => &img.rgb,
        other => panic!("expected an image, got {other:?}"),
    }
}

fn psnr(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let mse = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum::<f64>() / a.len() as f64;
    if mse == 0.0 { f64::INFINITY } else { 10.0 * (255.0f64 * 255.0 / mse).log10() }
}

/// §6.4: the reconstruction bound this implementation holds a float path to, against the
/// CPU reference, on the set's vector.
const MIN_PSNR_DB: f64 = 40.0;

/// The gate: bmshj2018's synthesis runs on the neural processor where the runtime's plan
/// places it there, the report says so with the plan, and the decode agrees with the CPU
/// reference at the capability's tier (`none`): the entropy stage's symbols are the
/// reference's, and the reconstruction is within §6.4's bound.
#[test]
fn bmshj2018_decodes_on_the_neural_processor_and_agrees_with_the_reference() {
    let Some(dir) = set_dir() else { return };
    let payload = std::fs::read(dir.join("vector.bmshj2018")).unwrap();
    let params = ParamSource::Unpinned(dir.join("decoder.onnx"));

    let (reference, cpu) = decode_payload_dispatched(CAP, &payload, &params, None, &Offer::cpu()).expect("the CPU decodes");
    let ran = cpu.ran.clone().expect("the CPU path ran");
    assert_eq!((ran.path, ran.tier, cpu.needed), (UnitPath::CPU_ONNX_RUNTIME, DeclaredTier::None, DeclaredTier::None));
    assert!(cpu.refused.is_empty());

    let (plan, operations) = npu_plan(&params).expect("the set holds a program").expect("the runtime loads it");
    println!("plan: {plan:?}");
    for (op, device) in &operations {
        println!("  {op}: {device}");
    }
    assert_eq!(plan.total() as usize, operations.len());

    let (decoded, report) = decode_payload_dispatched(CAP, &payload, &params, None, &Offer::native()).expect("dispatch decodes");
    let decoded = decoded.output;
    println!("{}", serde_json::to_string_pretty(&report.to_json()).unwrap());
    let ran = report.ran.clone().expect("a path ran");
    if plan.npu == 0 {
        assert_eq!(report.refused, vec![(UnitPath::NPU_MODEL_RUNTIME, Refusal::NotPlaced)], "a plan with nothing on the processor is not a run there");
        assert_eq!(ran.path, UnitPath::CPU_ONNX_RUNTIME);
        assert_eq!(image(&decoded), image(&reference.output), "the CPU path gave the reference");
        assert!(std::env::var("WAI_REQUIRE_NPU_PLACEMENT").as_deref() != Ok("1"), "WAI_REQUIRE_NPU_PLACEMENT=1, and the runtime placed no operation on the neural processor: {plan:?}");
        eprintln!("SKIPPED: the runtime placed no operation on the neural processor here ({plan:?}); the dispatch said so and decoded on the CPU");
        return;
    }
    assert_eq!(ran.path, UnitPath::NPU_MODEL_RUNTIME);
    assert_eq!(ran.placement, Some(plan));
    assert_eq!(plan.npu, plan.total(), "every operation placed on the neural processor: {operations:?}");
    assert_eq!(ran.tier, DeclaredTier::None, "a float path holds the registered tier and no more");
    assert!(report.refused.is_empty());
    let j = report.to_json();
    assert_eq!((j["ran"]["unit"].as_str(), j["ran"]["backend"].as_str(), j["ran"]["tier"].as_str()), (Some("npu"), Some("model-runtime"), Some("none")));

    // At the tier: the symbols are the entropy stage's, on the CPU, whichever path ran the
    // synthesis; and the reconstruction is within the bound of the reference.
    let cdfs: nd::bmshj2018::Cdfs = serde_json::from_slice(&std::fs::read(dir.join("cdfs_q3.json")).unwrap()).unwrap();
    let latent = nd::bmshj2018::latent(&payload, &cdfs).expect("the entropy stage decodes");
    println!("symbols digest: {} ({} symbols)", latent.symbols_digest(), latent.symbols.len());
    let (npu_rgb, ref_rgb) = (image(&decoded), image(&reference.output));
    let p = psnr(npu_rgb, ref_rgb);
    let max = npu_rgb.iter().zip(ref_rgb).map(|(&a, &b)| (a as i32 - b as i32).abs()).max().unwrap();
    let differ = npu_rgb.iter().zip(ref_rgb).filter(|(a, b)| a != b).count();
    println!("against the CPU reference: {p:.2} dB, max |Δ| {max}, {differ} of {} samples differ", npu_rgb.len());
    assert!(p >= MIN_PSNR_DB, "{p:.2} dB < {MIN_PSNR_DB} dB");
    if let Ok(out) = std::env::var("WAI_NPU_OUT") {
        std::fs::create_dir_all(&out).unwrap();
        std::fs::write(Path::new(&out).join("reference.rgb"), ref_rgb).unwrap();
        std::fs::write(Path::new(&out).join("npu.rgb"), npu_rgb).unwrap();
        std::fs::write(Path::new(&out).join("symbols.blake3"), latent.symbols_digest()).unwrap();
        std::fs::write(Path::new(&out).join("report.json"), serde_json::to_string_pretty(&j).unwrap()).unwrap();
    }
}

/// The neural processor is refused for a decode that needs `decode-equivalence`: a pin
/// that states it, and `intent = replicate`, leave a float capability no path at all.
#[test]
fn a_decode_that_needs_decode_equivalence_refuses_every_float_path() {
    let Some(dir) = set_dir() else { return };
    let payload = std::fs::read(dir.join("vector.bmshj2018")).unwrap();
    let params = ParamSource::Unpinned(dir.join("decoder.onnx"));
    let refused = |r: Result<_, DecodeError>| match r {
        Err(DecodeError::Dispatch(j)) => serde_json::from_str::<serde_json::Value>(&j).unwrap(),
        Err(e) => panic!("expected a dispatch refusal, got {e}"),
        Ok(_) => panic!("expected a dispatch refusal, and it decoded"),
    };
    let j = refused(decode_payload_dispatched(CAP, &payload, &params, Some("decode-equivalence"), &Offer::native()));
    assert_eq!(j["needed"], "decode-equivalence");
    assert_eq!(j["ran"], serde_json::Value::Null);
    let reasons: Vec<(&str, &str)> = j["refused"].as_array().unwrap().iter().map(|r| (r["unit"].as_str().unwrap(), r["reason"].as_str().unwrap())).collect();
    assert_eq!(reasons, vec![("npu", "dispatch_tier_refused"), ("cpu", "dispatch_tier_refused"), ("cpu", "dispatch_arithmetic_mismatch")]);

    let env = Wai::new(
        Manifest::new("image", "replicate", ModelRequirement::new(CAP), Conditioning::new("bmshj2018_rans")),
        payload.clone(),
    );
    let reg = ModelRegistry::new().register(CAP, dir.join("decoder.onnx"));
    let j = refused(decode_envelope_dispatched(&env, &reg, &Offer::native()));
    assert_eq!(j["refused"][0], serde_json::json!({"unit": "npu", "backend": "model-runtime", "arithmetic": "float", "reason": "dispatch_tier_refused"}));

    // `improve` over the same payload decodes, and says which path ran.
    let env = Wai::new(
        Manifest::new("image", "improve", ModelRequirement::new(CAP), Conditioning::new("bmshj2018_rans")),
        payload,
    );
    let (_, report) = decode_envelope_dispatched(&env, &reg, &Offer::native()).expect("improve decodes");
    assert_eq!(report.needed, DeclaredTier::None);
    assert!(report.ran.is_some());
}

/// An integer capability (`decode-equivalence`): the neural processor's float path and ONNX
/// Runtime are refused, and the integer interpreter decodes the corpus vector to its
/// committed hash. Runs everywhere.
#[test]
fn an_integer_capability_refuses_the_neural_processor_and_decodes_on_the_integer_path() {
    let cap = "wai.neural.int_synth";
    let src = Path::new(env!("CARGO_MANIFEST_DIR")).join("../neural-conformance/vectors/gradient");
    let spec = std::fs::read(src.join("int_synth.json")).unwrap();
    let latent = std::fs::read(src.join("latent.bin")).unwrap();
    let env = Wai::new(
        Manifest::new("image", "improve", ModelRequirement::new(cap),
                      Conditioning::new("int_synth").with_prior(Prior::new("gradient", wai::pin::sha256_hex(&spec), "int-synth/i64"))),
        latent,
    );
    let reg = ModelRegistry::new().register(cap, src.join("model"));
    let (d, report) = decode_envelope_dispatched(&env, &reg, &Offer::native()).expect("the integer path decodes");
    let mut h = blake3::Hasher::new();
    h.update(b"wai:int-synth-rgb\x01");
    h.update(match &d.output {
        Decoded::Image(img) => &img.rgb,
        other => panic!("{other:?}"),
    });
    assert_eq!(h.finalize().to_hex().as_str(), "1282839b699ad3e79fc810a92ef4220fec7a57a573bb6c844709af4c1fb3f638");
    assert_eq!(report.needed, DeclaredTier::DecodeEquivalence);
    assert_eq!(report.refused, vec![(UnitPath::NPU_MODEL_RUNTIME, Refusal::TierRefused), (UnitPath::CPU_ONNX_RUNTIME, Refusal::TierRefused)]);
    let ran = report.ran.expect("the integer path ran");
    assert_eq!((ran.path, ran.tier), (UnitPath::CPU_INTEGER, DeclaredTier::DecodeEquivalence));
}