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 {
15 formatCommand,
16 parseArgs,
17 BENCH_COMMAND,
18 type RunSpec,
19} from '../src/bench/runSpec.ts';
20import type { Check, Log } from './analyticChecks.ts';
22/**
23 * Kernels each model's step compiles to outside its solve loop — one per
24 * element-wise line, where the argument of a transform counts as its own line
25 * (it cannot fuse into an external call).
26 */
27const EXPECTED_KERNELS: Record<string, number> = {
28 schnakenberg: 7,
29 brusselator: 7,
30 allencahn: 3,
31 'schnakenberg-alg4': 7,
32};
34/**
35 * What one unrolled iteration of the solve loop adds, total (not per
36 * species — the surface Laplace-Beltrami correction's per-species kernel
37 * count is a byproduct of exactly how its expression tree happens to fuse,
38 * not a clean per-species multiple, so this is measured per model rather
39 * than derived from `model.species.length`). Each species' correction is
40 * the flux-form Laplace-Beltrami matvec of
41 * docs/reduced-transforms.md Sec 4: the two sin-weighted
42 * derivative synths, the pointwise flux combination through p1/p2/q2, the
43 * two flux analyses, the re-shifted divergence and its r-scaled synthesis,
44 * plus the round-sphere eigenvalue added back — see models/schnakenberg.m
45 * and docs/richardson-iteration.md. `schnakenberg-alg4` keeps the original
46 * Cartesian-gradient form (Algorithm 3/4 of evolving_surface/notes/algos.tex)
47 * as a live reference, with its original counts.
48 */
49const KERNELS_PER_ITERATION: Record<string, number> = {
50 schnakenberg: 18,
51 brusselator: 18,
52 allencahn: 9,
53 'schnakenberg-alg4': 30,
54};
56const LMAX = 31;
57const STEPS = 40;
58const NITER = 1;
60export async function modelChecks(
61 device: GPUDevice,
62 check: Check,
63 log: Log,
64): Promise<void> {
65 check('models: registry populated', mModels.length === 4, `${mModels.length} models`);
67 // The app formats the run it is showing into a `npm run bench` command and
68 // the benchmark parses it back. That is only worth anything if the round
69 // trip is lossless — a knob that formatCommand forgets is a knob the desktop
70 // run would silently take a default for, and the two runs would differ while
71 // claiming to be the same. Every field of the spec, through both directions.
72 {
73 const spec: RunSpec = {
74 preset: 'schnak-fine',
75 lmax: 127,
76 seed: 12345,
77 steps: 777,
78 warmup: 13,
79 params: { a: 0.11, b: 0.91, D1: 5e-4, D2: 9e-3, dt: 0.04 },
80 geometry: 'peanut',
81 geometryParams: { waist: 0.45, stretch: 1.25 },
82 niter: 3,
83 };
84 const command = formatCommand(spec);
85 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
86 const same = JSON.stringify(back) === JSON.stringify(spec);
87 check(
88 'runSpec: the benchmark command round-trips every field',
89 same,
90 same ? command.slice(BENCH_COMMAND.length + 1) : `got ${JSON.stringify(back)}`,
91 );
92 }
94 for (const model of mModels) {
95 const session = await ModelSession.create({
96 device,
97 model,
98 params: defaultParams(model),
99 lmax: LMAX,
100 niter: NITER,
101 });
103 const plan = session.describe();
104 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
105 const xforms = plan.step.filter(
106 (l) => l.startsWith('synth') || l.startsWith('analys'),
107 ).length;
108 const expected = EXPECTED_KERNELS[model.key] + NITER * KERNELS_PER_ITERATION[model.key];
109 log(
110 ` ${model.key}.m -> ${plan.step.length} ops/step ` +
111 `(${kernels} generated kernels, ${xforms} transforms, ${NITER} solve iter)`,
112 );
113 check(
114 `${model.key}: element-wise lines fused into one kernel each`,
115 kernels === expected,
116 `${kernels} kernels (expected ${expected})`,
117 );
119 session.seed(1);
120 session.step(STEPS);
122 // Every rendered field must be finite and have developed some contrast.
123 for (const field of model.species) {
124 const values = await session.read(field);
125 let lo = Infinity;
126 let hi = -Infinity;
127 let finite = true;
128 for (const v of values) {
129 if (!Number.isFinite(v)) finite = false;
130 if (v < lo) lo = v;
131 if (v > hi) hi = v;
132 }
133 check(
134 `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
135 finite && hi - lo > 1e-6,
136 finite
137 ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
138 : 'contains NaN or Infinity',
139 );
140 }
142 session.destroy();
143 }
145 // The oversampled readback: readSpecies must be the state synthesized on the
146 // display grid. Comparing against the display plan's own upload path
147 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
148 // coefficient copy against a known-good route through the same kernels.
149 {
150 const model = mModels.find((m) => m.key === 'allencahn')!;
151 const session = await ModelSession.create({
152 device,
153 model,
154 params: defaultParams(model),
155 lmax: LMAX,
156 oversample: 2,
157 });
158 session.seed(1);
159 session.step(STEPS);
161 const fine = await session.readSpecies(0);
162 const { nlat, nphi } = session.viewSht.cfg;
163 check(
164 'oversample: species field is on the 2x display grid',
165 nlat === 2 * session.cfg.nlat &&
166 nphi === 2 * session.cfg.nphi &&
167 fine.length === nlat * nphi,
168 `render ${nlat}×${nphi}, ${fine.length} values`,
169 );
171 const qlm = await session.read('U');
172 const expected = await session.viewSht.synth(qlm);
173 let maxDiff = 0;
174 for (let i = 0; i < fine.length; i++) {
175 const d = Math.abs(fine[i] - expected[i]);
176 if (d > maxDiff) maxDiff = d;
177 }
178 check(
179 'oversample: readSpecies matches synth of the read-back state',
180 maxDiff <= 1e-6,
181 `max |diff| = ${maxDiff.toExponential(2)}`,
182 );
184 // A timing burst must be invisible: the state is snapshotted and restored
185 // around it, and model time does not advance.
186 const tBefore = session.t;
187 const stepsBefore = session.steps;
188 const ms = await session.measure(8);
189 const after = await session.read('U');
190 let identical = qlm.length === after.length;
191 if (identical) {
192 for (let i = 0; i < qlm.length; i++) {
193 if (qlm[i] !== after[i]) {
194 identical = false;
195 break;
196 }
197 }
198 }
199 check(
200 'measure: a timing burst leaves state, t and steps untouched',
201 identical && session.t === tBefore && session.steps === stepsBefore,
202 identical
203 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
204 : 'state changed',
205 );
207 // Changing the oversampling in place is display-only: the state survives
208 // and the render grid drops back to the solver's.
209 await session.setOversample(1);
210 const qlmAfterSwap = await session.read('U');
211 let stateSurvived = qlmAfterSwap.length === after.length;
212 if (stateSurvived) {
213 for (let i = 0; i < after.length; i++) {
214 if (qlmAfterSwap[i] !== after[i]) {
215 stateSurvived = false;
216 break;
217 }
218 }
219 }
220 session.step(1); // recompute the view fields on the solver grid
221 const coarse = await session.readSpecies(0);
222 check(
223 'setOversample: swaps the render grid without touching the state',
224 stateSurvived &&
225 session.viewSht === session.sht &&
226 coarse.length === session.cfg.nlat * session.cfg.nphi,
227 stateSurvived
228 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
229 : 'state changed',
230 );
232 session.destroy();
233 }
234}