open-standards/wai/wai-pack-gpu/tests/integer_only.rs

//! packaging §11.6 item 1: a packager runs no step on a backend that computes otherwise than
//! the step's definition, such as in floating point where the step is defined in integers.
//! The packer's kernels (the motion search, SSE and BLAKE3's chunk function) are held here,
//! on naga 30.0.1's IR (the translator wgpu compiles them with), to the first two conditions
//! of determinism-tiers §5, *Claiming the tier on a GPU*, which
//! `wai-ops-gpu/tests/wgsl_rules.rs` checks for the decoders' kernels. No adapter is needed.
//!
//! 1. Each kernel parses, and validates with every validation flag and no capability: so no
//!    `f16`, no 64-bit type, no subgroup or quad operation, and no extension.
//! 2. Every scalar, in the type arena and in the type of every expression, is `bool`, `i32`
//!    or `u32`, and every literal is one of those: no floating-point value exists, and every
//!    conversion is between those three.
//!
//! The kernels checked are the ones the crate compiles (`Sources::shipped`), and they are
//! every `.wgsl` file under `src/wgsl`. Then each planted violation (a floating-point value,
//! a floating-point vector, a floating-point built-in, a bitcast to `f32`, `f16`, a subgroup
//! reduction) is shown to be refused by the rule it breaks.
//!
//! The other conditions of determinism-tiers §5 are a decoder's, and are not checked here: a
//! packager's backend is held to the step's output by the qualifying set it runs on each
//! device before a step runs there (packaging §11.6 item 1, `wai_pack::fabric::qualify`).

use std::path::Path;

use naga::valid::{Capabilities, ModuleInfo, ValidationFlags, Validator};
use naga::{Expression, Literal, Module, Scalar, ScalarKind, TypeInner};
use wai_pack_gpu::Sources;

type Check<T = ()> = Result<T, String>;

/// Rule 1.
fn compile(text: &str) -> Check<(Module, ModuleInfo)> {
    let module = naga::front::wgsl::parse_str(text).map_err(|e| format!("rule 1: parse: {}", e.emit_to_string(text)))?;
    let info = Validator::new(ValidationFlags::all(), Capabilities::empty()).validate(&module).map_err(|e| format!("rule 1: validation: {}", e.emit_to_string(text)))?;
    Ok((module, info))
}

fn scalar_ok(s: Scalar) -> bool {
    s.kind == ScalarKind::Bool || (matches!(s.kind, ScalarKind::Sint | ScalarKind::Uint) && s.width == 4)
}

fn type_ok(t: &TypeInner) -> Check {
    match *t {
        TypeInner::Scalar(s) | TypeInner::Vector { scalar: s, .. } | TypeInner::Atomic(s) | TypeInner::ValuePointer { scalar: s, .. } => {
            if scalar_ok(s) { Ok(()) } else { Err(format!("{:?} of width {}", s.kind, s.width)) }
        }
        // Their element types are in the arena, and checked there.
        TypeInner::Pointer { .. } | TypeInner::Array { .. } | TypeInner::Struct { .. } => Ok(()),
        ref other => Err(format!("{other:?}")),
    }
}

fn literal_ok(l: &Literal) -> Check {
    match l {
        Literal::U32(_) | Literal::I32(_) | Literal::Bool(_) => Ok(()),
        other => Err(format!("the literal {other:?}")),
    }
}

/// Rules 1 and 2 on one kernel's text.
fn integer_only(text: &str) -> Check {
    let (m, info) = compile(text)?;
    for (_, t) in m.types.iter() {
        type_ok(&t.inner).map_err(|e| format!("rule 2: the type {} is not made of bool, i32 and u32: {e}", t.name.clone().unwrap_or_default()))?;
    }
    let mut functions: Vec<(String, &naga::Function, &naga::valid::FunctionInfo)> = m.functions.iter().map(|(h, f)| (f.name.clone().unwrap_or_default(), f, &info[h])).collect();
    for (i, e) in m.entry_points.iter().enumerate() {
        functions.push((e.name.clone(), &e.function, info.get_entry_point(i)));
    }
    for (name, f, fi) in functions {
        for (h, e) in f.expressions.iter() {
            type_ok(fi[h].ty.inner_with(&m.types)).map_err(|err| format!("rule 2: in `{name}`, {e:?} has a type that is not bool, i32 or u32: {err}"))?;
            if let Expression::Literal(l) = e {
                literal_ok(l).map_err(|err| format!("rule 2: in `{name}`: {err}"))?;
            }
        }
    }
    for (_, e) in m.global_expressions.iter() {
        if let Expression::Literal(l) = e {
            literal_ok(l).map_err(|err| format!("rule 2: at module scope: {err}"))?;
        }
    }
    Ok(())
}

fn shipped() -> Vec<(&'static str, String)> {
    let s = Sources::shipped();
    vec![("motion", s.motion), ("sse", s.sse), ("blake3", s.blake3)]
}

#[test]
fn every_kernel_computes_in_integers_alone() {
    // The kernels the crate compiles are every kernel it holds: a `.wgsl` file under
    // src/wgsl that no `Sources` member is would escape the check.
    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/wgsl");
    let mut held: Vec<String> = std::fs::read_dir(&dir)
        .expect("src/wgsl")
        .map(|e| e.expect("a directory entry").path())
        .filter(|p| p.extension().is_some_and(|x| x == "wgsl"))
        .map(|p| std::fs::read_to_string(&p).expect("a kernel's text"))
        .collect();
    held.sort();
    let mut compiled: Vec<String> = shipped().into_iter().map(|(_, t)| t).collect();
    compiled.sort();
    assert_eq!(held.len(), compiled.len(), "src/wgsl holds {} kernels; the crate compiles {}", held.len(), compiled.len());
    assert!(held == compiled, "a kernel under src/wgsl is not the text the crate compiles");
    for (name, text) in shipped() {
        integer_only(&text).unwrap_or_else(|e| panic!("{name}: {e}"));
        println!("pack-gpu integer rule: {name}: rules 1-2 hold");
    }
}

#[test]
fn each_planted_violation_is_refused() {
    let sse = Sources::shipped().sse;
    let anchor = "    OUT[t] = acc;\n}";
    assert_eq!(sse.matches(anchor).count(), 1);
    let plant = |code: &str| sse.replacen(anchor, &format!("{code}\n{anchor}"), 1);
    let cases: Vec<(&str, String, &str)> = vec![
        ("a floating-point value", plant("    let fl = f32(acc) * 0.5;\n    acc = u32(fl);"), "rule 2"),
        ("a floating-point vector", plant("    var v: vec4f;\n    acc = acc + u32(v.x);"), "rule 2"),
        ("a floating-point built-in", plant("    acc = u32(sqrt(f32(acc)));"), "rule 2"),
        ("a bitcast to f32", plant("    acc = u32(bitcast<f32>(acc));"), "rule 2"),
        ("an f16 value", format!("enable f16;\n{}", plant("    let h = f16(acc);\n    acc = u32(h);")), "rule 1"),
        ("a subgroup reduction", format!("enable subgroups;\n{}", plant("    acc = subgroupAdd(acc);")), "rule 1"),
    ];
    for (what, text, rule) in &cases {
        assert_ne!(text, &sse, "{what}: the plant did not change the text");
        match integer_only(text) {
            Ok(()) => panic!("{what}: the rules accepted it"),
            Err(e) => {
                assert!(e.starts_with(rule), "{what}: refused, but by `{e}`, not {rule}");
                println!("pack-gpu integer rule: planted {what}: refused: {}", e.lines().next().unwrap_or(""));
            }
        }
    }
}