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