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: 8,
29 brusselator: 8,
30 allencahn: 4,
31 'schnakenberg-alg4': 8,
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: 14,
51 brusselator: 14,
52 allencahn: 7,
53 // 30 before the correction gained its band projection (.* filt on dLu):
54 // that line fused into the state update in this model's expression shape,
55 // and no longer does — one extra 2 x nlm kernel per species per iteration.
56 'schnakenberg-alg4': 32,
57};
59const LMAX = 31;
60const STEPS = 40;
61const NITER = 1;
63export async function modelChecks(
64 device: GPUDevice,
65 check: Check,
66 log: Log,
67): Promise<void> {
68 check('models: registry populated', mModels.length === 4, `${mModels.length} models`);
70 // The app formats the run it is showing into a `npm run bench` command and
71 // the benchmark parses it back. That is only worth anything if the round
72 // trip is lossless — a knob that formatCommand forgets is a knob the desktop
73 // run would silently take a default for, and the two runs would differ while
74 // claiming to be the same. Every field of the spec, through both directions.
75 {
76 const spec: RunSpec = {
77 preset: 'schnak-fine',
78 lmax: 127,
79 seed: 12345,
80 steps: 777,
81 warmup: 13,
82 params: { a: 0.11, b: 0.91, D1: 5e-4, D2: 9e-3, dt: 0.04 },
83 geometry: 'peanut',
84 geometryParams: { waist: 0.45, stretch: 1.25 },
85 niter: 3,
86 };
87 const command = formatCommand(spec);
88 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
89 const same = JSON.stringify(back) === JSON.stringify(spec);
90 check(
91 'runSpec: the benchmark command round-trips every field',
92 same,
93 same ? command.slice(BENCH_COMMAND.length + 1) : `got ${JSON.stringify(back)}`,
94 );
95 }
97 for (const model of mModels) {
98 const session = await ModelSession.create({
99 device,
100 model,
101 params: defaultParams(model),
102 lmax: LMAX,
103 niter: NITER,
104 });
106 const plan = session.describe();
107 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
108 const xforms = plan.step.filter(
109 (l) => l.startsWith('synth') || l.startsWith('analys'),
110 ).length;
111 const expected = EXPECTED_KERNELS[model.key] + NITER * KERNELS_PER_ITERATION[model.key];
112 log(
113 ` ${model.key}.m -> ${plan.step.length} ops/step ` +
114 `(${kernels} generated kernels, ${xforms} transforms, ${NITER} solve iter)`,
115 );
116 check(
117 `${model.key}: element-wise lines fused into one kernel each`,
118 kernels === expected,
119 `${kernels} kernels (expected ${expected})`,
120 );
122 await session.seed(1);
123 session.step(STEPS);
125 // Every rendered field must be finite and have developed some contrast.
126 for (const field of model.species) {
127 const values = await session.read(field);
128 let lo = Infinity;
129 let hi = -Infinity;
130 let finite = true;
131 for (const v of values) {
132 if (!Number.isFinite(v)) finite = false;
133 if (v < lo) lo = v;
134 if (v > hi) hi = v;
135 }
136 check(
137 `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
138 finite && hi - lo > 1e-6,
139 finite
140 ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
141 : 'contains NaN or Infinity',
142 );
143 }
145 session.destroy();
146 }
148 // Batched transforms are an encoding of the same arithmetic, so a run with
149 // batching disabled (SHT_BATCH=0 compiles scalar-only plans) must reproduce
150 // the default run to shader-compiler latitude, and the default run must
151 // actually be batching (the describe() lines say so). This is the guard
152 // that the planner's adjacency grouping rewires buffers correctly — a lane
153 // bound to the wrong field would miss by O(1), not O(1e-6).
154 {
155 const model = mModelByKey('schnakenberg')!;
156 const params = defaultParams(model);
157 const states: Float32Array[] = [];
158 let batchedLanes = 0;
159 for (const batch of [undefined, 0]) {
160 const g = globalThis as Record<string, unknown>;
161 if (batch !== undefined) g.SHT_BATCH = batch;
162 try {
163 const session = await ModelSession.create({
164 device, model, params, lmax: LMAX, niter: NITER,
165 });
166 if (batch === undefined) {
167 batchedLanes = session
168 .describe()
169 .step.filter((l) => l.includes('[batch lane')).length;
170 }
171 await session.seed(1);
172 session.step(STEPS);
173 states.push(await session.read('U'));
174 session.destroy();
175 } finally {
176 delete g.SHT_BATCH;
177 }
178 }
179 // Every batchable run at one solve iteration: the u/v syntheses and the
180 // reaction analyses outside the loop (2 + 2), the four gradient
181 // syntheses, two theta-flux analyses, two divergence syntheses and two
182 // final analyses inside it (4 + 2 + 2 + 2; the phi flux goes through
183 // dphig, which has no Legendre stage to batch). Lane counts are
184 // batch-width invariant: a x4 run is one batch at K = 4 and two at
185 // K = 2, but the lanes annotated are the same 14 either way.
186 check(
187 'batch: the compiled step batches every adjacent transform pair',
188 batchedLanes === 14,
189 `${batchedLanes} batched transform lanes (expected 14)`,
190 );
191 let worst = 0;
192 for (let i = 0; i < states[0].length; i++) {
193 worst = Math.max(worst, Math.abs(states[0][i] - states[1][i]));
194 }
195 check(
196 'batch: batched and scalar plans agree through a real run',
197 worst < 1e-4,
198 `max |U_batched - U_scalar| = ${worst.toExponential(2)} after ${STEPS} steps`,
199 );
200 }
202 // Misusing the grouped-transform syntax is refused at compile time with a
203 // message that says how to write it, not silently mis-planned: every input
204 // must get an output (each one costs a transform), whether the mismatch is
205 // an under-bound assignment or an ignored slot.
206 {
207 const model = mModelByKey('allencahn')!;
208 const cases: [string, string, string][] = [
209 [
210 'a single output bound to a grouped call',
211 'Ftu = synth(vtu, vpu);',
212 'bind each one',
213 ],
214 [
215 'an ignored output slot',
216 // Fpu is reassigned so the only error left is the dropped slot
217 // itself, which the planner refuses (numbl would otherwise catch
218 // the undefined 'Fpu' first, masking the check under test).
219 '[Ftu, ~] = synth(vtu, vpu);\n Fpu = Ftu;',
220 'must be bound',
221 ],
222 ];
223 for (const [what, bad, expect] of cases) {
224 const source = model.source.replace('[Ftu, Fpu] = synth(vtu, vpu);', bad);
225 let message = '';
226 try {
227 const session = await ModelSession.create({
228 device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
229 });
230 session.destroy();
231 } catch (e) {
232 message = e instanceof Error ? e.message : String(e);
233 }
234 check(
235 `batch: ${what} is refused at compile time`,
236 message.includes(expect),
237 message ? `refused: ${message.slice(0, 76)}…` : 'compiled anyway',
238 );
239 }
240 }
242 // The oversampled readback: readSpecies must be the state synthesized on the
243 // display grid. Comparing against the display plan's own upload path
244 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
245 // coefficient copy against a known-good route through the same kernels.
246 {
247 const model = mModels.find((m) => m.key === 'allencahn')!;
248 const session = await ModelSession.create({
249 device,
250 model,
251 params: defaultParams(model),
252 lmax: LMAX,
253 oversample: 2,
254 });
255 await session.seed(1);
256 session.step(STEPS);
258 const fine = await session.readSpecies(0);
259 const { nlat, nphi } = session.viewSht.cfg;
260 check(
261 'oversample: species field is on the 2x display grid',
262 nlat === 2 * session.cfg.nlat &&
263 nphi === 2 * session.cfg.nphi &&
264 fine.length === nlat * nphi,
265 `render ${nlat}×${nphi}, ${fine.length} values`,
266 );
268 const qlm = await session.read('U');
269 const expected = await session.viewSht.synth(qlm);
270 let maxDiff = 0;
271 for (let i = 0; i < fine.length; i++) {
272 const d = Math.abs(fine[i] - expected[i]);
273 if (d > maxDiff) maxDiff = d;
274 }
275 check(
276 'oversample: readSpecies matches synth of the read-back state',
277 maxDiff <= 1e-6,
278 `max |diff| = ${maxDiff.toExponential(2)}`,
279 );
281 // A timing burst must be invisible: the state is snapshotted and restored
282 // around it, and model time does not advance.
283 const tBefore = session.t;
284 const stepsBefore = session.steps;
285 const ms = await session.measure(8);
286 const after = await session.read('U');
287 let identical = qlm.length === after.length;
288 if (identical) {
289 for (let i = 0; i < qlm.length; i++) {
290 if (qlm[i] !== after[i]) {
291 identical = false;
292 break;
293 }
294 }
295 }
296 check(
297 'measure: a timing burst leaves state, t and steps untouched',
298 identical && session.t === tBefore && session.steps === stepsBefore,
299 identical
300 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
301 : 'state changed',
302 );
304 // Changing the oversampling in place is display-only: the state survives
305 // and the render grid drops back to the solver's.
306 await session.setOversample(1);
307 const qlmAfterSwap = await session.read('U');
308 let stateSurvived = qlmAfterSwap.length === after.length;
309 if (stateSurvived) {
310 for (let i = 0; i < after.length; i++) {
311 if (qlmAfterSwap[i] !== after[i]) {
312 stateSurvived = false;
313 break;
314 }
315 }
316 }
317 session.step(1); // recompute the view fields on the solver grid
318 const coarse = await session.readSpecies(0);
319 check(
320 'setOversample: swaps the render grid without touching the state',
321 stateSurvived &&
322 session.viewSht === session.sht &&
323 coarse.length === session.cfg.nlat * session.cfg.nphi,
324 stateSurvived
325 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
326 : 'state changed',
327 );
329 session.destroy();
330 }
331}