1/**
2 * Correctness of the .m -> WGSL path, against the TypeScript solver.
3 *
4 * src/solver/ is no longer what the app runs, but it is an independent
5 * implementation of the same IMEX scheme, which makes it the oracle here: run
6 * both from the same seeded perturbation, through the same fp32 transforms, and
7 * compare the spectral state.
8 *
9 * The only difference between the two is where the reaction and the IMEX update
10 * happen — f64 on the CPU for the reference, fp32 in generated WGSL for the .m.
11 * The pattern-forming regime amplifies small differences, so this checks a
12 * short run.
13 */
14import { GpuBackend } from '../src/solver/backend.ts';
15import { Simulation, gridForLmax, makeRandn } from '../src/solver/simulation.ts';
16import { models, defaultParams } from '../src/solver/models.ts';
17import { ShtPlan } from '../src/sht/sht.ts';
18import { GpuModel } from '../src/mgpu/model.ts';
19import { mModels, type MModel } from '../src/mgpu/registry.ts';
21type Check = (name: string, ok: boolean, detail: string) => void;
22type Log = (s: string) => void;
24function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
25 let num = 0;
26 let den = 0;
27 for (let i = 0; i < a.length; i++) {
28 const d = a[i] - b[i];
29 num += d * d;
30 den += b[i] * b[i];
31 }
32 return Math.sqrt(num / Math.max(den, 1e-300));
33}
35const LMAX = 31;
36const STEPS = 10;
38/**
39 * Kernels the step of each model should compile to — one per element-wise line
40 * of MATLAB. This is a fusion guard: numbl's lowering emits one statement per
41 * *operator*, and the inline pass folds those back into per-line expression
42 * trees. If that stops happening the results stay correct but every operator
43 * becomes its own dispatch, which is exactly the silent regression to catch.
44 */
45const EXPECTED_KERNELS: Record<string, number> = {
46 schnakenberg: 5,
47 brusselator: 5,
48 allencahn: 2,
49};
51/** One model: compile it, run it, and compare against the reference solver. */
52async function checkModel(
53 device: GPUDevice,
54 m: MModel,
55 check: Check,
56 log: Log,
57): Promise<{ mgpuMs: number; refMs: number } | null> {
58 const spec = models.find((x) => x.key === m.key);
59 if (!spec) {
60 check(`${m.key}: reference model exists`, false, 'no matching ModelSpec');
61 return null;
62 }
63 const params = defaultParams(spec);
64 const { nlat, nphi } = gridForLmax(LMAX, m.pdeg);
65 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
66 const npts = nlat * nphi;
68 const sht = await ShtPlan.create(device, cfg);
69 const gpu = await GpuModel.create({
70 device,
71 sht,
72 cfg,
73 source: m.source,
74 paramNames: m.params.map((p) => p.key),
75 state: m.state,
76 view: m.species,
77 });
78 gpu.setParams(params);
80 const plan = gpu.describe();
81 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
82 const xforms = plan.step.filter(
83 (l) => l.startsWith('synth') || l.startsWith('analys'),
84 ).length;
85 log(
86 ` ${m.key}.m -> ${plan.step.length} ops/step ` +
87 `(${kernels} generated kernels, ${xforms} transforms)`,
88 );
89 const expected = EXPECTED_KERNELS[m.key];
90 check(
91 `${m.key}: element-wise lines fused into one kernel each`,
92 kernels === expected,
93 `${kernels} kernels (expected ${expected})`,
94 );
96 // One randn per grid point, in index order. Rounded to f32 once and fed to
97 // BOTH sides, so the comparison is about the compute path, not the seed.
98 const randn = makeRandn(1);
99 const noise = new Float32Array(npts);
100 for (let i = 0; i < npts; i++) noise[i] = m.seedAmp * randn();
102 gpu.init(noise);
104 // Reference, seeded from the same perturbation by handing the model's own
105 // init the identical sequence.
106 const backend = await GpuBackend.create(device, cfg);
107 const ref = new Simulation(backend, spec, params);
108 {
109 let i = 0;
110 const feed = (): number => noise[i++] / m.seedAmp;
111 const grids = m.state.map(() => new Float64Array(npts));
112 spec.init(params, ref.x, ref.y, ref.z, feed, grids);
113 for (let k = 0; k < m.state.length; k++) {
114 ref.U[k].set(await backend.analys(grids[k]));
115 }
116 }
118 let worstInit = 0;
119 for (let k = 0; k < m.state.length; k++) {
120 worstInit = Math.max(worstInit, relL2(await gpu.read(m.state[k]), ref.U[k]));
121 }
122 check(
123 `${m.key}: init matches reference`,
124 worstInit < 1e-5,
125 `rel L2 ${worstInit.toExponential(2)}`,
126 );
128 for (let s = 0; s < STEPS; s++) await ref.step();
129 gpu.step(STEPS);
131 let worst = 0;
132 let nan = false;
133 for (let k = 0; k < m.state.length; k++) {
134 const got = await gpu.read(m.state[k]);
135 worst = Math.max(worst, relL2(got, ref.U[k]));
136 for (const v of got) if (!Number.isFinite(v)) nan = true;
137 }
138 check(
139 `${m.key}: .m vs reference after ${STEPS} steps`,
140 worst < 2e-3 && !nan,
141 `rel L2 ${worst.toExponential(2)}${nan ? ', NaN!' : ''}`,
142 );
144 // Guard against "both sides computed nothing".
145 const field = await gpu.read(m.species[0]);
146 let peak = 0;
147 for (const v of field) peak = Math.max(peak, Math.abs(v));
148 check(
149 `${m.key}: rendered field is non-trivial`,
150 peak > 1e-4,
151 `max |${m.species[0]}| ${peak.toExponential(2)}`,
152 );
154 // Step rate. The reference maps a staging buffer on every transform, so it
155 // pays four driver round-trips per step; the .m path keeps everything in GPU
156 // buffers and submits once.
157 const TIMED = 50;
158 const t0 = performance.now();
159 gpu.step(TIMED);
160 await device.queue.onSubmittedWorkDone();
161 const mgpuMs = (performance.now() - t0) / TIMED;
163 const t1 = performance.now();
164 for (let s = 0; s < TIMED; s++) await ref.step();
165 const refMs = (performance.now() - t1) / TIMED;
167 gpu.destroy();
168 backend.destroy();
169 sht.destroy();
170 return { mgpuMs, refMs };
171}
173export async function mgpuChecks(
174 device: GPUDevice,
175 check: Check,
176 log: Log,
177): Promise<void> {
178 check(
179 'models: registry populated',
180 mModels.length === models.length,
181 `${mModels.length} .m models`,
182 );
184 for (const m of mModels) {
185 const timing = await checkModel(device, m, check, log);
186 if (!timing) continue;
187 const { mgpuMs, refMs } = timing;
188 log(
189 ` step rate: .m ${mgpuMs.toFixed(2)} ms vs reference ` +
190 `${refMs.toFixed(2)} ms (${(refMs / mgpuMs).toFixed(1)}x)`,
191 );
192 // A soft bound, not a performance target: on a software rasterizer the
193 // transforms dominate and the round-trips this path avoids are a small
194 // share of the total, so the ratio understates what it is worth on real
195 // hardware. The check is only that executing the .m did not make it worse.
196 check(
197 `${m.key}: step rate no worse than the readback path`,
198 mgpuMs < refMs * 1.15,
199 `${mgpuMs.toFixed(2)} vs ${refMs.toFixed(2)} ms/step`,
200 );
201 }
202}