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 type { Check, Log } from './analyticChecks.ts';
16/** Kernels the step of each model should compile to — one per element-wise line. */
17const EXPECTED_KERNELS: Record<string, number> = {
18 schnakenberg: 5,
19 brusselator: 5,
20 allencahn: 2,
21};
23const LMAX = 31;
24const STEPS = 40;
26export async function modelChecks(
27 device: GPUDevice,
28 check: Check,
29 log: Log,
30): Promise<void> {
31 check('models: registry populated', mModels.length === 3, `${mModels.length} models`);
33 for (const model of mModels) {
34 const session = await ModelSession.create({
35 device,
36 model,
37 params: defaultParams(model),
38 lmax: LMAX,
39 });
41 const plan = session.describe();
42 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
43 const xforms = plan.step.filter(
44 (l) => l.startsWith('synth') || l.startsWith('analys'),
45 ).length;
46 log(
47 ` ${model.key}.m -> ${plan.step.length} ops/step ` +
48 `(${kernels} generated kernels, ${xforms} transforms)`,
49 );
50 check(
51 `${model.key}: element-wise lines fused into one kernel each`,
52 kernels === EXPECTED_KERNELS[model.key],
53 `${kernels} kernels (expected ${EXPECTED_KERNELS[model.key]})`,
54 );
56 session.seed(1);
57 session.step(STEPS);
59 // Every rendered field must be finite and have developed some contrast.
60 for (const field of model.species) {
61 const values = await session.read(field);
62 let lo = Infinity;
63 let hi = -Infinity;
64 let finite = true;
65 for (const v of values) {
66 if (!Number.isFinite(v)) finite = false;
67 if (v < lo) lo = v;
68 if (v > hi) hi = v;
69 }
70 check(
71 `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
72 finite && hi - lo > 1e-6,
73 finite
74 ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
75 : 'contains NaN or Infinity',
76 );
77 }
79 session.destroy();
80 }
82 // The oversampled readback: readSpecies must be the state synthesized on the
83 // display grid. Comparing against the display plan's own upload path
84 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
85 // coefficient copy against a known-good route through the same kernels.
86 {
87 const model = mModels.find((m) => m.key === 'allencahn')!;
88 const session = await ModelSession.create({
89 device,
90 model,
91 params: defaultParams(model),
92 lmax: LMAX,
93 oversample: 2,
94 });
95 session.seed(1);
96 session.step(STEPS);
98 const fine = await session.readSpecies(0);
99 const { nlat, nphi } = session.viewSht.cfg;
100 check(
101 'oversample: species field is on the 2x display grid',
102 nlat === 2 * session.cfg.nlat &&
103 nphi === 2 * session.cfg.nphi &&
104 fine.length === nlat * nphi,
105 `render ${nlat}×${nphi}, ${fine.length} values`,
106 );
108 const qlm = await session.read('U');
109 const expected = await session.viewSht.synth(qlm);
110 let maxDiff = 0;
111 for (let i = 0; i < fine.length; i++) {
112 const d = Math.abs(fine[i] - expected[i]);
113 if (d > maxDiff) maxDiff = d;
114 }
115 check(
116 'oversample: readSpecies matches synth of the read-back state',
117 maxDiff <= 1e-6,
118 `max |diff| = ${maxDiff.toExponential(2)}`,
119 );
121 // A timing burst must be invisible: the state is snapshotted and restored
122 // around it, and model time does not advance.
123 const tBefore = session.t;
124 const stepsBefore = session.steps;
125 const ms = await session.measure(8);
126 const after = await session.read('U');
127 let identical = qlm.length === after.length;
128 if (identical) {
129 for (let i = 0; i < qlm.length; i++) {
130 if (qlm[i] !== after[i]) {
131 identical = false;
132 break;
133 }
134 }
135 }
136 check(
137 'measure: a timing burst leaves state, t and steps untouched',
138 identical && session.t === tBefore && session.steps === stepsBefore,
139 identical
140 ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
141 : 'state changed',
142 );
144 // Changing the oversampling in place is display-only: the state survives
145 // and the render grid drops back to the solver's.
146 await session.setOversample(1);
147 const qlmAfterSwap = await session.read('U');
148 let stateSurvived = qlmAfterSwap.length === after.length;
149 if (stateSurvived) {
150 for (let i = 0; i < after.length; i++) {
151 if (qlmAfterSwap[i] !== after[i]) {
152 stateSurvived = false;
153 break;
154 }
155 }
156 }
157 session.step(1); // recompute the view fields on the solver grid
158 const coarse = await session.readSpecies(0);
159 check(
160 'setOversample: swaps the render grid without touching the state',
161 stateSurvived &&
162 session.viewSht === session.sht &&
163 coarse.length === session.cfg.nlat * session.cfg.nphi,
164 stateSurvived
165 ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
166 : 'state changed',
167 );
169 session.destroy();
170 }