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