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 solver: 'bicgstab',
82 };
83 const command = formatCommand(spec);
84 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
85 const same = JSON.stringify(back) === JSON.stringify(spec);
86 check(
87 'runSpec: the benchmark command round-trips every field',
88 same,
89 same ? command.slice(BENCH_COMMAND.length + 1) : `got ${JSON.stringify(back)}`,
90 );
91 }
93 for (const model of mModels) {
94 const session = await ModelSession.create({
95 device,
96 model,
97 params: defaultParams(model),
98 lmax: LMAX,
99 niter: NITER,
100 });
102 const plan = session.describe();
103 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
104 const xforms = plan.step.filter(
105 (l) => l.startsWith('synth') || l.startsWith('analys'),
106 ).length;
107 const expected = EXPECTED_KERNELS[model.key] + NITER * KERNELS_PER_ITERATION[model.key];
108 log(
109 ` ${model.key}.m -> ${plan.step.length} ops/step ` +
110 `(${kernels} generated kernels, ${xforms} transforms, ${NITER} solve iter)`,
111 );
112 check(
113 `${model.key}: element-wise lines fused into one kernel each`,
114 kernels === expected,
115 `${kernels} kernels (expected ${expected})`,
116 );
118 session.seed(1);
119 session.step(STEPS);
121 // Every rendered field must be finite and have developed some contrast.
122 for (const field of model.species) {
123 const values = await session.read(field);
124 let lo = Infinity;
125 let hi = -Infinity;
126 let finite = true;
127 for (const v of values) {
128 if (!Number.isFinite(v)) finite = false;
129 if (v < lo) lo = v;
130 if (v > hi) hi = v;
131 }
132 check(
133 `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
134 finite && hi - lo > 1e-6,
135 finite
136 ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
137 : 'contains NaN or Infinity',
138 );
139 }
141 session.destroy();
142 }
144 // User-defined subroutines: a .m may define its own functions (and call the
145 // shared solver/operator library), and each call is expanded into the caller
146 // at compile time (src/mgpu/inlineCalls.ts). This model exercises the
147 // shapes the shipped models do not: a multi-output function, a
148 // scalar-returning function, a function reassigning its own parameter, and
149 // a solver-like local whose loop bound arrives as the `niter` argument.
150 {
151 const model = mModels.find((m) => m.key === 'allencahn')!;
152 const source = `
153function [U, u] = init(noise)
154 U = analys(noise);
155 u = synth(U);
156end
158function [Un, u] = step(U, lam, eps2, dt, niter)
159 u = synth(U);
160 [p, q] = react(u, dt);
161 s = gain(eps2, dt);
162 Bu = U + s * analys(p - q);
163 Un = solveid(Bu, lam, dt, niter);
164end
166function [p, q] = react(x, c)
167 p = x + c * (x .* x);
168 q = c * (x .* x);
169end
171function y = gain(a, b)
172 y = a + 2 * b;
173end
175function X = solveid(B, lam, c, n)
176 X = B ./ (1 + c * lam);
177 for k = 1:n
178 X = (B + c * (0 * X)) ./ (1 + c * lam);
179 end
180end
181`;
182 const session = await ModelSession.create({
183 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 2,
184 });
185 session.seed(1);
186 session.step(STEPS);
187 const values = await session.read('u');
188 let finite = true;
189 for (const v of values) if (!Number.isFinite(v)) finite = false;
190 check(
191 'subroutines: a model composed of user functions compiles and runs',
192 finite,
193 `${session.describe().step.length} ops/step after expansion`,
194 );
195 session.destroy();
196 }
198 // The reduction op and GPU-resident scalars: `dot` runs as a single
199 // reduction dispatch into a 1-element buffer, scalars computed from its
200 // result compile to 1-element kernels, and a single-element value
201 // broadcasts into element-wise expressions as `in[0]`. These are the
202 // primitives the Krylov solver is made of, checked directly against the
203 // CPU here so a solver-level failure has somewhere smaller to point.
204 {
205 const model = mModels.find((m) => m.key === 'allencahn')!;
206 const source = `
207function [U, u] = init(noise)
208 U = analys(noise);
209 u = synth(U);
210end
212function [Un, u] = step(U, lam, wlm, eps2, dt, niter)
213 u = synth(U);
214 s = dot(U, U);
215 Uw = U .* wlm;
216 sw = dot(Uw, lam);
217 s2 = 2 * s;
218 s3 = s2 - s;
219 Un = (s * U) ./ s;
220end
221`;
222 const session = await ModelSession.create({
223 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
224 });
225 session.seed(1);
226 session.step(1);
227 const U = await session.read('U');
228 const nlm = U.length / 2;
229 const cfg = session.cfg;
231 let cpuS = 0;
232 for (let i = 0; i < U.length; i++) cpuS += U[i] * U[i];
233 const gpuS = (await session.read('s'))[0];
234 check(
235 'dot: matches the CPU sum',
236 Math.abs(gpuS - cpuS) <= 1e-5 * Math.abs(cpuS),
237 `gpu ${gpuS.toExponential(6)} vs cpu ${cpuS.toExponential(6)}`,
238 );
240 const wlm = weightMask(cfg, nlm);
241 const lam = eigenvalues(cfg, nlm);
242 let cpuSw = 0;
243 for (let i = 0; i < U.length; i++) cpuSw += U[i] * wlm[i] * lam[i];
244 const gpuSw = (await session.read('sw'))[0];
245 check(
246 'dot: the wlm-weighted inner product matches the CPU',
247 Math.abs(gpuSw - cpuSw) <= 1e-5 * Math.abs(cpuSw),
248 `gpu ${gpuSw.toExponential(6)} vs cpu ${cpuSw.toExponential(6)}`,
249 );
251 // 2s - s is exact in any IEEE arithmetic, so the whole scalar chain
252 // (reduction -> 1-element kernels -> readback) must return s's bits.
253 const gpuS3 = (await session.read('s3'))[0];
254 check('dot: scalar arithmetic on the result is exact', gpuS3 === gpuS,
255 `s3 ${gpuS3.toExponential(6)} vs s ${gpuS.toExponential(6)}`);
257 const Un = await session.read('Un');
258 let worst = 0;
259 let scale = 0;
260 for (let i = 0; i < U.length; i++) {
261 worst = Math.max(worst, Math.abs(Un[i] - U[i]));
262 scale = Math.max(scale, Math.abs(U[i]));
263 }
264 check(
265 'dot: a 1-element value broadcasts into an element-wise kernel',
266 worst <= 1e-6 * scale,
267 `(s*U)./s vs U: worst |d| = ${worst.toExponential(2)}`,
268 );
269 session.destroy();
270 }
272 // The indexed-access ops (getslab/setslab on a bank of spectral fields,
273 // getat/setat on a small matrix): functional updates the planner compiles
274 // to static-offset buffer copies. Everything below has an exact expected
275 // value, so the offsets themselves are what is being checked.
276 {
277 const model = mModels.find((m) => m.key === 'allencahn')!;
278 const source = `
279function [U, u] = init(noise)
280 U = analys(noise);
281 u = synth(U);
282end
284function [Un, u] = step(U, lam, eps2, dt, nlm, niter)
285 u = synth(U);
286 A = zeros(2, 2);
287 s1 = dot(U, U);
288 s2 = 2 * s1;
289 A = setat(A, s1, 1, 1);
290 A = setat(A, s2, 2, 2);
291 a11 = getat(A, 1, 1);
292 a22 = getat(A, 2, 2);
293 a21 = getat(A, 2, 1);
294 chk = a22 - 2 * a11 + a21;
295 VB = zeros(2, nlm * 2);
296 VB = setslab(VB, U, 2);
297 U2 = getslab(VB, 2);
298 Z1 = getslab(VB, 1);
299 Un = U2 + Z1;
300end
301`;
302 const session = await ModelSession.create({
303 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
304 });
305 session.seed(1);
306 session.step(1);
307 // a22 - 2*a11 + a21 = 2*s - 2*s + 0, exactly, if every element landed
308 // where its indices say.
309 const chk = (await session.read('chk'))[0];
310 check('indexing: matrix elements round-trip through setat/getat', chk === 0,
311 `a22 - 2*a11 + a21 = ${chk}`);
312 // The slab written at 2 must come back; the slab at 1 must still be zero.
313 const U = await session.read('U');
314 const Un = await session.read('Un');
315 let same = U.length === Un.length;
316 for (let i = 0; same && i < U.length; i++) if (Un[i] !== U[i]) same = false;
317 check('indexing: a spectral field round-trips through setslab/getslab', same,
318 same ? 'getslab(setslab(VB, U, 2), 2) + zeros = U, element for element' : 'mismatch');
319 session.destroy();
320 }
322 // A recursive function cannot unroll into a fixed op sequence, and must be
323 // refused with a message that says so, not hang the compiler.
324 {
325 const model = mModels.find((m) => m.key === 'allencahn')!;
326 const source = `
327function [U, u] = init(noise)
328 U = analys(noise);
329 u = synth(U);
330end
332function [Un, u] = step(U, lam, eps2, dt, niter)
333 u = synth(U);
334 Un = f(U);
335end
337function y = f(x)
338 y = f(x) + 1;
339end
340`;
341 let message = '';
342 try {
343 const session = await ModelSession.create({
344 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
345 });
346 session.destroy();
347 } catch (e) {
348 message = e instanceof Error ? e.message : String(e);
349 }
350 check(
351 'subroutines: recursion is refused at compile time',
352 message.includes('recursion'),
353 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
354 );
355 }
357 // The oversampled readback: readSpecies must be the state synthesized on the
358 // display grid. Comparing against the display plan's own upload path
359 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
360 // coefficient copy against a known-good route through the same kernels.
361 {
362 const model = mModels.find((m) => m.key === 'allencahn')!;
363 const session = await ModelSession.create({
364 device,
365 model,
366 params: defaultParams(model),
367 lmax: LMAX,
368 oversample: 2,
369 });
370 session.seed(1);
371 session.step(STEPS);
373 const fine = await session.readSpecies(0);
374 const { nlat, nphi } = session.viewSht.cfg;
375 check(
376 'oversample: species field is on the 2x display grid',
377 nlat === 2 * session.cfg.nlat &&
378 nphi === 2 * session.cfg.nphi &&
379 fine.length === nlat * nphi,
380 `render ${nlat}×${nphi}, ${fine.length} values`,
381 );
383 const qlm = await session.read('U');
384 const expected = await session.viewSht.synth(qlm);
385 let maxDiff = 0;
386 for (let i = 0; i < fine.length; i++) {
387 const d = Math.abs(fine[i] - expected[i]);
388 if (d > maxDiff) maxDiff = d;
389 }
390 check(
391 'oversample: readSpecies matches synth of the read-back state',
392 maxDiff <= 1e-6,
393 `max |diff| = ${maxDiff.toExponential(2)}`,
394 );
396 // A timing burst must be invisible: the state is snapshotted and restored
397 // around it, and model time does not advance.
398 const tBefore = session.t;
399 const stepsBefore = session.steps;
400 const ms = await session.measure(8);
401 const after = await session.read('U');
402 let identical = qlm.length === after.length;
403 if (identical) {
404 for (let i = 0; i < qlm.length; i++) {
405 if (qlm[i] !== after[i]) {
406 identical = false;
407 break;
408 }
409 }
410 }
411 check(
412 'measure: a timing burst leaves state, t and steps untouched',
413 identical && session.t === tBefore && session.steps === stepsBefore,
414 identical
415 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
416 : 'state changed',
417 );
419 // Changing the oversampling in place is display-only: the state survives
420 // and the render grid drops back to the solver's.
421 await session.setOversample(1);
422 const qlmAfterSwap = await session.read('U');
423 let stateSurvived = qlmAfterSwap.length === after.length;
424 if (stateSurvived) {
425 for (let i = 0; i < after.length; i++) {
426 if (qlmAfterSwap[i] !== after[i]) {
427 stateSurvived = false;
428 break;
429 }
430 }
431 }
432 session.step(1); // recompute the view fields on the solver grid
433 const coarse = await session.readSpecies(0);
434 check(
435 'setOversample: swaps the render grid without touching the state',
436 stateSurvived &&
437 session.viewSht === session.sht &&
438 coarse.length === session.cfg.nlat * session.cfg.nphi,
439 stateSurvived
440 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
441 : 'state changed',
442 );
444 session.destroy();
445 }
446}