/ concept-collection / turing-sphere-2
Sign in
concept-collection / turing-sphere-2
turing-sphere-2 / scripts / bench.ts
270 lines · 8.6 KBBlameHistoryRaw
1/**
2 * Command-line benchmark: run exactly the simulation the browser app is
3 * running — same solver, same transforms, same parameters — on desktop WebGPU
4 * (Google Dawn, via the optional `webgpu` package) or on the f64 CPU
5 * reference, and report ms/step. The app prints the matching command under
6 * its stats line; copy it and run it here for an apples-to-apples comparison.
7 *
8 * node scripts/bench.ts --preset schnak-spots --lmax 63 --backend webgpu \
9 * --steps 500 --seed 1 --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 */
13import {
14 GpuBackend,
15 CpuBackend,
16 requestShtDevice,
17 describeAdapter,
18 type ShtBackend,
19} from '../src/solver/backend.ts';
20import { Simulation } from '../src/solver/simulation.ts';
21import { presets } from '../src/solver/models.ts';
22import {
23 parseArgs,
24 modelForSpec,
25 configForSpec,
26 resolvePreset,
27 formatCommand,
28 BENCH_COMMAND,
29 DEFAULT_LMAX,
30 DEFAULT_SEED,
31 DEFAULT_STEPS,
32 DEFAULT_WARMUP,
33 DEFAULT_BACKEND,
34 type RunSpec,
35} from '../src/bench/runSpec.ts';
37const USAGE = `usage: ${BENCH_COMMAND} [options]
39 --preset <key> ${presets.map((p) => p.key).join(' | ')}
40 (default ${presets[0].key})
41 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
42 --backend <kind> webgpu | cpu (default ${DEFAULT_BACKEND})
43 --steps <n> timed steps (default ${DEFAULT_STEPS})
44 --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
45 --seed <n> initial-noise seed (default ${DEFAULT_SEED})
46 --<param> <v> any parameter of the preset's model, e.g. --dt 0.05
47 --json machine-readable output
48 --help
50The browser app shows the command for whatever it is currently simulating;
51copy it from under the stats line to compare the same run here.`;
53function fail(msg: string, code = 1): never {
54 console.error(`bench: ${msg}`);
55 process.exit(code);
57const errMsg = (e: unknown): string => (e instanceof Error ? e.message : String(e));
59// ---------------------------------------------------------------- arguments
60const argv = process.argv.slice(2);
61if (argv.includes('--help') || argv.includes('-h')) {
62 console.log(USAGE);
63 process.exit(0);
65const wantJson = argv.includes('--json');
66let spec: RunSpec;
67try {
68 spec = parseArgs(argv.filter((a) => a !== '--json'));
69} catch (e) {
70 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
73// ---------------------------------------------------------------- WebGPU
74/**
75 * Install Dawn under the globals the transform code expects (navigator.gpu,
76 * GPUBufferUsage, ...), so src/ runs here unchanged — including
77 * requestShtDevice(), which is the same device request the browser makes.
78 * The specifier is indirect so that typechecking does not require the
79 * optional package to be installed.
80 */
81async function installWebGpu(): Promise<string> {
82 const specifier = 'webgpu';
83 let mod: {
84 create: (flags: string[]) => GPU;
85 globals: Record<string, unknown>;
86 };
87 try {
88 mod = await import(specifier);
89 } catch {
90 throw new Error(
91 'desktop WebGPU needs the optional `webgpu` package (prebuilt Google Dawn):\n' +
92 ' npm install webgpu\n' +
93 'or run with --backend cpu.',
94 );
95 }
96 Object.assign(globalThis, mod.globals);
97 // DAWN_FLAGS is ';'-separated because individual Dawn options take
98 // comma-separated lists, e.g. 'enable-dawn-features=allow_unsafe_apis,timestamp_quantization'
99 const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
100 Object.defineProperty(globalThis, 'navigator', {
101 value: { gpu: mod.create(dawnFlags) },
102 configurable: true,
103 writable: true,
104 });
105 const { version } = await import(`${specifier}/package.json`, {
106 with: { type: 'json' },
107 }).then(
108 (m) => m.default as { version: string },
109 () => ({ version: '?' }),
110 );
111 return `node-webgpu ${version} (Google Dawn)`;
114// ---------------------------------------------------------------- statistics
115interface Timing {
116 meanMs: number;
117 medianMs: number;
118 p05Ms: number;
119 p95Ms: number;
120 minMs: number;
121 totalMs: number;
122 stepsPerSec: number;
125function timing(samples: Float64Array): Timing {
126 const sorted = Float64Array.from(samples).sort();
127 const q = (p: number): number =>
128 sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
129 let total = 0;
130 for (const v of samples) total += v;
131 const mean = total / samples.length;
132 return {
133 meanMs: mean,
134 medianMs: q(0.5),
135 p05Ms: q(0.05),
136 p95Ms: q(0.95),
137 minMs: sorted[0],
138 totalMs: total,
139 stepsPerSec: 1000 / mean,
140 };
143function fieldRange(v: ArrayLike<number>): { min: number; max: number } {
144 let min = Infinity;
145 let max = -Infinity;
146 for (let i = 0; i < v.length; i++) {
147 if (v[i] < min) min = v[i];
148 if (v[i] > max) max = v[i];
149 }
150 return { min, max };
153// ---------------------------------------------------------------- run
154const model = modelForSpec(spec);
155const { preset } = resolvePreset(spec.preset);
156const cfg = configForSpec(spec);
158let device: GPUDevice | null = null;
159let backend: ShtBackend | null = null;
160let runtime = 'CPU (direct summation, f64)';
161let adapter = '';
163try {
164 if (spec.backend === 'webgpu') {
165 runtime = await installWebGpu();
166 device = await requestShtDevice();
167 adapter = await describeAdapter(device);
168 backend = await GpuBackend.create(device, cfg);
169 } else {
170 backend = new CpuBackend(cfg);
171 }
173 const sim = new Simulation(backend, model, spec.params);
174 await sim.init(spec.seed);
176 if (!wantJson) {
177 const kind =
178 spec.backend === 'webgpu'
179 ? `WebGPU fp32${adapter ? ` — ${adapter}` : ''}`
180 : 'CPU f64';
181 console.log(`turing-sphere bench — solver only, no rendering\n`);
182 console.log(` preset ${preset.label} (model ${model.key}: ${model.species.join(', ')})`);
183 console.log(
184 ` params ${model.params.map((p) => `${p.key}=${spec.params[p.key]}`).join(' ')}`,
185 );
186 console.log(
187 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${backend.nlm.toLocaleString()}`,
188 );
189 console.log(` backend ${kind}\n ${runtime}`);
190 console.log(` run ${spec.warmup} warmup + ${spec.steps} timed steps, seed ${spec.seed}\n`);
191 }
193 for (let s = 0; s < spec.warmup; s++) await sim.step();
195 const samples = new Float64Array(spec.steps);
196 const progress = !wantJson && process.stderr.isTTY;
197 let lastReport = performance.now();
198 let running = 0;
199 for (let s = 0; s < spec.steps; s++) {
200 const t0 = performance.now();
201 await sim.step();
202 samples[s] = performance.now() - t0;
203 running += samples[s];
204 if (progress && performance.now() - lastReport > 1000) {
205 process.stderr.write(
206 `\r\x1b[K ${s + 1}/${spec.steps} steps · ${(running / (s + 1)).toFixed(2)} ms/step`,
207 );
208 lastReport = performance.now();
209 }
210 }
211 if (progress) process.stderr.write('\r\x1b[K');
213 const t = timing(samples);
214 const range = fieldRange(sim.V[0]);
215 let finite = true;
216 for (const v of sim.V[0]) if (!Number.isFinite(v)) finite = false;
218 if (wantJson) {
219 console.log(
220 JSON.stringify(
221 {
222 command: formatCommand(spec),
223 spec,
224 model: model.key,
225 backend: { kind: spec.backend, adapter, runtime },
226 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm: backend.nlm },
227 timing: t,
228 state: {
229 t: sim.t,
230 steps: sim.stepCount,
231 species: model.species[0],
232 min: range.min,
233 max: range.max,
234 contrast: range.max - range.min,
235 finite,
236 },
237 },
238 null,
239 2,
240 ),
241 );
242 } else {
243 console.log(
244 ` ${t.meanMs.toFixed(2)} ms/step ${t.stepsPerSec.toFixed(1)} steps/s ` +
245 `${(spec.params.dt * t.stepsPerSec).toFixed(2)} model time/s`,
246 );
247 console.log(
248 ` median ${t.medianMs.toFixed(2)} · p05 ${t.p05Ms.toFixed(2)} · ` +
249 `p95 ${t.p95Ms.toFixed(2)} · min ${t.minMs.toFixed(2)} ms ` +
250 `(${(t.totalMs / 1000).toFixed(1)} s total)`,
251 );
252 console.log(
253 ` after ${sim.stepCount} steps: t = ${sim.t.toFixed(2)}, ` +
254 `${model.species[0]} ∈ [${range.min.toFixed(4)}, ${range.max.toFixed(4)}] ` +
255 `(contrast ${(range.max - range.min).toFixed(4)})${finite ? '' : ' — NOT FINITE'}`,
256 );
257 console.log(
258 `\n Compare with the ms/step in the app's stats line. That one is also the\n` +
259 ` solver alone, but measured while the page renders the spheres.`,
260 );
261 }
263 backend.destroy();
264 device?.destroy();
265 process.exit(finite ? 0 : 1);
266} catch (e) {
267 backend?.destroy();
268 device?.destroy();
269 fail(errMsg(e));
moveopenescclose