1import { describe, expect, it } from "vitest";
2import { NodeRunner } from "../src/node";
3import {
4 FD_METHOD,
5 LIBRARY_FILES,
6 buildGpuStages,
7 buildRunBundle,
8 parsePlateResult,
9 withEdits,
10 type ProblemSpec,
11 type RunResult,
12} from "../src/index";
13import { requestGpuDevice, type GpuHandle } from "../src/gpu/fdLoop";
14import { maxAbs, relL2, testSpec } from "./helpers";
16/**
17 * Desktop WebGPU for Node through the optional `webgpu` package (Google Dawn):
18 * `npm i --no-save webgpu`. Without the package, or on a machine without a GPU
19 * adapter, the suite is skipped rather than failed, so CI can run it
20 * opportunistically.
21 */
22async function dawn(): Promise<GpuHandle | null> {
23 try {
24 const specifier = "webgpu";
25 const mod = (await import(specifier)) as {
26 create: (flags: string[]) => GPU;
27 globals: Record<string, unknown>;
28 };
29 Object.assign(globalThis, mod.globals);
30 Object.defineProperty(globalThis, "navigator", {
31 value: { gpu: mod.create([]) },
32 configurable: true,
33 writable: true,
34 });
35 return await requestGpuDevice();
36 } catch (e) {
37 console.log(
38 `gpu.test: no WebGPU device under Node (${e instanceof Error ? e.message : e}); skipping`,
39 );
40 return null;
41 }
42}
44const handle = await dawn();
45const cpu = new NodeRunner();
46const gpu = new NodeRunner({ gpu: () => Promise.resolve(handle!) });
48async function both(spec: ProblemSpec, edits: Record<string, string> = {}) {
49 const stages = buildGpuStages(spec, FD_METHOD)!;
50 const g = await gpu.run({
51 bundle: withEdits(buildRunBundle(spec, FD_METHOD, { gpu: true }), edits),
52 gpu: stages,
53 });
54 const c = await cpu.run({ bundle: withEdits(buildRunBundle(spec, FD_METHOD), edits) });
55 if (!g.ok) throw new Error(`gpu run failed: ${g.error}\n${g.output.join("")}`);
56 if (!c.ok) throw new Error(`cpu run failed: ${c.error}\n${c.output.join("")}`);
57 return { g, c, rg: parsePlateResult(g.value), rc: parsePlateResult(c.value) };
58}
60describe.skipIf(!handle)("explicit finite differences on WebGPU", () => {
61 it.each(["simply-supported", "clamped", "free", "pinned-center"] as const)(
62 "reproduces the double-precision loop for %s edges",
63 async (bc) => {
64 const spec = testSpec();
65 spec.bc = bc;
66 const { g, rg, rc } = await both(spec);
67 expect(g.compute?.loop).toBe("webgpu");
68 expect(rg.info.extra.loop).toBe("webgpu");
69 expect(rg.info.stable).toBe(true);
70 expect(rg.info.notes.some((n) => /GPU/.test(n))).toBe(true);
71 expect(rg.t.length).toBe(rc.t.length);
72 expect(rg.frames.t.length).toBe(rc.frames.t.length);
73 const eW = relL2(rc.probe.w, rg.probe.w);
74 const eF = relL2(rc.frames.w.data, rg.frames.w.data);
75 const eVm = relL2(rc.probe.vm, rg.probe.vm);
76 const eA = relL2(rc.audio.p, rg.audio.p);
77 const eEnv = relL2(rc.envelope.vm.data, rg.envelope.vm.data);
78 console.log(
79 `gpu.test ${bc}: rel L2 probe w ${eW.toExponential(2)}, frames ${eF.toExponential(2)}, probe vm ${eVm.toExponential(2)}, audio ${eA.toExponential(2)}, envelope ${eEnv.toExponential(2)}, ${rg.info.nsteps} steps in ${g.compute?.gpuMs?.toFixed(0)} ms on the GPU`,
80 );
81 expect(eW).toBeLessThan(1e-3);
82 expect(eF).toBeLessThan(1e-3);
83 expect(eVm).toBeLessThan(1e-2);
84 expect(eA).toBeLessThan(1e-2);
85 expect(eEnv).toBeLessThan(1e-2);
86 expect(maxAbs(rg.probe.w)).toBeGreaterThan(0);
87 },
88 );
90 it.each(["simply-supported", "clamped", "free", "pinned-center"] as const)(
91 "reproduces the double-precision loop on an ellipse with %s edges",
92 async (bc) => {
93 const spec = testSpec();
94 spec.plate.shape = "ellipse";
95 spec.bc = bc;
96 const { g, rg, rc } = await both(spec);
97 expect(g.compute?.loop).toBe("webgpu");
98 expect(rg.info.stable).toBe(true);
99 expect(rg.frames.w.data.filter((v) => Number.isNaN(v)).length).toBe(
100 rc.frames.w.data.filter((v) => Number.isNaN(v)).length,
101 );
102 const fin = (a: Float64Array) => a.filter(Number.isFinite);
103 const eW = relL2(rc.probe.w, rg.probe.w);
104 const eF = relL2(fin(rc.frames.w.data), fin(rg.frames.w.data));
105 const eVm = relL2(rc.probe.vm, rg.probe.vm);
106 const eA = relL2(rc.audio.p, rg.audio.p);
107 console.log(
108 `gpu.test ellipse ${bc}: rel L2 probe w ${eW.toExponential(2)}, frames ${eF.toExponential(2)}, probe vm ${eVm.toExponential(2)}, audio ${eA.toExponential(2)}`,
109 );
110 expect(eW).toBeLessThan(1e-3);
111 expect(eF).toBeLessThan(1e-3);
112 expect(eVm).toBeLessThan(1e-2);
113 expect(eA).toBeLessThan(1e-2);
114 expect(maxAbs(rg.probe.w)).toBeGreaterThan(0);
115 },
116 );
118 it("stops at the blow-up where the double-precision loop does", async () => {
119 const spec = testSpec();
120 spec.fd.cfl = 1.03;
121 const { rg, rc } = await both(spec);
122 expect(rg.info.stable).toBe(false);
123 expect(rc.info.stable).toBe(false);
124 expect(rg.info.diverged_step).toBeGreaterThan(0);
125 // the unstable mode grows from rounding noise once the strike is over, so the two runs need not agree to the step
126 expect(Math.abs(rg.info.diverged_step - rc.info.diverged_step)).toBeLessThanOrEqual(
127 Math.max(10, 0.1 * rc.info.nsteps),
128 );
129 expect(rg.probe.w.every(Number.isFinite)).toBe(true);
130 expect(rg.frames.w.data.every(Number.isFinite)).toBe(true);
131 expect(rg.frames.t.length).toBe(rc.frames.t.length);
132 expect(rg.info.notes.some((n) => /diverged/.test(n))).toBe(true);
133 });
135 it("falls back to numbl when the loop file is edited, and says so", async () => {
136 const spec = testSpec();
137 const loop = LIBRARY_FILES.find((f) => f.name === "plate_fd_loop.m")!.source;
138 const { g, rg, rc } = await both(spec, {
139 "plate_fd_loop.m": loop + "\n% edited in the code view\n",
140 });
141 expect(g.compute?.loop).toBe("numbl");
142 expect(g.compute?.note).toMatch(/edited/);
143 expect(rg.info.extra.loop).toBe("numbl");
144 expect(relL2(rc.probe.w, rg.probe.w)).toBeLessThan(1e-12);
145 });
146});
148describe("explicit finite differences without a GPU", () => {
149 it("runs the whole bundle in numbl when the GPU is asked for but absent", async () => {
150 const spec = testSpec();
151 const r: RunResult = await cpu.run({
152 bundle: buildRunBundle(spec, FD_METHOD, { gpu: true }),
153 gpu: buildGpuStages(spec, FD_METHOD)!,
154 });
155 expect(r.ok).toBe(true);
156 if (!r.ok) return;
157 expect(r.compute?.loop).toBe("numbl");
158 expect(r.compute?.note).toMatch(/not available/);
159 const res = parsePlateResult(r.value);
160 expect(res.info.extra.loop).toBe("numbl");
161 expect(res.info.stable).toBe(true);
162 });
164 it("writes the staged script and the plain script for the same problem", () => {
165 const spec = testSpec();
166 const plain = buildRunBundle(spec, FD_METHOD).files[0].source;
167 const staged = buildRunBundle(spec, FD_METHOD, { gpu: true }).files[0].source;
168 expect(plain).toMatch(/result = plate_fd\(params\);/);
169 expect(staged).toMatch(/S = plate_fd_setup\(params\);/);
170 expect(staged).toMatch(/R = plate_fd_loop\(S\);/);
171 expect(staged).toMatch(/result = plate_fd_output\(S, R\);/);
172 const stages = buildGpuStages(spec, FD_METHOD)!;
173 expect(stages.loopFile).toBe("plate_fd_loop.m");
174 expect(stages.setup).toMatch(/find\(S\.Km\)/);
175 expect(stages.output).toMatch(/plate_fd_output\(S, R\)/);
176 });
177});