open-standards/wai/wai-web/src/quantum_webgpu.rs

//! The fabric's WebGPU member in a page (`wai::quantum_fabric_webgpu`), bound
//! for JS in the separate `quantum_webgpu` bundle.
//!
//! [`quantum_webgpu_discover`] opens the adapter the page is offered as a
//! member and probes its arithmetic against the CPU's in this page. The member
//! then certifies the gate kernel against the CPU reference on a corpus
//! ([`QuantumWebGpu::certify`]) and runs `WQC` containers to their statevector
//! hash ([`QuantumWebGpu::statevector_hash_hex`]).
//!
//! The kernel is the same WGSL text the native member runs, compiled here by
//! the browser's own WGSL compiler, so a bit-exact certificate from a page is
//! evidence from an independent translator. [`QuantumWebGpu::tiled`] gives the
//! same adapter running the tile implementation, whose barrier the browser's
//! compiler holds to WGSL's uniformity rules, which the native translator does
//! not check. The corpus digest covers inputs
//! and outputs only, so a bit-exact member reproduces the CPU reference's
//! pinned digests. Certificates are earned in the page that runs the work.
//!
//! Every method that uses the device returns a promise over a clone of the
//! member's handle, so no borrow of the object is held across an await.

use std::rc::Rc;

use wai::quantum::{from_wqc, Gateset, QuantumError};
use wai::quantum_fabric::{Corpus, Parity};
use wai::quantum_fabric_webgpu::{discover, WebGpuError, WebGpuMember};
use wasm_bindgen::prelude::*;
use wasm_bindgen_futures::future_to_promise;

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// The class word of a refusal: the reference's own class for a circuit it
/// refuses (`malformed`, `invalid`, `contract`, `too-many-qubits`, as
/// `reject.json` lists them), else the member's reason.
fn class(e: &WebGpuError) -> &'static str {
    match e {
        WebGpuError::Circuit(q) => q.kind(),
        WebGpuError::TooWide { .. } => "too-wide",
        WebGpuError::OutOfRange { .. } => "out-of-range",
        WebGpuError::NoAdapter(_) => "no-adapter",
        WebGpuError::Unsuitable { .. } => "unsuitable",
        WebGpuError::Device(_) => "device",
        WebGpuError::Memory(_) => "memory",
        WebGpuError::Timeout => "timeout",
        WebGpuError::Pending => "pending",
        WebGpuError::Internal(_) => "internal",
        WebGpuError::Probe { .. } => "probe",
        _ => "member",
    }
}

/// A promise's rejection: an `Error` whose message is
/// `<stage> (<class>): <reason>`, the form the quantum sink throws.
fn refusal(stage: &str, class: &str, reason: &dyn std::fmt::Display) -> JsValue {
    JsError::new(&format!("{stage} ({class}): {reason}")).into()
}

fn member_refusal(stage: &str, e: &WebGpuError) -> JsValue {
    refusal(stage, class(e), e)
}

fn parse_refusal(e: &QuantumError) -> JsValue {
    refusal("WQC parse", e.kind(), e)
}

fn parity_label(p: &Parity) -> &'static str {
    match p {
        Parity::BitExact => "bit-exact",
        Parity::Within { .. } => "within",
        Parity::Statistical { .. } => "statistical",
        Parity::Uncertified => "uncertified",
    }
}

/// The browser's adapter as a fabric member, its arithmetic probed in this page.
#[wasm_bindgen]
pub struct QuantumWebGpu {
    member: Rc<WebGpuMember>,
}

/// The browser's adapter as a fabric member, opened and probed in this page.
/// Rejects with `no adapter: <why>` when the page is offered none, or with the
/// adapter's refusal (`adapter <key> (<name>) refused: <why>`).
#[wasm_bindgen]
pub async fn quantum_webgpu_discover() -> Result<QuantumWebGpu, JsError> {
    let d = discover().await;
    if let Some(member) = d.found.into_iter().next() {
        return Ok(QuantumWebGpu { member: Rc::new(member) });
    }
    match d.refused.first() {
        Some(r) => Err(JsError::new(&format!("adapter {} ({:?}) refused: {}", r.key, r.adapter, r.error))),
        None => Err(JsError::new("no adapter: the page's WebGPU offered none")),
    }
}

#[wasm_bindgen]
impl QuantumWebGpu {
    /// The member id: `local/gpu-webgpu-<d12>`, or `local/swgpu-webgpu-<d12>`
    /// for a software adapter.
    pub fn member_id(&self) -> String {
        self.member.id().to_string()
    }

    /// The key a runner declares for it.
    pub fn key(&self) -> String {
        self.member.key().to_string()
    }

