FE-1340: Add an experimental WebGPU compute backend for experiments

.changeset/webgpu-experiment-backend.md

@@ -0,0 +1,6 @@
+---
+"@hashintel/petrinaut-core": patch
+"@hashintel/petrinaut": patch
+---
+
+Add an experimental WebGPU compute backend for experiments, chosen per experiment behind a user setting. It runs the net's lowered HIR on the device, declines nets it cannot run so they fall back to the CPU, and agrees with the CPU in distribution rather than seed for seed. A Compilation panel, also behind a setting, shows what the compiler made of each condition, kernel and equation.

libs/@hashintel/petrinaut-core/package.json

@@ -27,7 +27,7 @@
       "import": "./dist/index.js"
     },
     "./experiments": {
-      "types": "./dist/experiments.d.d.ts",
+      "types": "./dist/experiments.d.ts",
       "import": "./dist/experiments.js"
     },
     "./ai": {
@@ -70,13 +70,17 @@
       "types": "./dist/workers/simulation.d.ts",
       "import": "./dist/workers/simulation.js"
     },
+    "./webgpu": {
+      "types": "./dist/webgpu.d.ts",
+      "import": "./dist/webgpu.js"
+    },
     "./package.json": "./package.json"
   },
   "publishConfig": {
     "access": "public"
   },
   "scripts": {
-    "build": "vite build",
+    "build": "vite build && node scripts/check-browser-safe-entries.mjs",
     "fix:eslint": "oxlint --fix --type-aware --report-unused-disable-directives-severity=error .",
     "lint:eslint": "oxlint --type-aware --report-unused-disable-directives-severity=error .",
     "lint:tsc": "tsgo --noEmit",
@@ -97,6 +101,7 @@
     "@types/js-yaml": "^4",
     "@types/node": "22.18.13",
     "@typescript/native-preview": "7.0.0-dev.20260511.1",
+    "@webgpu/types": "0.1.71",
     "oxlint": "1.63.0",
     "oxlint-tsgolint": "0.22.1",
     "rolldown": "1.2.6",

libs/@hashintel/petrinaut-core/scripts/check-browser-safe-entries.mjs

@@ -0,0 +1,141 @@
+/**
+ * Fails if a browser-facing entry point reaches Node-only code.
+ *
+ * This exists because that mistake broke the frontend build twice in a row, and
+ * neither `lint:tsc`, `lint:eslint` nor `test:unit` can see it — the import is
+ * perfectly valid TypeScript. Only bundling the consumer catches it, and a full
+ * `@apps/hash-frontend` build takes ~9 minutes, which is too slow to run before
+ * every push.
+ *
+ * The failure mode it guards: an entry that transitively imports the HIR
+ * frontend pulls in the TypeScript compiler, whose `require("module")` webpack
+ * cannot resolve for the browser, so the consuming app fails with
+ * `Module not found: Can't resolve 'module'`.
+ *
+ *   node scripts/check-browser-safe-entries.mjs
+ *
+ * Run after `yarn build`, since it inspects `dist`.
+ */
+import { readFileSync } from "node:fs";
+import { dirname, resolve } from "node:path";
+
+const packageRoot = resolve(dirname(new URL(import.meta.url).pathname), "..");
+
+/**
+ * Entries a browser bundle may import, and what they must never reach.
+ *
+ * `hir` and `compiled-model` are deliberately absent: they are Node/worker
+ * entries and are *expected* to bundle the compiler.
+ */
+const BROWSER_SAFE_ENTRIES = ["index.js", "webgpu.js", "hir-runtime.js"];
+
+/** Bare specifiers that mean "this cannot run in a browser". */
+const NODE_ONLY = new Set([
+  "module",
+  "fs",
+  "fs/promises",
+  "path",
+  "os",
+  "crypto",
+  "child_process",
+  "worker_threads",
+  "url",
+  "util",
+  "typescript",
+]);
+
+/**
+ * Follows relative imports from `entry` and returns every bare specifier
+ * reached, mapped to the first file that imported it.
+ *
+ * @param {string} entry
+ * @returns {Map<string, string>}
+ */
+function collectBareImports(entry) {
+  /** @type {Set<string>} */
+  const seen = new Set();
+  /** @type {Map<string, string>} */
+  const bare = new Map();
+  /** @type {string[]} */
+  const queue = [entry];
+
+  for (let file = queue.pop(); file !== undefined; file = queue.pop()) {
+    if (seen.has(file)) {
+      continue;
+    }
+    seen.add(file);
+
+    let source;
+    try {
+      source = readFileSync(file, "utf8");
+    } catch {
+      continue;
+    }
+
+    // Static and dynamic import specifiers, plus CJS requires, as they appear in
+    // the built output.
+    /** @type {string[]} */
+    const specifiers = [];
+    for (const pattern of [
+      /(?:from|import)\s*\(?\s*["']([^"']+)["']/gu,
+      /require\(\s*["']([^"']+)["']\s*\)/gu,
+    ]) {
+      for (const [, specifier] of source.matchAll(pattern)) {
+        specifiers.push(specifier);
+      }
+    }
+
+    for (const specifier of specifiers) {
+      if (specifier.startsWith(".")) {
+        const target = resolve(dirname(file), specifier);
+        queue.push(target, `${target}.js`);
+      } else {
+        const bareName = specifier.startsWith("node:")
+          ? specifier.slice("node:".length)
+          : specifier;
+        if (!bare.has(bareName)) {
+          bare.set(bareName, file);
+        }
+      }
+    }
+  }
+
+  return bare;
+}
+
+let failed = false;
+
+for (const entry of BROWSER_SAFE_ENTRIES) {
+  const entryPath = resolve(packageRoot, "dist", entry);
+  const bare = collectBareImports(entryPath);
+  const offenders = [...bare].filter(([name]) => NODE_ONLY.has(name));
+
+  if (offenders.length === 0) {
+    const external = [...bare.keys()];
+    process.stdout.write(
+      `  ok   ${entry}${external.length > 0 ? ` (external: ${external.join(", ")})` : " (no external imports)"}\n`,
+    );
+    continue;
+  }
+
+  failed = true;
+  process.stdout.write(`  FAIL ${entry}\n`);
+  for (const [name, importer] of offenders) {
+    process.stdout.write(
+      `         reaches "${name}" via ${importer.replace(`${packageRoot}/`, "")}\n`,
+    );
+  }
+}
+
+if (failed) {
+  process.stdout.write(
+    "\nA browser-facing entry reaches Node-only code. Consuming apps will fail to\n" +
+      "bundle with `Module not found`. Move the Node-only dependency behind a\n" +
+      "separate entry point rather than guarding the import site.\n",
+  );
+  process.exit(1);
+}
+
+process.stdout.write(
+  "\nAll browser-facing entries are free of Node-only imports.\n",
+);

libs/@hashintel/petrinaut-core/src/ai.ts

@@ -97,6 +97,7 @@ export const petrinautDocNames = [
   "actual-mode",
   "ai-assistant",
   "visual-settings",
+  "compilation-output",
   "examples",
 ] as const;
 
@@ -124,7 +125,9 @@ export const petrinautDocSummaries: Record<PetrinautDocName, string> = {
   "ai-assistant":
     "In-app AI assistant: opening the panel, one text and Voice mode transcript/composer, waveform start, inline Voice state and provenance, typed handoff, consent/recovery, prompt chips, tool cards, read-only/simulate-mode rules, host configuration.",
   "visual-settings":
-    "Animations, keep-panels-mounted, minimap, snap-to-grid, compact vs classic nodes, partial selection, tree view, arc rendering style.",
+    "Animations, keep-panels-mounted, minimap, snap-to-grid, compact vs classic nodes, partial selection, tree view, arc rendering style, compute backend, compilation output.",
+  "compilation-output":
+    "The Compilation bottom-panel tab: enabling it, the GPU verdict line, structural blockers, shader emission failures, per-item GPU/CPU/untested/no-HIR/unused status, and HIR node counts.",
   examples:
     "Walkthroughs of the built-in examples and the scenarios/metrics each ships with: SIR, Supply Chain, Deployment Pipeline, Production Machines, Satellites in Orbit, Probabilistic Satellites Launcher.",
 };

libs/@hashintel/petrinaut-core/src/hir.ts

@@ -21,6 +21,7 @@ export {
 } from "./hir/analyze";
 export {
   compileHirArtifacts,
+  type CompileHirArtifactsOptions,
   type HirCompileFailure,
   type HirCompileResult,
 } from "./hir/compile";

libs/@hashintel/petrinaut-core/src/hir/compile.ts

@@ -101,9 +101,23 @@ function lowerAndCheck(
  * not scalarize to a buffer program are reported in `failures` (mirrored by
  * the LSP as error diagnostics); such items cannot simulate.
  */
+export type CompileHirArtifactsOptions = {
+  /**
+   * Also carry the lowered HIR tree on each artifact.
+   *
+   * Off by default: the tree roughly triples artifact size (measured +197% on
+   * the supply-chain example), and artifacts are posted to every Monte Carlo
+   * shard worker, none of which need it. Only the WebGPU backend does, because
+   * it generates a shader from the HIR and cannot lower the net itself — that
+   * would pull the TypeScript frontend into the browser bundle.
+   */
+  includeHir?: boolean;
+};
+
 export function compileHirArtifacts(
   sdcpn: SDCPN,
   extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS,
+  options: CompileHirArtifactsOptions = {},
 ): HirCompileResult {
   const sanitized = sanitizeSDCPNForExtensions(sdcpn, extensions);
   // The four sub-records are keyed by item ids from the net definition, so
@@ -168,7 +182,10 @@ export function compileHirArtifacts(
         });
         continue;
       }
-      artifacts.dynamics[de.id] = { source };
+      artifacts.dynamics[de.id] = {
+        source,
+        ...(options.includeHir ? { hir: item.fn } : {}),
+      };
     }
 
     for (const transition of transitions) {
@@ -203,6 +220,7 @@ export function compileHirArtifacts(
             });
           } else {
             artifacts.lambdas[transition.id] = {
+              ...(options.includeHir ? { hir: item.fn } : {}),
               source: program.source,
               inputSlotCount: program.inputSlotCount,
             };
@@ -241,6 +259,7 @@ export function compileHirArtifacts(
               source: program.source,
               inputSlotCount: program.inputSlotCount,
               outputByteCount: program.outputByteCount,
+              ...(options.includeHir ? { hir: item.fn } : {}),
             };
           }
         }

