open-standards/wai/wai-ops-gpu/browser/src/lib.rs

//! The graph vectors of `ops-conformance/` on a browser's WebGPU adapter: the same
//! kernels as natively, compiled by the browser's own WGSL compiler, each decode held to
//! the committed expectation and to the interpreter running in the page (wasm32).
//!
//! [`run_corpus`] takes every vector file as one text, each file after a line
//! `### <name>`, and returns the report: one line per adapter, `browser webgpu: …`,
//! then a line per mismatch or device error (at most 20), or `browser webgpu: no adapter
//! (<why>)`.
//!
//! [`run_semantics`] takes WGSL's integer semantics vectors
//! (`wai/wgsl-conformance/integer.jsonl`) and holds them twice: to the reference
//! (`wai_ops_gpu::semantics::reference`) compiled to wasm32 and run in the page, and to
//! every adapter's WebGPU, the kernel compiled by the browser's own WGSL compiler.
//! [`semantics_reference`] is the first half alone, for an engine with no WebGPU.

use wai_ops::json::{self, Value, unhex};
use wai_ops::vectors::read_tensor;
use wai_ops::{Profile, Sequence, Tensor};
use wai_ops_gpu::{GpuSequence, Width, support};
use wasm_bindgen::prelude::*;

#[derive(Clone, Debug, PartialEq, Eq)]
enum Expect {
    Outputs(Vec<Tensor>),
    Refuse(String),
}

struct Decode {
    inputs: Vec<Tensor>,
    scalars: Vec<i64>,
    expect: Expect,
}

struct Vector {
    id: String,
    graph: Vec<u8>,
    decodes: Vec<Decode>,
}

fn parse(text: &str) -> Result<Vec<Vector>, String> {
    let mut out = Vec::new();
    for line in text.split_terminator('\n') {
        if line.starts_with("### ") || line.is_empty() {
            continue;
        }
        let v = json::parse(line.as_bytes()).map_err(|e| format!("JSON at byte {}", e.at))?;
        let s = |k: &str| v.get(k).and_then(Value::as_str).map(str::to_string).ok_or(format!("no {k}"));
        let id = s("id")?;
        let graph = unhex(&s("graph")?).ok_or("graph hex")?;
        let mut decodes = Vec::new();
        for d in v.get("decodes").and_then(Value::as_arr).ok_or("decodes")? {
            let inputs = d.get("inputs").and_then(Value::as_arr).ok_or("inputs")?.iter().map(read_tensor).collect::<Result<Vec<_>, _>>()?;
            let scalars = d.get("scalars").and_then(Value::as_arr).ok_or("scalars")?.iter().map(|x| x.as_dec_i64().ok_or("scalar")).collect::<Result<Vec<_>, _>>()?;
            let e = d.get("expect").ok_or("expect")?;
            let expect = match e.get("refuse") {
                Some(c) => Expect::Refuse(c.as_str().ok_or("refuse")?.to_string()),
                None => Expect::Outputs(e.get("outputs").and_then(Value::as_arr).ok_or("outputs")?.iter().map(read_tensor).collect::<Result<Vec<_>, _>>()?),
            };
            decodes.push(Decode { inputs, scalars, expect });
        }
        out.push(Vector { id, graph, decodes });
    }
    Ok(out)
}