    /// `"gpu"`, or `"cpu"` for a software adapter.
    pub fn class(&self) -> String {
        self.member.class().label().to_string()
    }

    /// What discovery found: adapter, translator, granted limits. Not signed. A
    /// page's WebGPU may report no adapter description; the member id then rests
    /// on the rest of the identity, and this says so.
    pub fn detail(&self) -> String {
        let detail = self.member.detail();
        if self.member.identity().name.is_empty() {
            format!("{detail}; adapter identity withheld")
        } else {
            detail.to_string()
        }
    }

    /// The widest circuit this member's buffers hold.
    pub fn max_qubits(&self) -> u8 {
        self.member.max_qubits()
    }

    /// The kernel implementation the member runs (`wgsl-lanes32/<16 hex>`, or
    /// `wgsl-lanes32-tile/<16 hex>` for a tile member).
    pub fn implementation(&self) -> String {
        self.member.implementation().to_string()
    }

    /// The same adapter running the tile implementation: the tile kernel
    /// compiled by the page's WebGPU, its arithmetic probed in this page, and
    /// circuits planned in tile runs. Resolves to a `QuantumWebGpu` that
    /// certifies and runs as this one does; rejects with
    /// `tiled (<class>): <reason>`.
    pub fn tiled(&self) -> js_sys::Promise {
        let member = Rc::clone(&self.member);
        future_to_promise(async move {
            let tiled = member.tiled().await.map_err(|e| member_refusal("tiled", &e))?;
            Ok(QuantumWebGpu { member: Rc::new(tiled) }.into())
        })
    }

    /// Certify the gate kernel of `gateset` (1 for `wai.quantum.circuit`, 2 for
    /// `wai.quantum.circuit2`) on this member over a corpus: `"random-v1"`
    /// (`seed`, `cases`; `wide` must be 0) or `"fabric-v1"` (`seed`, `cases`,
    /// `wide`). `seed` is a `BigInt`. Resolves to the certificate as JSON:
    /// `{kernel, implementation, member, generator, parity, cases, corpus_hex,
    /// certificate_hash_hex}`; rejects with `certify (<class>): <reason>`.
    pub fn certify(&self, gateset: u8, corpus: String, seed: u64, cases: u32, wide: u32) -> js_sys::Promise {
        let member = Rc::clone(&self.member);
        future_to_promise(async move {
            let gs = match gateset {
                1 => Gateset::V1,
                2 => Gateset::V2,
                g => return Err(refusal("certify", "argument", &format!("gateset {g}: 1 is wai.quantum.circuit, 2 is wai.quantum.circuit2"))),
            };
            let corpus = match corpus.as_str() {
                "random-v1" if wide == 0 => Corpus::Random { seed, cases: cases as usize },
                "random-v1" => return Err(refusal("certify", "argument", &"random-v1 has no wide circuits; pass wide = 0")),
                "fabric-v1" => Corpus::Fabric { seed, cases: cases as usize, wide: wide as usize },
                other => return Err(refusal("certify", "argument", &format!("corpus {other:?}: random-v1 or fabric-v1"))),
            };
            let cert = member.certify(gs, corpus).await.map_err(|e| member_refusal("certify", &e))?;
            let json = serde_json::json!({
                "kernel": cert.kernel,
                "implementation": cert.implementation,
                "member": cert.member,
                "generator": cert.generator,
                "parity": parity_label(&cert.parity),
                "cases": cert.cases,
                "corpus_hex": hex(&cert.corpus),
                "certificate_hash_hex": hex(&cert.hash()),
            });
            Ok(JsValue::from_str(&json.to_string()))
        })
    }

    /// Parse a `WQC1` or `WQC2` container and run it on the member. Resolves to
    /// the statevector hash in hex, the hash the native simulator gives; rejects
    /// with `<stage> (<class>): <reason>`: `WQC parse` with the class
    /// `reject.json` lists for a container the reference reader refuses, or
    /// `simulate` with the reference's class for a circuit it refuses
    /// (`too-many-qubits`), or the member's (`too-wide` for more qubits than
    /// its buffers hold, `out-of-range` where the reference's result leaves 32
    /// bits, `device`, `memory`, ...).
    pub fn statevector_hash_hex(&self, wqc: Vec<u8>) -> js_sys::Promise {
        let member = Rc::clone(&self.member);
        future_to_promise(async move {
            let (circuit, _measure) = from_wqc(&wqc).map_err(|e| parse_refusal(&e))?;
            drop(wqc);
            let state = member.simulate(&circuit).await.map_err(|e| member_refusal("simulate", &e))?;
            Ok(JsValue::from_str(&hex(&state.statevector_hash())))
        })
    }
}