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, mModelByKey, 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 // Batched transforms are an encoding of the same arithmetic, so a run with
146 // batching disabled (SHT_BATCH=0 compiles scalar-only plans) must reproduce
147 // the default run to shader-compiler latitude, and the default run must
148 // actually be batching (the describe() lines say so). This is the guard
149 // that the planner's adjacency grouping rewires buffers correctly — a lane
150 // bound to the wrong field would miss by O(1), not O(1e-6).
151 {
152 const model = mModelByKey('schnakenberg')!;
153 const params = defaultParams(model);
154 const states: Float32Array[] = [];
155 let batchedLanes = 0;
156 for (const batch of [undefined, 0]) {
157 const g = globalThis as Record<string, unknown>;
158 if (batch !== undefined) g.SHT_BATCH = batch;
159 try {
160 const session = await ModelSession.create({
161 device, model, params, lmax: LMAX, niter: NITER,
162 });
163 if (batch === undefined) {
164 batchedLanes = session
165 .describe()
166 .step.filter((l) => l.includes('[batch lane')).length;
167 }
168 session.seed(1);
169 session.step(STEPS);
170 states.push(await session.read('U'));
171 session.destroy();
172 } finally {
173 delete g.SHT_BATCH;
174 }
175 }
176 // Every batchable run at one solve iteration: the u/v syntheses and the
177 // reaction analyses outside the loop (2 + 2), the four gradient
178 // syntheses, four flux analyses, two divergence syntheses and two final
179 // analyses inside it (4 + 4 + 2 + 2). Lane counts are batch-width
180 // invariant: a x4 run is one batch at K = 4 and two at K = 2, but the
181 // lanes annotated are the same 16 either way.
182 check(
183 'batch: the compiled step batches every adjacent transform pair',
184 batchedLanes === 16,
185 `${batchedLanes} batched transform lanes (expected 16)`,
186 );
187 let worst = 0;
188 for (let i = 0; i < states[0].length; i++) {
189 worst = Math.max(worst, Math.abs(states[0][i] - states[1][i]));
190 }
191 check(
192 'batch: batched and scalar plans agree through a real run',
193 worst < 1e-4,
194 `max |U_batched - U_scalar| = ${worst.toExponential(2)} after ${STEPS} steps`,
195 );
196 }
198 // Misusing the grouped-transform syntax is refused at compile time with a
199 // message that says how to write it, not silently mis-planned: every input
200 // must get an output (each one costs a transform), whether the mismatch is
201 // an under-bound assignment or an ignored slot.
202 {
203 const model = mModelByKey('allencahn')!;
204 const cases: [string, string, string][] = [
205 [
206 'a single output bound to a grouped call',
207 'Ftu = synth(vtu, vpu);',
208 'bind each one',
209 ],
210 [
211 'an ignored output slot',
212 // Fpu is reassigned so the only error left is the dropped slot
213 // itself, which the planner refuses (numbl would otherwise catch
214 // the undefined 'Fpu' first, masking the check under test).
215 '[Ftu, ~] = synth(vtu, vpu);\n Fpu = Ftu;',
216 'must be bound',
217 ],
218 ];
219 for (const [what, bad, expect] of cases) {
220 const source = model.source.replace('[Ftu, Fpu] = synth(vtu, vpu);', bad);
221 let message = '';
222 try {
223 const session = await ModelSession.create({
224 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
225 });
226 session.destroy();
227 } catch (e) {
228 message = e instanceof Error ? e.message : String(e);
229 }
230 check(
231 `batch: ${what} is refused at compile time`,
232 message.includes(expect),
233 message ? `refused: ${message.slice(0, 76)}…` : 'compiled anyway',
234 );
235 }
236 }
238 // The oversampled readback: readSpecies must be the state synthesized on the
239 // display grid. Comparing against the display plan's own upload path
240 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
241 // coefficient copy against a known-good route through the same kernels.
242 {
243 const model = mModels.find((m) => m.key === 'allencahn')!;
244 const session = await ModelSession.create({
245 device,
246 model,
247 params: defaultParams(model),
248 lmax: LMAX,
249 oversample: 2,
250 });
251 session.seed(1);
252 session.step(STEPS);
254 const fine = await session.readSpecies(0);
255 const { nlat, nphi } = session.viewSht.cfg;
256 check(
257 'oversample: species field is on the 2x display grid',
258 nlat === 2 * session.cfg.nlat &&
259 nphi === 2 * session.cfg.nphi &&
260 fine.length === nlat * nphi,
261 `render ${nlat}×${nphi}, ${fine.length} values`,
262 );
264 const qlm = await session.read('U');
265 const expected = await session.viewSht.synth(qlm);
266 let maxDiff = 0;
267 for (let i = 0; i < fine.length; i++) {
268 const d = Math.abs(fine[i] - expected[i]);
269 if (d > maxDiff) maxDiff = d;
270 }
271 check(
272 'oversample: readSpecies matches synth of the read-back state',
273 maxDiff <= 1e-6,
274 `max |diff| = ${maxDiff.toExponential(2)}`,
275 );
277 // A timing burst must be invisible: the state is snapshotted and restored
278 // around it, and model time does not advance.
279 const tBefore = session.t;
280 const stepsBefore = session.steps;
281 const ms = await session.measure(8);
282 const after = await session.read('U');
283 let identical = qlm.length === after.length;
284 if (identical) {
285 for (let i = 0; i < qlm.length; i++) {
286 if (qlm[i] !== after[i]) {
287 identical = false;
288 break;
289 }
290 }
291 }
292 check(
293 'measure: a timing burst leaves state, t and steps untouched',
294 identical && session.t === tBefore && session.steps === stepsBefore,
295 identical
296 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
297 : 'state changed',
298 );
300 // Changing the oversampling in place is display-only: the state survives
301 // and the render grid drops back to the solver's.
302 await session.setOversample(1);
303 const qlmAfterSwap = await session.read('U');
304 let stateSurvived = qlmAfterSwap.length === after.length;
305 if (stateSurvived) {
306 for (let i = 0; i < after.length; i++) {
307 if (qlmAfterSwap[i] !== after[i]) {
308 stateSurvived = false;
309 break;
310 }
311 }
312 }
313 session.step(1); // recompute the view fields on the solver grid
314 const coarse = await session.readSpecies(0);
315 check(
316 'setOversample: swaps the render grid without touching the state',
317 stateSurvived &&
318 session.viewSht === session.sht &&
319 coarse.length === session.cfg.nlat * session.cfg.nphi,
320 stateSurvived
321 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
322 : 'state changed',
323 );
325 session.destroy();
326 }
327}