open-standards/wai/wai-web/demo/quantum_webgpu_worker.js

// The fabric's WebGPU member, run in a module worker (quantum_webgpu.html starts
// it). It opens the adapter this worker is offered as a member, probing its
// arithmetic against the CPU's, then:
//
// - certifies the gate kernel of both gatesets on random-v1 and fabric-v1, the
//   four certificates quantum-conformance/fabric/expected.json pins;
// - runs every vector of quantum-conformance/vectors/ and v2/vectors/ on the
//   member to its statevector hash;
// - does both again on the same adapter running the tile implementation
//   (member.tiled()), whose kernel this worker's WGSL compiler also holds to the
//   uniformity rules (messages with tile: true);
// - offers it every file of quantum-conformance/reject/, each of which must be
//   refused with its listed class;
// - offers it a circuit one qubit wider than it holds: refused as too-wide, or as
//   too-many-qubits where it already holds the reference's maximum.
//
// Each result is posted to the page as it arrives; the last message is
// {type: "done"}. Query parameters:
//
//   corpus=<url>             the quantum-conformance directory
//                            (default ../../quantum-conformance/, as served from wai/)
//   limits=webgpu-defaults   the adapter offers only the WebGPU default buffer limits
//                            (128 MiB per storage binding, 256 MiB per buffer), as many
//                            pages are offered; the worker then runs the too-wide check
//                            alone. An adapter that holds the reference's maximum can
//                            never be offered a valid circuit too wide for it; this
//                            shows the refusal on the same device.

import init, { quantum_webgpu_discover } from "./pkg-webgpu/wai_quantum_webgpu.js";

const params = new URL(self.location.href).searchParams;
const CORPUS = new URL(params.get("corpus") || "../../quantum-conformance/", self.location.href);
const LIMITS = params.get("limits");

// The widest state the reference simulates (wai::quantum::MAX_QUBITS).
const MAX_QUBITS = 26;
// The WebGPU default limits, in bytes.
const DEFAULT_STORAGE_BINDING = 134217728;
const DEFAULT_BUFFER = 268435456;

const post = (m) => self.postMessage(m);
const message = (e) => String((e && e.message) || e);

async function fetchBytes(url) {
  const r = await fetch(url);
  if (!r.ok) throw new Error(`${url}: HTTP ${r.status}`);
  return new Uint8Array(await r.arrayBuffer());
}

async function fetchJson(url) {
  const r = await fetch(url);
  if (!r.ok) throw new Error(`${url}: HTTP ${r.status}`);
  return r.json();
}

// The same container with its contract's n_qubits set to `n`: a valid circuit
// whose ops touch only the qubits they did, now `n` qubits wide. Sections keep
// their order; their offsets are recomputed.
function widen(wqc, n) {
  const view = new DataView(wqc.buffer, wqc.byteOffset, wqc.byteLength);
  const count = view.getUint16(4, true);
  const blob = 6 + 9 * count;
  const sections = [];
  for (let k = 0; k < count; k++) {
    const at = 6 + 9 * k;
    const off = view.getUint32(at + 1, true);
    const len = view.getUint32(at + 5, true);
    sections.push({ kind: view.getUint8(at), bytes: wqc.subarray(blob + off, blob + off + len) });
  }
  const contract = sections.find((s) => s.kind === 1);
  if (!contract) throw new Error("the container has no contract section");
  const text = new TextDecoder().decode(contract.bytes);
  const wider = text.replace(/"n_qubits":\d+/, `"n_qubits":${n}`);
  if (wider === text) throw new Error("the contract names no n_qubits");
  contract.bytes = new TextEncoder().encode(wider);
  const out = new Uint8Array(blob + sections.reduce((t, s) => t + s.bytes.length, 0));
  const w = new DataView(out.buffer);
  out.set(wqc.subarray(0, 6));
  let off = 0;
  sections.forEach((s, k) => {
    const at = 6 + 9 * k;
    w.setUint8(at, s.kind);
    w.setUint32(at + 1, off, true);
    w.setUint32(at + 5, s.bytes.length, true);
    out.set(s.bytes, blob + off);
    off += s.bytes.length;
  });
  return out;
}

