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