1/**
2 * Every model the app offers: that it compiles, what it compiles to, and that it
3 * runs stably and produces a pattern.
4 *
5 * Numerical correctness of the pipeline is analyticChecks.ts's job. This file is
6 * about the models themselves and about the compilation staying as intended — in
7 * particular the kernel count, which is a fusion guard: numbl's lowering emits
8 * one statement per *operator*, and its inline pass folds those back into
9 * per-line expression trees. If that stops happening the results stay correct
10 * but every operator becomes its own dispatch, which is invisible except here.
11 */
12import { ModelSession } from '../src/mgpu/session.ts';
13import { mModels, defaultParams } from '../src/mgpu/registry.ts';
14import { eigenvalues, weightMask } from '../src/mgpu/model.ts';
15import {
16 formatCommand,
17 parseArgs,
18 BENCH_COMMAND,
19 type RunSpec,
20} from '../src/bench/runSpec.ts';
21import type { Check, Log } from './analyticChecks.ts';
23/**
24 * Kernels each model's step compiles to outside its solve loop — one per
25 * element-wise line, where the argument of a transform counts as its own line
26 * (it cannot fuse into an external call).
27 */
28const EXPECTED_KERNELS: Record<string, number> = {
29 schnakenberg: 7,
30 brusselator: 7,
31 allencahn: 3,
32};
34/**
35 * What one unrolled iteration of the solve loop adds — 14 kernels per
36 * species: 12 in lib/dlap.m's operator (the gradient contraction, the three
37 * re-analysed components, the five-step divergence accumulation), plus
38 * solvers/richardson.m's dtD*lam divisor temp and its update divide. Each
39 * species' correction is Algorithm 3 of evolving_surface/notes/algos.tex: a
40 * surface gradient (dtheta/dphi contracted through the metric), reanalysed
41 * per Cartesian component and differentiated again, recombined into the
42 * divergence, plus the round-sphere eigenvalue added back — see
43 * solvers/richardson.m, lib/dlap.m and docs/richardson-iteration.md.
44 * (Before the solver was factored out, the monolithic models compiled to 15
45 * per species: the interleaved species order kept the second species'
46 * divisor from fusing into its divide.)
47 */
48const KERNELS_PER_ITERATION: Record<string, number> = {
49 schnakenberg: 28,
50 brusselator: 28,
51 allencahn: 14,
52};
54const LMAX = 31;
55const STEPS = 40;
56const NITER = 1;
58export async function modelChecks(
59 device: GPUDevice,
60 check: Check,
61 log: Log,
62): Promise<void> {
63 check('models: registry populated', mModels.length === 3, `${mModels.length} models`);
65 // The app formats the run it is showing into a `npm run bench` command and
66 // the benchmark parses it back. That is only worth anything if the round
67 // trip is lossless — a knob that formatCommand forgets is a knob the desktop
68 // run would silently take a default for, and the two runs would differ while
69 // claiming to be the same. Every field of the spec, through both directions.
70 {
71 const spec: RunSpec = {
72 preset: 'schnak-fine',
73 lmax: 127,
74 seed: 12345,
75 steps: 777,
76 warmup: 13,
77 params: { a: 0.11, b: 0.91, D1: 5e-4, D2: 9e-3, dt: 0.04 },
78 geometry: 'peanut',
79 geometryParams: { waist: 0.45, stretch: 1.25 },
80 niter: 3,
81 };
82 const command = formatCommand(spec);
83 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
84 const same = JSON.stringify(back) === JSON.stringify(spec);
85 check(
86 'runSpec: the benchmark command round-trips every field',
87 same,
88 same ? command.slice(BENCH_COMMAND.length + 1) : `got ${JSON.stringify(back)}`,
89 );
90 }
92 for (const model of mModels) {
93 const session = await ModelSession.create({
94 device,
95 model,
96 params: defaultParams(model),
97 lmax: LMAX,
98 niter: NITER,
99 });
101 const plan = session.describe();
102 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
103 const xforms = plan.step.filter(
104 (l) => l.startsWith('synth') || l.startsWith('analys'),
105 ).length;
106 const expected = EXPECTED_KERNELS[model.key] + NITER * KERNELS_PER_ITERATION[model.key];
107 log(
108 ` ${model.key}.m -> ${plan.step.length} ops/step ` +
109 `(${kernels} generated kernels, ${xforms} transforms, ${NITER} solve iter)`,
110 );
111 check(
112 `${model.key}: element-wise lines fused into one kernel each`,
113 kernels === expected,
114 `${kernels} kernels (expected ${expected})`,
115 );
117 session.seed(1);
118 session.step(STEPS);
120 // Every rendered field must be finite and have developed some contrast.
121 for (const field of model.species) {
122 const values = await session.read(field);
123 let lo = Infinity;
124 let hi = -Infinity;
125 let finite = true;
126 for (const v of values) {
127 if (!Number.isFinite(v)) finite = false;
128 if (v < lo) lo = v;
129 if (v > hi) hi = v;
130 }
131 check(
132 `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
133 finite && hi - lo > 1e-6,
134 finite
135 ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
136 : 'contains NaN or Infinity',
137 );
138 }
140 session.destroy();
141 }
143 // User-defined subroutines: a .m may define its own functions (and call the
144 // shared solver/operator library), and each call is expanded into the caller
145 // at compile time (src/mgpu/inlineCalls.ts). This model exercises the
146 // shapes the shipped models do not: a multi-output function, a
147 // scalar-returning function, a function reassigning its own parameter, and
148 // a solver-like local whose loop bound arrives as the `niter` argument.
149 {
150 const model = mModels.find((m) => m.key === 'allencahn')!;
151 const source = `
152function [U, u] = init(noise)
153 U = analys(noise);
154 u = synth(U);
155end
157function [Un, u] = step(U, lam, eps2, dt, niter)
158 u = synth(U);
159 [p, q] = react(u, dt);
160 s = gain(eps2, dt);
161 Bu = U + s * analys(p - q);
162 Un = solveid(Bu, lam, dt, niter);
163end
165function [p, q] = react(x, c)
166 p = x + c * (x .* x);
167 q = c * (x .* x);
168end
170function y = gain(a, b)
171 y = a + 2 * b;
172end
174function X = solveid(B, lam, c, n)
175 X = B ./ (1 + c * lam);
176 for k = 1:n
177 X = (B + c * (0 * X)) ./ (1 + c * lam);
178 end
179end
180`;
181 const session = await ModelSession.create({
182 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 2,
183 });
184 session.seed(1);
185 session.step(STEPS);
186 const values = await session.read('u');
187 let finite = true;
188 for (const v of values) if (!Number.isFinite(v)) finite = false;
189 check(
190 'subroutines: a model composed of user functions compiles and runs',
191 finite,
192 `${session.describe().step.length} ops/step after expansion`,
193 );
194 session.destroy();
195 }
197 // The reduction op and GPU-resident scalars: `dot` runs as a single
198 // reduction dispatch into a 1-element buffer, scalars computed from its
199 // result compile to 1-element kernels, and a single-element value
200 // broadcasts into element-wise expressions as `in[0]`. These are the
201 // primitives the Krylov solver is made of, checked directly against the
202 // CPU here so a solver-level failure has somewhere smaller to point.
203 {
204 const model = mModels.find((m) => m.key === 'allencahn')!;
205 const source = `
206function [U, u] = init(noise)
207 U = analys(noise);
208 u = synth(U);
209end
211function [Un, u] = step(U, lam, wlm, eps2, dt, niter)
212 u = synth(U);
213 s = dot(U, U);
214 Uw = U .* wlm;
215 sw = dot(Uw, lam);
216 s2 = 2 * s;
217 s3 = s2 - s;
218 Un = (s * U) ./ s;
219end
220`;
221 const session = await ModelSession.create({
222 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
223 });
224 session.seed(1);
225 session.step(1);
226 const U = await session.read('U');
227 const nlm = U.length / 2;
228 const cfg = session.cfg;
230 let cpuS = 0;
231 for (let i = 0; i < U.length; i++) cpuS += U[i] * U[i];
232 const gpuS = (await session.read('s'))[0];
233 check(
234 'dot: matches the CPU sum',
235 Math.abs(gpuS - cpuS) <= 1e-5 * Math.abs(cpuS),
236 `gpu ${gpuS.toExponential(6)} vs cpu ${cpuS.toExponential(6)}`,
237 );
239 const wlm = weightMask(cfg, nlm);
240 const lam = eigenvalues(cfg, nlm);
241 let cpuSw = 0;
242 for (let i = 0; i < U.length; i++) cpuSw += U[i] * wlm[i] * lam[i];
243 const gpuSw = (await session.read('sw'))[0];
244 check(
245 'dot: the wlm-weighted inner product matches the CPU',
246 Math.abs(gpuSw - cpuSw) <= 1e-5 * Math.abs(cpuSw),
247 `gpu ${gpuSw.toExponential(6)} vs cpu ${cpuSw.toExponential(6)}`,
248 );
250 // 2s - s is exact in any IEEE arithmetic, so the whole scalar chain
251 // (reduction -> 1-element kernels -> readback) must return s's bits.
252 const gpuS3 = (await session.read('s3'))[0];
253 check('dot: scalar arithmetic on the result is exact', gpuS3 === gpuS,
254 `s3 ${gpuS3.toExponential(6)} vs s ${gpuS.toExponential(6)}`);
256 const Un = await session.read('Un');
257 let worst = 0;
258 let scale = 0;
259 for (let i = 0; i < U.length; i++) {
260 worst = Math.max(worst, Math.abs(Un[i] - U[i]));
261 scale = Math.max(scale, Math.abs(U[i]));
262 }
263 check(
264 'dot: a 1-element value broadcasts into an element-wise kernel',
265 worst <= 1e-6 * scale,
266 `(s*U)./s vs U: worst |d| = ${worst.toExponential(2)}`,
267 );
268 session.destroy();
269 }
271 // The indexed-access ops (getslab/setslab on a bank of spectral fields,
272 // getat/setat on a small matrix): functional updates the planner compiles
273 // to static-offset buffer copies. Everything below has an exact expected
274 // value, so the offsets themselves are what is being checked.
275 {
276 const model = mModels.find((m) => m.key === 'allencahn')!;
277 const source = `
278function [U, u] = init(noise)
279 U = analys(noise);
280 u = synth(U);
281end
283function [Un, u] = step(U, lam, eps2, dt, nlm, niter)
284 u = synth(U);
285 A = zeros(2, 2);
286 s1 = dot(U, U);
287 s2 = 2 * s1;
288 A = setat(A, s1, 1, 1);
289 A = setat(A, s2, 2, 2);
290 a11 = getat(A, 1, 1);
291 a22 = getat(A, 2, 2);
292 a21 = getat(A, 2, 1);
293 chk = a22 - 2 * a11 + a21;
294 VB = zeros(2, nlm * 2);
295 VB = setslab(VB, U, 2);
296 U2 = getslab(VB, 2);
297 Z1 = getslab(VB, 1);
298 Un = U2 + Z1;
299end
300`;
301 const session = await ModelSession.create({
302 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
303 });
304 session.seed(1);
305 session.step(1);
306 // a22 - 2*a11 + a21 = 2*s - 2*s + 0, exactly, if every element landed
307 // where its indices say.
308 const chk = (await session.read('chk'))[0];
309 check('indexing: matrix elements round-trip through setat/getat', chk === 0,
310 `a22 - 2*a11 + a21 = ${chk}`);
311 // The slab written at 2 must come back; the slab at 1 must still be zero.
312 const U = await session.read('U');
313 const Un = await session.read('Un');
314 let same = U.length === Un.length;
315 for (let i = 0; same && i < U.length; i++) if (Un[i] !== U[i]) same = false;
316 check('indexing: a spectral field round-trips through setslab/getslab', same,
317 same ? 'getslab(setslab(VB, U, 2), 2) + zeros = U, element for element' : 'mismatch');
318 session.destroy();
319 }
321 // A recursive function cannot unroll into a fixed op sequence, and must be
322 // refused with a message that says so, not hang the compiler.
323 {
324 const model = mModels.find((m) => m.key === 'allencahn')!;
325 const source = `
326function [U, u] = init(noise)
327 U = analys(noise);
328 u = synth(U);
329end
331function [Un, u] = step(U, lam, eps2, dt, niter)
332 u = synth(U);
333 Un = f(U);
334end
336function y = f(x)
337 y = f(x) + 1;
338end
339`;
340 let message = '';
341 try {
342 const session = await ModelSession.create({
343 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
344 });
345 session.destroy();
346 } catch (e) {
347 message = e instanceof Error ? e.message : String(e);
348 }
349 check(
350 'subroutines: recursion is refused at compile time',
351 message.includes('recursion'),
352 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
353 );
354 }
356 // The oversampled readback: readSpecies must be the state synthesized on the
357 // display grid. Comparing against the display plan's own upload path
358 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
359 // coefficient copy against a known-good route through the same kernels.
360 {
361 const model = mModels.find((m) => m.key === 'allencahn')!;
362 const session = await ModelSession.create({
363 device,
364 model,
365 params: defaultParams(model),
366 lmax: LMAX,
367 oversample: 2,
368 });
369 session.seed(1);
370 session.step(STEPS);
372 const fine = await session.readSpecies(0);
373 const { nlat, nphi } = session.viewSht.cfg;
374 check(
375 'oversample: species field is on the 2x display grid',
376 nlat === 2 * session.cfg.nlat &&
377 nphi === 2 * session.cfg.nphi &&
378 fine.length === nlat * nphi,
379 `render ${nlat}×${nphi}, ${fine.length} values`,
380 );
382 const qlm = await session.read('U');
383 const expected = await session.viewSht.synth(qlm);
384 let maxDiff = 0;
385 for (let i = 0; i < fine.length; i++) {
386 const d = Math.abs(fine[i] - expected[i]);
387 if (d > maxDiff) maxDiff = d;
388 }
389 check(
390 'oversample: readSpecies matches synth of the read-back state',
391 maxDiff <= 1e-6,
392 `max |diff| = ${maxDiff.toExponential(2)}`,
393 );
395 // A timing burst must be invisible: the state is snapshotted and restored
396 // around it, and model time does not advance.
397 const tBefore = session.t;
398 const stepsBefore = session.steps;
399 const ms = await session.measure(8);
400 const after = await session.read('U');
401 let identical = qlm.length === after.length;
402 if (identical) {
403 for (let i = 0; i < qlm.length; i++) {
404 if (qlm[i] !== after[i]) {
405 identical = false;
406 break;
407 }
408 }
409 }
410 check(
411 'measure: a timing burst leaves state, t and steps untouched',
412 identical && session.t === tBefore && session.steps === stepsBefore,
413 identical
414 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
415 : 'state changed',
416 );
418 // Changing the oversampling in place is display-only: the state survives
419 // and the render grid drops back to the solver's.
420 await session.setOversample(1);
421 const qlmAfterSwap = await session.read('U');
422 let stateSurvived = qlmAfterSwap.length === after.length;
423 if (stateSurvived) {
424 for (let i = 0; i < after.length; i++) {
425 if (qlmAfterSwap[i] !== after[i]) {
426 stateSurvived = false;
427 break;
428 }
429 }
430 }
431 session.step(1); // recompute the view fields on the solver grid
432 const coarse = await session.readSpecies(0);
433 check(
434 'setOversample: swaps the render grid without touching the state',
435 stateSurvived &&
436 session.viewSht === session.sht &&
437 coarse.length === session.cfg.nlat * session.cfg.nphi,
438 stateSurvived
439 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
440 : 'state changed',
441 );
443 session.destroy();
444 }
445}