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