// Every adapter this worker is then offered reports the WebGPU default buffer
// limits; the member requests what the adapter reports, so its device is granted
// those.
function offerDefaultLimits() {
  const gpu = navigator.gpu;
  const request = gpu.requestAdapter.bind(gpu);
  gpu.requestAdapter = async (options) => {
    const adapter = await request(options);
    if (adapter) {
      Object.defineProperty(adapter.limits, "maxStorageBufferBindingSize", { value: DEFAULT_STORAGE_BINDING });
      Object.defineProperty(adapter.limits, "maxBufferSize", { value: DEFAULT_BUFFER });
    }
    return adapter;
  };
}

// qft5 of circuit2, widened to one qubit more than the member holds.
async function tooWide(member, limits) {
  const n = member.max_qubits() + 1;
  const bytes = widen(await fetchBytes(new URL("v2/vectors/qft5.wqc", CORPUS)), n);
  const want = n > MAX_QUBITS ? "too-many-qubits" : "too-wide";
  try {
    const hash = await member.statevector_hash_hex(bytes);
    post({ type: "wide", limits, n, max_qubits: member.max_qubits(), want, accepted: hash });
  } catch (e) {
    post({ type: "wide", limits, n, max_qubits: member.max_qubits(), want, error: message(e) });
  }
}

// The pinned certificates, then every vector of both trees, on `member`.
async function certifyAndRun(member, expected, tile) {
  for (const c of expected.certificates) {
    const label = `${tile ? "tile " : ""}${c.kernel} ${c.corpus}`;
    try {
      const got = JSON.parse(await member.certify(c.gateset, c.corpus, BigInt(c.seed), c.cases, c.wide));
      post({ type: "certificate", tile, label, gateset: c.gateset, corpus: c.corpus, pin: c.corpus_hex, got });
    } catch (e) {
      post({ type: "certificate", tile, label, gateset: c.gateset, corpus: c.corpus, pin: c.corpus_hex, error: message(e) });
    }
  }
  for (const prefix of ["", "v2/"]) {
    const manifest = await fetchJson(new URL(`${prefix}expected_hashes.json`, CORPUS));
    for (const [file, entry] of Object.entries(manifest)) {
      const label = prefix + file;
      try {
        const bytes = await fetchBytes(new URL(`${prefix}vectors/${file}`, CORPUS));
        post({ type: "vector", tile, label, want: entry.statevector_hash, got: await member.statevector_hash_hex(bytes) });
      } catch (e) {
        post({ type: "vector", tile, label, want: entry.statevector_hash, error: message(e) });
      }
    }
  }
}

async function run() {
  if (!navigator.gpu) {
    post({ type: "no-adapter", why: "navigator.gpu is absent in this worker" });
    return;
  }
  if (LIMITS === "webgpu-defaults") offerDefaultLimits();
  else if (LIMITS) throw new Error(`limits=${LIMITS}: the one value is webgpu-defaults`);
  await init();

  let member;
  try {
    member = await quantum_webgpu_discover();
  } catch (e) {
    const why = message(e);
    post(why.startsWith("no adapter: ") ? { type: "no-adapter", why: why.slice("no adapter: ".length) } : { type: "refused", why });
    return;
  }
  post({
    type: "member",
    limits: LIMITS || "offered",
    id: member.member_id(),
    key: member.key(),
    class: member.class(),
    detail: member.detail(),
    max_qubits: member.max_qubits(),
    implementation: member.implementation(),
  });
  if (LIMITS) {
    await tooWide(member, LIMITS);
    return;
  }

  // The four certificates and every vector of both trees, on the per-gate
  // member and then on the tile member.
  const expected = await fetchJson(new URL("fabric/expected.json", CORPUS));
  await certifyAndRun(member, expected, false);
  let tiled;
  try {
    tiled = await member.tiled();
    post({ type: "tiled", id: tiled.member_id(), implementation: tiled.implementation() });
  } catch (e) {
    post({ type: "tiled", error: message(e) });
  }
  if (tiled) await certifyAndRun(tiled, expected, true);

  // Every reject file.
  const listing = await fetchJson(new URL("reject/reject.json", CORPUS));
  for (const [file, kind] of Object.entries(listing)) {
    const label = "reject/" + file;
    try {
      const bytes = await fetchBytes(new URL(`reject/${file}`, CORPUS));
      post({ type: "reject", label, want: kind, accepted: await member.statevector_hash_hex(bytes) });
    } catch (e) {
      post({ type: "reject", label, want: kind, error: message(e) });
    }
  }

  await tooWide(member, "offered");
}

run()
  .catch((e) => post({ type: "error", error: message(e) }))
  .finally(() => post({ type: "done" }));