2 * Command-line benchmark: run exactly what the browser runs — the same .m
3 * models, lowered by numbl and compiled to the same WGSL kernels, over the same
4 * transforms — on desktop WebGPU (Google Dawn, via the optional `webgpu`
5 * package), and report ms/step. The app prints the matching command under its
6 * stats line; copy it and run it here for an apples-to-apples comparison.
7 *
8 * npm run bench -- --preset schnak-spots --lmax 63 --steps 2000 --seed 1 \
9 * --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05
10 *
11 * The only thing missing here is the rendering: this is the solver alone.
12 *
13 * Two numbers are reported, because they answer different questions:
14 * - throughput: a batch of steps submitted together, awaited once. This is how
15 * the app runs, and what keeping the state in GPU buffers is for.
16 * - latency: one step per submit, each awaited. Comparable to a design that
17 * reads back every step, and the only way to get a per-step distribution.
18 */
19import { requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
20import { ModelSession } from '../src/mgpu/session.ts';
21import { presets } from '../src/mgpu/registry.ts';
22import { mGeometries, DEFAULT_GEOMETRY_KEY } from '../src/geom/registry.ts';
1015122Editable operator/solver files in the page, and a solver selectorJeremy Magland 23import { DEFAULT_SOLVER, solverKeys } from '../src/mgpu/libs.ts';
25 parseArgs,
26 modelForSpec,
27 resolvePreset,
28 geometryForSpec,
29 formatCommand,
30 BENCH_COMMAND,
31 DEFAULT_LMAX,
32 DEFAULT_NITER,
33 DEFAULT_SEED,
34 DEFAULT_STEPS,
35 DEFAULT_WARMUP,
36 type RunSpec,
37} from '../src/bench/runSpec.ts';
38import { digestOf, formatDigest } from '../src/mgpu/digest.ts';
39import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
40import { writeFileSync } from 'node:fs';
42const USAGE = `usage: ${BENCH_COMMAND} [options]
44 --preset <key> ${presets.map((p) => p.key).join(' | ')}
45 (default ${presets[0].key})
46 --geometry <key> ${mGeometries.map((g) => g.key).join(' | ')}
47 (default ${DEFAULT_GEOMETRY_KEY})
48 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
49 --niter <n> iterations of the implicit solve, unrolled into the compiled
50 step (default ${DEFAULT_NITER})
1015122Editable operator/solver files in the page, and a solver selectorJeremy Magland 51 --solver <key> ${solverKeys.join(' | ')} — which solver answers the
52 models' solve(...) call (default ${DEFAULT_SOLVER})
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 53 --steps <n> timed steps (default ${DEFAULT_STEPS})
54 --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
55 --seed <n> initial-noise seed (default ${DEFAULT_SEED})
56 --batch <n> steps per submit for the throughput number (default 16)
57 --digest after timing, re-run exactly --steps steps from the seed and
58 print a digest of the final state
59 --dump-state <f> like --digest, and write the state to <f> as JSON, for
60 scripts/compare-env.mjs to compare against a browser run
61 --<param> <v> any parameter of the preset's model, e.g. --dt 0.05
62 --g<param> <v> any parameter of the geometry, e.g. --gwaist 0.6
63 --json machine-readable output
64 --help
66The browser app shows the command for whatever it is currently simulating;
67copy it from under the stats line to compare the same run here.`;
69function fail(msg: string, code = 1): never {
70 console.error(`bench: ${msg}`);
71 process.exit(code);
72}
74// ---------------------------------------------------------------- arguments
75const argv = process.argv.slice(2);
76if (argv.includes('--help') || argv.includes('-h')) {
77 console.log(USAGE);
78 process.exit(0);
79}
80const wantJson = argv.includes('--json');
81let batch = 16;
82let dumpState: string | null = null;
83let wantDigest = false;
84const rest: string[] = [];
85for (let i = 0; i < argv.length; i++) {
86 const a = argv[i];
87 if (a === '--json') continue;
88 if (a === '--digest') {
89 wantDigest = true;
90 continue;
91 }
92 const valued = (name: string): string | null => {
93 if (a === `--${name}`) return argv[++i];
94 if (a.startsWith(`--${name}=`)) return a.slice(name.length + 3);
95 return null;
96 };
97 const b = valued('batch');
98 if (b !== null) {
99 batch = Number(b);
100 continue;
101 }
102 const d = valued('dump-state');
103 if (d !== null) {
104 dumpState = d;
105 wantDigest = true;
106 continue;
107 }
108 rest.push(a);
109}
110if (!Number.isInteger(batch) || batch < 1) fail(`--batch must be an integer >= 1`, 2);
112let spec: RunSpec;
113try {
114 spec = parseArgs(rest);
115} catch (e) {
116 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
117}
119// ---------------------------------------------------------------- statistics
120interface Timing {
121 meanMs: number;
122 medianMs: number;
123 p05Ms: number;
124 p95Ms: number;
125 minMs: number;
126 totalMs: number;
127 stepsPerSec: number;
128}
130function timing(samples: Float64Array): Timing {
131 const sorted = Float64Array.from(samples).sort();
132 const q = (p: number): number =>
133 sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
134 let total = 0;
135 for (const v of samples) total += v;
136 const mean = total / samples.length;
137 return {
138 meanMs: mean,
139 medianMs: q(0.5),
140 p05Ms: q(0.05),
141 p95Ms: q(0.95),
142 minMs: sorted[0],
143 totalMs: total,
144 stepsPerSec: 1000 / mean,
145 };
146}
148function fieldRange(v: ArrayLike<number>): { min: number; max: number } {
149 let min = Infinity;
150 let max = -Infinity;
151 for (let i = 0; i < v.length; i++) {
152 if (v[i] < min) min = v[i];
153 if (v[i] > max) max = v[i];
154 }
155 return { min, max };
156}
158// ---------------------------------------------------------------- run
159const model = modelForSpec(spec);
160const { preset } = resolvePreset(spec.preset);
161const geometry = geometryForSpec(spec);
163let device: GPUDevice | null = null;
164let session: ModelSession | null = null;
166try {
167 const runtime = await installWebGpu();
168 device = await requestShtDevice().catch((e: unknown) => {
169 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
170 });
171 const adapter = await describeAdapter(device);
173 session = await ModelSession.create({
174 device,
175 model,
176 params: spec.params,
177 lmax: spec.lmax,
178 geometry,
179 geometryParams: spec.geometryParams,
180 niter: spec.niter,
1015122Editable operator/solver files in the page, and a solver selectorJeremy Magland 181 solver: spec.solver,
183 session.seed(spec.seed);
185 const plan = session.describe();
186 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
187 const cfg = session.cfg;
189 if (!wantJson) {
190 console.log(`turing-surface bench — solver only, no rendering\n`);
191 console.log(` preset ${preset.label} (models/${model.key}.m: ${model.species.join(', ')})`);
192 console.log(
193 ` params ${model.params.map((p) => `${p.key}=${spec.params[p.key]}`).join(' ')}`,
194 );
195 const radius = session.geometry.radiusRange();
196 console.log(
197 ` geometry ${geometry.label} (geometries/${geometry.key}.m` +
198 (geometry.params.length
199 ? `: ${geometry.params.map((p) => `${p.key}=${spec.geometryParams[p.key]}`).join(' ')})`
200 : ')') +
201 ` radius ${radius.lo.toFixed(3)}–${radius.hi.toFixed(3)}`,
202 );
203 console.log(
204 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
1015122Editable operator/solver files in the page, and a solver selectorJeremy Magland 205 `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'} of ${spec.solver}`,
207 console.log(` compiled ${plan.step.length} GPU ops/step (${kernels} generated kernels)`);
208 console.log(` fourier ${session.sht.fourierMode.toUpperCase()} stage`);
209 console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
210 console.log(` run ${spec.warmup} warmup + ${spec.steps} timed steps, seed ${spec.seed}\n`);
211 }
213 const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
215 session.step(spec.warmup);
216 await done();
218 // --- throughput: batches submitted together, awaited once each ---
219 const batches = Math.max(1, Math.ceil(spec.steps / batch));
220 const progress = !wantJson && process.stderr.isTTY;
221 let lastReport = performance.now();
222 const tp0 = performance.now();
223 let stepsRun = 0;
224 let encodeMs = 0;
225 for (let b = 0; b < batches; b++) {
226 const n = Math.min(batch, spec.steps - stepsRun);
227 const e0 = performance.now();
228 session.step(n);
229 encodeMs += performance.now() - e0;
230 await done();
231 stepsRun += n;
232 if (progress && performance.now() - lastReport > 1000) {
233 const so_far = (performance.now() - tp0) / stepsRun;
234 process.stderr.write(
235 `\r\x1b[K ${stepsRun}/${spec.steps} steps · ${so_far.toFixed(2)} ms/step`,
236 );
237 lastReport = performance.now();
238 }
239 }
240 const throughputMs = (performance.now() - tp0) / stepsRun;
241 const encodePerStep = encodeMs / stepsRun;
242 if (progress) process.stderr.write('\r\x1b[K');
244 // --- latency: one step per submit, for the distribution ---
245 const latencySteps = Math.min(spec.steps, 200);
246 const samples = new Float64Array(latencySteps);
247 for (let s = 0; s < latencySteps; s++) {
248 const t0 = performance.now();
249 session.step(1);
250 await done();
251 samples[s] = performance.now() - t0;
252 }
253 const t = timing(samples);
255 const field = await session.read(model.species[0]);
256 const range = fieldRange(field);
257 let finite = true;
258 for (const v of field) if (!Number.isFinite(v)) finite = false;
260 // A reproducible state to compare across machines: exactly `--steps` steps
261 // from the seed, separate from the timed runs above (which step a different
262 // number of times to measure throughput and latency).
263 let digest = null;
264 let state: Float32Array | null = null;
265 if (wantDigest) {
266 session.seed(spec.seed);
267 session.step(spec.steps);
268 await done();
269 state = await session.read(model.state[0]);
270 digest = digestOf(state, session.sht.fourierMode, adapter);
271 }
273 if (wantJson) {
274 console.log(
275 JSON.stringify(
276 {
277 command: formatCommand(spec),
278 spec,
279 model: model.key,
280 backend: { adapter, runtime, precision: 'fp32' },
281 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm: session.sht.nlm },
282 compiled: { opsPerStep: plan.step.length, kernels },
283 digest,
284 throughput: {
285 batch,
286 msPerStep: throughputMs,
287 stepsPerSec: 1000 / throughputMs,
288 encodeMsPerStep: encodePerStep,
289 },
290 latency: t,
291 state: {
292 t: session.t,
293 steps: session.steps,
294 species: model.species[0],
295 min: range.min,
296 max: range.max,
297 contrast: range.max - range.min,
298 finite,
299 },
300 },
301 null,
302 2,
303 ),
304 );
305 } else {
306 console.log(
307 ` ${throughputMs.toFixed(2)} ms/step ${(1000 / throughputMs).toFixed(1)} steps/s ` +
308 `${(spec.params.dt * (1000 / throughputMs)).toFixed(2)} model time/s` +
309 ` (batches of ${batch})`,
310 );
311 console.log(
312 ` of which CPU command encoding: ${encodePerStep.toFixed(3)} ms/step ` +
313 `(${((100 * encodePerStep) / throughputMs).toFixed(0)}% — the rest is the GPU)`,
314 );
315 console.log(
316 ` one step per submit: ${t.meanMs.toFixed(2)} ms mean · median ${t.medianMs.toFixed(2)} · ` +
317 `p05 ${t.p05Ms.toFixed(2)} · p95 ${t.p95Ms.toFixed(2)} · min ${t.minMs.toFixed(2)}`,
318 );
319 console.log(
320 ` after ${session.steps} steps: t = ${session.t.toFixed(2)}, ` +
321 `${model.species[0]} ∈ [${range.min.toFixed(4)}, ${range.max.toFixed(4)}] ` +
322 `(contrast ${(range.max - range.min).toFixed(4)})${finite ? '' : ' — NOT FINITE'}`,
323 );
324 if (digest) {
325 console.log(`\n state after ${spec.steps} steps from seed ${spec.seed}:`);
326 console.log(` ${formatDigest(digest)}`);
327 }
328 console.log(
329 `\n The app's stats line reports the same solver number (batched steps,\n` +
330 ` nothing read back) plus a separate ms/frame that carries the readback\n` +
331 ` and the rendering. Compare solver with solver.`,
332 );
333 }
335 if (dumpState && state && digest) {
336 writeFileSync(
337 dumpState,
338 JSON.stringify({ command: formatCommand(spec), spec, digest, state: [...state] }),
339 );
340 if (!wantJson) console.log(`\n wrote ${dumpState}`);
341 }
343 session.destroy();
344 device.destroy();
345 process.exit(finite ? 0 : 1);
346} catch (e) {
347 session?.destroy();
348 device?.destroy();
349 fail(errMsg(e));
350}