libs/@hashintel/petrinaut-core/src/hir/instantiate.ts

@@ -10,6 +10,7 @@
 import { cloneUserKeyedRecord } from "../validation/record-keys";
 
 import type { RuntimeDistribution } from "../simulation/authoring/user-code/distribution";
+import type { HirFunction } from "./hir";
 
 export type HirParameterValues = Record<string, number | boolean>;
 
@@ -88,17 +89,31 @@ export type HirLambdaArtifact = {
   source: string;
   /** Expected `indices.length` — engine-side sanity check. */
   inputSlotCount: number;
+  /**
+   * The lowered HIR the program was emitted from.
+   *
+   * Carried so a second backend can compile the same code without re-running the
+   * TypeScript frontend. That matters for the WebGPU backend specifically: it
+   * runs in the browser, and lowering would pull the TypeScript compiler (and
+   * its Node builtins) into the browser bundle. HIR is a JSON-serializable tree,
+   * so it crosses the worker boundary with the rest of the artifact.
+   */
+  hir?: HirFunction;
 };
 
 export type HirKernelArtifact = {
   source: string;
   inputSlotCount: number;
   /** Expected staging byte length — engine-side sanity check. */
   outputByteCount: number;
+  /** The lowered HIR the program was emitted from — see `HirLambdaArtifact.hir`. */
+  hir?: HirFunction;
 };
 
 export type HirDynamicsArtifact = {
   source: string;
+  /** The lowered HIR the program was emitted from — see `HirLambdaArtifact.hir`. */
+  hir?: HirFunction;
 };
 
 export type HirMetricArtifact = {

libs/@hashintel/petrinaut-core/src/index.ts

@@ -223,6 +223,9 @@ export {
   createWorkerTransport,
   deriveRunSeed,
 } from "./simulation";
+// Dependency-free WebGPU capability check. The backend itself lives behind the
+// `./webgpu` entry point, which bundles the HIR frontend and must not reach UI.
+export { isWebGpuAvailable } from "./webgpu/support";
 export type {
   BackpressureConfig,
   CreateMonteCarloExperimentConfig,
@@ -313,6 +316,7 @@ export type {
 // --- HIR (type-only from the main entry; the compiler itself stays in the
 // LSP worker, runtime instantiation in ./hir-runtime) ---
 export type {
+  CompileHirArtifactsOptions,
   HirArtifacts,
   HirCompileFailure,
   HirCompileResult,

libs/@hashintel/petrinaut-core/src/lsp/language-client.ts

@@ -9,6 +9,7 @@ import {
 import type { PetrinautExtensionSettings } from "../extensions";
 // Type-only: must not pull the compiler (`typescript`) into client bundles.
 import type { HirCompileResult, ScenarioHir } from "../hir";
+import type { CompileHirArtifactsOptions } from "../hir/compile";
 import type { AdHocSynthesisContext } from "../simulation/authoring/scenario/ad-hoc/ad-hoc-scenario";
 import type { ReadableStore } from "../store";
 import type { Scenario, SDCPN } from "../types/sdcpn";
@@ -99,6 +100,12 @@ export interface LanguageClient {
     this: void,
     sdcpn: SDCPN,
     extensions?: PetrinautExtensionSettings,
+    /**
+     * Pass `{ includeHir: true }` when the caller needs the HIR tree — only the
+     * WebGPU backend does. It roughly triples artifact size, so it is off by
+     * default.
+     */
+    options?: CompileHirArtifactsOptions,
   ): Promise<HirCompileResult>;
 
   /**
@@ -363,10 +370,11 @@ export function createLanguageClient(
         position,
       });
     },
-    requestHirArtifacts(sdcpn, extensions) {
+    requestHirArtifacts(sdcpn, extensions, options) {
       return sendRequest<HirCompileResult>("sdcpn/compileHirArtifacts", {
         sdcpn,
         extensions,
+        options,
       });
     },
     requestScenarioHir(scenario, adHocContext) {

libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts

@@ -566,7 +566,11 @@ workerRuntime.onMessage((data) => {
         const { id } = data;
         respond(
           id,
-          compileHirArtifacts(data.params.sdcpn, data.params.extensions),
+          compileHirArtifacts(
+            data.params.sdcpn,
+            data.params.extensions,
+            data.params.options,
+          ),
         );
         break;
       }

libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts

@@ -13,6 +13,7 @@
  * upstream package directly.
  */
 import type { PetrinautExtensionSettings } from "../../extensions";
+import type { CompileHirArtifactsOptions } from "../../hir/compile";
 import type {
   AdHocScenarioState,
   AdHocSynthesisContext,
@@ -172,6 +173,7 @@ type ClientRequest =
       params: {
         sdcpn: SDCPN;
         extensions?: PetrinautExtensionSettings;
+        options?: CompileHirArtifactsOptions;
       };
     }
   | {