1/**
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 {
23 parseArgs,
24 modelForSpec,
25 resolvePreset,
26 formatCommand,
27 BENCH_COMMAND,
28 DEFAULT_LMAX,
29 DEFAULT_SEED,
30 DEFAULT_STEPS,
31 DEFAULT_WARMUP,
32 type RunSpec,
33} from '../src/bench/runSpec.ts';
34import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
36const USAGE = `usage: ${BENCH_COMMAND} [options]
38 --preset <key> ${presets.map((p) => p.key).join(' | ')}
39 (default ${presets[0].key})
40 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
41 --steps <n> timed steps (default ${DEFAULT_STEPS})
42 --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
43 --seed <n> initial-noise seed (default ${DEFAULT_SEED})
44 --batch <n> steps per submit for the throughput number (default 16)
45 --<param> <v> any parameter of the preset's model, e.g. --dt 0.05
46 --json machine-readable output
47 --help
49The browser app shows the command for whatever it is currently simulating;
50copy it from under the stats line to compare the same run here.`;
52function fail(msg: string, code = 1): never {
53 console.error(`bench: ${msg}`);
54 process.exit(code);
55}
57// ---------------------------------------------------------------- arguments
58const argv = process.argv.slice(2);
59if (argv.includes('--help') || argv.includes('-h')) {
60 console.log(USAGE);
61 process.exit(0);
62}
63const wantJson = argv.includes('--json');
64let batch = 16;
65const rest: string[] = [];
66for (let i = 0; i < argv.length; i++) {
67 if (argv[i] === '--json') continue;
68 if (argv[i] === '--batch') {
69 batch = Number(argv[++i]);
70 continue;
71 }
72 if (argv[i].startsWith('--batch=')) {
73 batch = Number(argv[i].slice('--batch='.length));
74 continue;
75 }
76 rest.push(argv[i]);
77}
78if (!Number.isInteger(batch) || batch < 1) fail(`--batch must be an integer >= 1`, 2);
80let spec: RunSpec;
81try {
82 spec = parseArgs(rest);
83} catch (e) {
84 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
85}
87// ---------------------------------------------------------------- statistics
88interface Timing {
89 meanMs: number;
90 medianMs: number;
91 p05Ms: number;
92 p95Ms: number;
93 minMs: number;
94 totalMs: number;
95 stepsPerSec: number;
96}
98function timing(samples: Float64Array): Timing {
99 const sorted = Float64Array.from(samples).sort();
100 const q = (p: number): number =>
101 sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
102 let total = 0;
103 for (const v of samples) total += v;
104 const mean = total / samples.length;
105 return {
106 meanMs: mean,
107 medianMs: q(0.5),
108 p05Ms: q(0.05),
109 p95Ms: q(0.95),
110 minMs: sorted[0],
111 totalMs: total,
112 stepsPerSec: 1000 / mean,
113 };
114}
116function fieldRange(v: ArrayLike<number>): { min: number; max: number } {
117 let min = Infinity;
118 let max = -Infinity;
119 for (let i = 0; i < v.length; i++) {
120 if (v[i] < min) min = v[i];
121 if (v[i] > max) max = v[i];
122 }
123 return { min, max };
124}
126// ---------------------------------------------------------------- run
127const model = modelForSpec(spec);
128const { preset } = resolvePreset(spec.preset);
130let device: GPUDevice | null = null;
131let session: ModelSession | null = null;
133try {
134 const runtime = await installWebGpu();
135 device = await requestShtDevice().catch((e: unknown) => {
136 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
137 });
138 const adapter = await describeAdapter(device);
140 session = await ModelSession.create({
141 device,
142 model,
143 params: spec.params,
144 lmax: spec.lmax,
145 });
146 session.seed(spec.seed);
148 const plan = session.describe();
149 const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
150 const cfg = session.cfg;
152 if (!wantJson) {
153 console.log(`turing-sphere bench — solver only, no rendering\n`);
154 console.log(` preset ${preset.label} (models/${model.key}.m: ${model.species.join(', ')})`);
155 console.log(
156 ` params ${model.params.map((p) => `${p.key}=${spec.params[p.key]}`).join(' ')}`,
157 );
158 console.log(
159 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${session.sht.nlm.toLocaleString()}`,
160 );
161 console.log(` compiled ${plan.step.length} GPU ops/step (${kernels} generated kernels)`);
162 console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
163 console.log(` run ${spec.warmup} warmup + ${spec.steps} timed steps, seed ${spec.seed}\n`);
164 }
166 const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
168 session.step(spec.warmup);
169 await done();
171 // --- throughput: batches submitted together, awaited once each ---
172 const batches = Math.max(1, Math.ceil(spec.steps / batch));
173 const progress = !wantJson && process.stderr.isTTY;
174 let lastReport = performance.now();
175 const tp0 = performance.now();
176 let stepsRun = 0;
177 for (let b = 0; b < batches; b++) {
178 const n = Math.min(batch, spec.steps - stepsRun);
179 session.step(n);
180 await done();
181 stepsRun += n;
182 if (progress && performance.now() - lastReport > 1000) {
183 const so_far = (performance.now() - tp0) / stepsRun;
184 process.stderr.write(
185 `\r\x1b[K ${stepsRun}/${spec.steps} steps · ${so_far.toFixed(2)} ms/step`,
186 );
187 lastReport = performance.now();
188 }
189 }
190 const throughputMs = (performance.now() - tp0) / stepsRun;
191 if (progress) process.stderr.write('\r\x1b[K');
193 // --- latency: one step per submit, for the distribution ---
194 const latencySteps = Math.min(spec.steps, 200);
195 const samples = new Float64Array(latencySteps);
196 for (let s = 0; s < latencySteps; s++) {
197 const t0 = performance.now();
198 session.step(1);
199 await done();
200 samples[s] = performance.now() - t0;
201 }
202 const t = timing(samples);
204 const field = await session.read(model.species[0]);
205 const range = fieldRange(field);
206 let finite = true;
207 for (const v of field) if (!Number.isFinite(v)) finite = false;
209 if (wantJson) {
210 console.log(
211 JSON.stringify(
212 {
213 command: formatCommand(spec),
214 spec,
215 model: model.key,
216 backend: { adapter, runtime },
217 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm: session.sht.nlm },
218 compiled: { opsPerStep: plan.step.length, kernels },
219 throughput: { batch, msPerStep: throughputMs, stepsPerSec: 1000 / throughputMs },
220 latency: t,
221 state: {
222 t: session.t,
223 steps: session.steps,
224 species: model.species[0],
225 min: range.min,
226 max: range.max,
227 contrast: range.max - range.min,
228 finite,
229 },
230 },
231 null,
232 2,
233 ),
234 );
235 } else {
236 console.log(
237 ` ${throughputMs.toFixed(2)} ms/step ${(1000 / throughputMs).toFixed(1)} steps/s ` +
238 `${(spec.params.dt * (1000 / throughputMs)).toFixed(2)} model time/s` +
239 ` (batches of ${batch})`,
240 );
241 console.log(
242 ` one step per submit: ${t.meanMs.toFixed(2)} ms mean · median ${t.medianMs.toFixed(2)} · ` +
243 `p05 ${t.p05Ms.toFixed(2)} · p95 ${t.p95Ms.toFixed(2)} · min ${t.minMs.toFixed(2)}`,
244 );
245 console.log(
246 ` after ${session.steps} steps: t = ${session.t.toFixed(2)}, ` +
247 `${model.species[0]} ∈ [${range.min.toFixed(4)}, ${range.max.toFixed(4)}] ` +
248 `(contrast ${(range.max - range.min).toFixed(4)})${finite ? '' : ' — NOT FINITE'}`,
249 );
250 console.log(
251 `\n Compare with the ms/step in the app's stats line: same .m, same kernels,\n` +
252 ` but measured while the page renders the spheres.`,
253 );
254 }
256 session.destroy();
257 device.destroy();
258 process.exit(finite ? 0 : 1);
259} catch (e) {
260 session?.destroy();
261 device?.destroy();
262 fail(errMsg(e));
263}