/// Run every vector on every adapter the page is offered.
#[wasm_bindgen]
pub async fn run_corpus(bundle: String) -> String {
    let vs = match parse(&bundle) {
        Ok(v) => v,
        Err(e) => return format!("browser webgpu: the corpus did not parse: {e}"),
    };
    let d = wai_ops_gpu::discover().await;
    let mut lines = Vec::new();
    for (a, why) in &d.refused {
        lines.push(format!("browser webgpu: adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
    }
    if d.found.is_empty() {
        lines.push(format!("browser webgpu: no adapter ({} refused)", d.refused.len()));
        return lines.join("\n");
    }
    for gpu in &d.found {
        let (mut v32, mut d32, mut v64, mut d64, mut cpu, mut refused, mut refusals) = (0u32, 0u32, 0u32, 0u32, 0u32, 0u32, 0u32);
        let mut bad: Vec<String> = Vec::new();
        let (mut errors, mut mismatches) = (0u32, 0u32);
        for v in &vs {
            let Ok(l) = wai_ops::graph::load(&v.graph, Profile::Acc64) else {
                refused += 1;
                continue;
            };
            if support(&l).is_err() {
                cpu += 1;
                continue;
            }
            let mut widths = vec![Width::Acc64];
            if l.meets(Profile::Acc32) {
                widths.push(Width::Acc32);
            }
            for w in widths {
                if w == Width::Acc32 { v32 += 1 } else { v64 += 1 }
                let mut seq = GpuSequence::new();
                let mut iseq = Sequence::new();
                for (k, dd) in v.decodes.iter().enumerate() {
                    if w == Width::Acc32 { d32 += 1 } else { d64 += 1 }
                    refusals += matches!(dd.expect, Expect::Refuse(_)) as u32;
                    let interp = match wai_ops::interp::run(&l, &dd.inputs, &dd.scalars, &mut iseq) {
                        Ok(t) => Expect::Outputs(t),
                        Err(r) => Expect::Refuse(r.code.as_str().to_string()),
                    };
                    let got = match gpu.run_at(&l, w, &dd.inputs, &dd.scalars, &mut seq).await {
                        Ok(Ok(t)) => Expect::Outputs(t),
                        Ok(Err(r)) => Expect::Refuse(r.code.as_str().to_string()),
                        Err(e) => {
                            errors += 1;
                            if bad.len() < 20 {
                                bad.push(format!("browser webgpu: DEVICE ERROR {} decode {k} at {}: {e}", v.id, w.as_str()));
                            }
                            break;
                        }
                    };
                    if got != dd.expect || interp != dd.expect {
                        mismatches += 1;
                        if bad.len() < 20 {
                            bad.push(format!("browser webgpu: MISMATCH {} decode {k} at {} (the interpreter in the page {} the committed result)", v.id, w.as_str(), if interp == dd.expect { "gives" } else { "does NOT give" }));
                        }
                        break;
                    }
                }
            }
        }
        let a = gpu.adapter();
        lines.push(format!(
            "browser webgpu: key={} {} acc32 vectors={v32} decodes={d32}; acc64 vectors={v64} decodes={d64}; refusals={refusals}; cpu={cpu} load_refused={refused} mismatches={mismatches} errors={errors}",
            a.key(),
            a.describe(),
        ));
        lines.extend(bad);
    }
    lines.join("\n")
}

/// The committed semantics vectors against the reference compiled to wasm32: one line,
/// `wasm reference: cases=N equal`, or the first difference.
#[wasm_bindgen]
pub fn semantics_reference(vectors: String) -> String {
    use wai_ops_gpu::semantics;
    let cases = match semantics::parse_vector_file(&vectors) {
        Ok(c) => c,
        Err(e) => return format!("wasm reference: the vectors did not parse: {e}"),
    };
    let ptr_bits = usize::BITS;
    for (c, want) in &cases {
        let got = semantics::reference(c);
        if got != *want {
            return format!("wasm reference: MISMATCH {} a={:#010x} b={:#010x} c={:#010x} d={:#010x}: the file says {want:#010x}, the reference {got:#010x}", c.op.name(), c.a, c.b, c.c, c.d);
        }
    }
    format!("wasm reference: target=wasm32 usize_bits={ptr_bits} cases={} equal", cases.len())
}

/// [`semantics_reference`], then the committed cases and SplitMix64 cases of every
/// operation on every adapter the page is offered, against the reference. One line per
/// adapter, `browser webgpu semantics: key=…`.
#[wasm_bindgen]
pub async fn run_semantics(vectors: String) -> String {
    use wai_ops_gpu::semantics;
    let mut lines = vec![semantics_reference(vectors.clone())];
    let Ok(committed) = semantics::parse_vector_file(&vectors) else {
        return lines.join("\n");
    };
    let mut cases: Vec<semantics::Case> = committed.iter().map(|(c, _)| *c).collect();
    cases.extend(semantics::random_cases(1024, 0xA417_5E01));
    let expect: Vec<u32> = cases.iter().map(semantics::reference).collect();
    let d = wai_ops_gpu::discover().await;
    for (a, why) in &d.refused {
        lines.push(format!("browser webgpu semantics: adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
    }
    if d.found.is_empty() {
        lines.push(format!("browser webgpu semantics: no adapter ({} refused)", d.refused.len()));
        return lines.join("\n");
    }
    for g in &d.found {
        let got = match semantics::run(g, &cases).await {
            Ok(v) => v,
            Err(e) => {
                lines.push(format!("browser webgpu semantics: DEVICE ERROR key={}: {e}", g.adapter().key()));
                continue;
            }
        };
        let o = semantics::compare(&cases, &expect, &got);
        for x in o.guarded_differ.iter().take(20) {
            lines.push(format!("browser webgpu semantics: MISMATCH {x}"));
        }
        let mut shown = std::collections::BTreeSet::new();
        for x in &o.unguarded_differ {
            if shown.insert(x.case.op) {
                lines.push(format!("browser webgpu semantics: unguarded {x}"));
            }
        }
        lines.push(format!(
            "browser webgpu semantics: key={} {} committed={} cases={} guarded={} differ={} unguarded={} unguarded_differ={} unguarded_ops={}",
            g.adapter().key(),
            g.adapter().describe(),
            committed.len(),
            cases.len(),
            o.guarded,
            o.guarded_differ.len(),
            o.unguarded,
            o.unguarded_differ.len(),
            if o.unguarded_ops().is_empty() { "none".to_string() } else { o.unguarded_ops().join(",") },
        ));
    }
    lines.join("\n")
}