1/**
2 * The two things this project adds to turing-sphere: a surface, and a `for`
3 * loop in the compiled step.
4 *
5 * The surface is checked against what it is supposed to be — the sphere really
6 * is the unit sphere and really is degree 1, a deformed shape really has the
7 * radius profile its .m says, and the coefficients really do evaluate to the
8 * same surface on a finer grid.
9 *
10 * The loop is checked for the property the whole design rests on: it is
11 * unrolled into the fixed op sequence, so more iterations means more GPU ops —
12 * and, while the geometry correction inside it is identically zero, the answer
13 * must be *bit for bit* independent of how many times it runs. That is a
14 * stronger statement than "close enough": if the placeholder were ever
15 * something that merely rounds to zero, or if the loop were miscompiled to
16 * read a stale buffer, these would differ in the last bits and this fails.
17 */
18import { ShtPlan } from '../src/sht/sht.ts';
19import { gridForLmax, lmIndex } from '../src/sht/layout.ts';
20import { ModelSession } from '../src/mgpu/session.ts';
21import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
22import { Geometry } from '../src/geom/geometry.ts';
23import {
24 mGeometries,
25 mGeometryByKey,
26 defaultGeometryParams,
27 SPHERE_KEY,
28} from '../src/geom/registry.ts';
29import { ModelCompileError } from '../src/mgpu/errors.ts';
30import type { Check, Log } from './analyticChecks.ts';
32const LMAX = 31;
33const STEPS = 20;
35/** Build one geometry on its own transform plan, for inspection. */
36async function buildGeometry(device: GPUDevice, key: string) {
37 const g = mGeometryByKey(key)!;
38 const { nlat, nphi } = gridForLmax(LMAX, 3);
39 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
40 const sht = await ShtPlan.create(device, cfg);
41 const geometry = await Geometry.create({
42 device,
43 sht,
44 cfg,
45 source: g.source,
46 paramNames: g.params.map((p) => p.key),
47 params: defaultGeometryParams(g),
48 });
49 return { g, sht, cfg, geometry };
50}
52export async function geometryChecks(
53 device: GPUDevice,
54 check: Check,
55 log: Log,
56): Promise<void> {
57 // ---- every geometry compiles and closes ---------------------------------
58 for (const spec of mGeometries) {
59 const { sht, geometry } = await buildGeometry(device, spec.key);
60 let finite = true;
61 for (const a of [geometry.x, geometry.y, geometry.z]) {
62 for (const v of a) if (!Number.isFinite(v)) finite = false;
63 }
64 const { lo, hi } = geometry.radiusRange();
65 check(
66 `geometry: ${spec.key}.m evaluates to a finite surface`,
67 finite && lo > 1e-3,
68 `radius ${lo.toFixed(4)}–${hi.toFixed(4)}`,
69 );
70 sht.destroy();
71 }
73 // ---- the sphere is the unit sphere, exactly, and is degree 1 ------------
74 {
75 const { sht, geometry } = await buildGeometry(device, SPHERE_KEY);
77 let maxRadiusErr = 0;
78 for (let i = 0; i < geometry.x.length; i++) {
79 const r = Math.hypot(geometry.x[i], geometry.y[i], geometry.z[i]);
80 maxRadiusErr = Math.max(maxRadiusErr, Math.abs(r - 1));
81 }
82 // Tolerance is fp32 through a full analysis/synthesis round trip, not the
83 // geometry: the exact answer is representable, and what is measured here
84 // is the transforms' own round-off. It is set by the loosest stack this
85 // runs on — SwiftShader in CI is an order of magnitude worse than Dawn on
86 // real hardware (4e-4 against 2e-5). A geometry that was actually wrong
87 // would miss by O(1), so the slack costs nothing.
88 check(
89 'geometry: sphere.m has radius 1 everywhere',
90 maxRadiusErr < 2e-3,
91 `max |r - 1| = ${maxRadiusErr.toExponential(2)}`,
92 );
94 // x, y, z of the unit sphere are the three degree-1 harmonics and nothing
95 // else, so analysing them must leave every other coefficient at zero.
96 // This is what makes the sphere case exact rather than merely accurate:
97 // there is no content for the band limit to throw away.
98 const degreeOne = new Set([
99 lmIndex(LMAX, 1, 0),
100 lmIndex(LMAX, 1, 1),
101 ]);
102 let leak = 0;
103 for (const coeffs of [geometry.X, geometry.Y, geometry.Z]) {
104 for (let i = 0; i < coeffs.length / 2; i++) {
105 if (degreeOne.has(i)) continue;
106 leak = Math.max(leak, Math.abs(coeffs[2 * i]), Math.abs(coeffs[2 * i + 1]));
107 }
108 }
109 check(
110 'geometry: sphere.m is exactly degree 1 in the harmonics',
111 leak < 1e-3,
112 `max |coefficient| outside l = 1 is ${leak.toExponential(2)}`,
113 );
114 sht.destroy();
115 }
117 // ---- a deformed surface matches its own formula, on any grid ------------
118 {
119 const { g, sht, cfg, geometry } = await buildGeometry(device, 'peanut');
120 const p = defaultGeometryParams(g);
122 // peanut.m written out: r = 1 - waist*sin(theta)^2 scales the unit sphere,
123 // and z is then stretched, so the distance from the origin depends on
124 // theta alone. Checking every point against this closed form checks the
125 // whole path at once — the compiled shape kernel, the analysis into
126 // coefficients, the synthesis back — and, because the formula has no phi
127 // in it, that the surface really is a surface of revolution.
128 const peanutRadius = (ct: number): number => {
129 const st2 = Math.max(0, 1 - ct * ct);
130 const r = 1 - p.waist * st2;
131 return r * Math.hypot(Math.sqrt(st2), (1 + p.stretch) * ct);
132 };
134 const onGrid = (
135 cosTheta: Float64Array,
136 nlat: number,
137 nphi: number,
138 at: (i: number) => number,
139 ): number => {
140 let worst = 0;
141 for (let i = 0; i < nlat; i++) {
142 const want = peanutRadius(cosTheta[i]);
143 for (let j = 0; j < nphi; j++) {
144 worst = Math.max(worst, Math.abs(at(i * nphi + j) - want));
145 }
146 }
147 return worst;
148 };
150 const coarse = onGrid(sht.cosTheta, cfg.nlat, cfg.nphi, (k) =>
151 Math.hypot(geometry.x[k], geometry.y[k], geometry.z[k]),
152 );
153 check(
154 'geometry: peanut.m matches its own radial formula on the solver grid',
155 coarse < 1e-3,
156 `max |dr| = ${coarse.toExponential(2)}`,
157 );
159 // And the same on a finer grid, from the same coefficients. This is what
160 // "the rendered surface is the surface being solved on" means: display
161 // oversampling evaluates the embedding at more points, it does not
162 // subdivide or smooth it. The 2x Gauss latitudes share no point with the
163 // 1x ones, so agreeing here is agreeing everywhere, not at samples.
164 const fine = await ShtPlan.create(device, {
165 lmax: cfg.lmax,
166 mmax: cfg.mmax,
167 nlat: 2 * cfg.nlat,
168 nphi: 2 * cfg.nphi,
169 });
170 const finePos = await geometry.positionsOn(fine);
171 const refined = onGrid(fine.cosTheta, 2 * cfg.nlat, 2 * cfg.nphi, (k) =>
172 Math.hypot(finePos[3 * k], finePos[3 * k + 1], finePos[3 * k + 2]),
173 );
174 check(
175 'geometry: the same coefficients give the same surface on a 2x grid',
176 refined < 1e-3,
177 `max |dr| = ${refined.toExponential(2)} at ${2 * cfg.nlat}×${2 * cfg.nphi} points`,
178 );
179 fine.destroy();
180 sht.destroy();
181 }
183 // ---- the unrolled loop: more ops, identical answer ----------------------
184 {
185 const model = mModelByKey('schnakenberg')!;
186 const params = defaultParams(model);
187 const counts = [0, 1, 4];
188 const ops: number[] = [];
189 const states: Float32Array[] = [];
191 for (const niter of counts) {
192 const session = await ModelSession.create({
193 device, model, params, lmax: LMAX, niter,
194 });
195 ops.push(session.describe().step.length);
196 session.seed(1);
197 session.step(STEPS);
198 states.push(await session.read('U'));
199 session.destroy();
200 }
202 log(` schnakenberg.m ops/step by solve iterations: ${
203 counts.map((n, i) => `${n} -> ${ops[i]}`).join(', ')
204 }`);
205 check(
206 'loop: each solve iteration adds GPU operations',
207 ops[0] < ops[1] && ops[1] < ops[2],
208 `${ops.join(' < ')} ops for ${counts.join(', ')} iterations`,
209 );
210 // Unrolling has to be exactly linear in the trip count: the body planned
211 // once per iteration, no more and no less. Two dispatches per species per
212 // iteration — the placeholder line and the update that reads it.
213 const perIteration = ops[1] - ops[0];
214 const want = 2 * model.species.length;
215 check(
216 'loop: unrolling is exactly linear in the trip count',
217 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
218 `${perIteration} ops per iteration (expected ${want}), ` +
219 `${ops[2] - ops[0]} for 4 iterations`,
220 );
222 let identical = true;
223 let worst = 0;
224 for (let k = 1; k < states.length; k++) {
225 if (states[k].length !== states[0].length) identical = false;
226 for (let i = 0; i < states[0].length; i++) {
227 if (states[k][i] !== states[0][i]) identical = false;
228 worst = Math.max(worst, Math.abs(states[k][i] - states[0][i]));
229 }
230 }
231 check(
232 'loop: the geometry correction is exactly zero, so the answer does not move',
233 identical,
234 identical
235 ? `bit-identical after ${STEPS} steps at ${counts.join('/')} iterations`
236 : `states differ by up to ${worst.toExponential(2)}`,
237 );
238 }
240 // ---- a loop whose length is not known at compile time is refused --------
241 {
242 const model = mModelByKey('allencahn')!;
243 // `dt` is a tunable parameter, so it reaches the compiler with no value:
244 // the plan cannot know how many iterations to emit.
245 const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
246 let message = '';
247 try {
248 const session = await ModelSession.create({
249 device, model, params: defaultParams(model), lmax: LMAX, source: bad, niter: 1,
250 });
251 session.destroy();
252 } catch (e) {
253 message = e instanceof ModelCompileError ? e.message : `wrong error type: ${e}`;
254 }
255 check(
256 'loop: a runtime loop bound is refused at compile time',
257 message.includes('known when the model is compiled'),
258 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
259 );
260 }
262 // ---- swapping the surface leaves the simulation alone ------------------
263 {
264 const model = mModelByKey('schnakenberg')!;
265 const session = await ModelSession.create({
266 device, model, params: defaultParams(model), lmax: LMAX,
267 });
268 session.seed(1);
269 session.step(STEPS);
270 const before = await session.read('U');
272 const peanut = mGeometryByKey('peanut')!;
273 await session.setGeometry(peanut, defaultGeometryParams(peanut));
274 const after = await session.read('U');
276 let survived = before.length === after.length;
277 for (let i = 0; survived && i < before.length; i++) {
278 if (before[i] !== after[i]) survived = false;
279 }
280 const { lo, hi } = session.geometry.radiusRange();
281 check(
282 'geometry: swapping the surface mid-run does not disturb the state',
283 survived && session.geometryModel.key === 'peanut' && hi - lo > 0.1,
284 survived
285 ? `state identical, now on ${session.geometryModel.key} (radius ${lo.toFixed(3)}–${hi.toFixed(3)})`
286 : 'state changed',
287 );
288 session.destroy();
289 }
290}
292/** Index of the entry minimizing `score`, over the first `n` entries. */
293function argMin(
294 xs: Float64Array | Float32Array,
295 n: number,
296 score: (v: number) => number,
297): number {
298 let best = 0;
299 let bestScore = Infinity;
300 for (let i = 0; i < n; i++) {
301 const s = score(xs[i]);
302 if (s < bestScore) {
303 bestScore = s;
304 best = i;
305 }
306 }
307 return best;
308}