1/**
2 * The transforms alone, on desktop WebGPU — the number to put next to upstream
3 * SHTNS.
4 *
5 * npm run bench:sht -- --lmax 63 --steps 2000
6 *
7 * `npm run bench` measures a whole timestep of a .m model. This measures one
8 * spectral -> grid -> spectral round trip and nothing else, which is what
9 * bench/shtns/shtbench{,_gpu} --mode transform measures on the other side. The
10 * solver does one of these per species per step, and profiling of the reference
11 * implementation puts them at ~96% of its compute, so this is the comparison
12 * that actually decides how fast the solver can be.
13 *
14 * The grid comes from the same rule the app uses, through the same
15 * parseArgs/configForSpec as `npm run bench`, so --preset and --lmax mean here
16 * exactly what they mean there. Nothing about the model is used beyond its
17 * dealiasing degree.
18 *
19 * Like the solver benchmark it reports throughput (a batch of round trips
20 * submitted together, awaited once) and latency (one per submit, for the
21 * distribution).
22 */
23import { ShtPlan, requestShtDevice, describeAdapter, type ShtBinding } from '../src/sht/sht.ts';
24import { lmIndex } from '../src/sht/layout.ts';
25import { makeRand } from '../src/mgpu/noise.ts';
26import { digestOf, formatDigest, relL2 } from '../src/mgpu/digest.ts';
27import {
28 parseArgs,
29 configForSpec,
30 modelForSpec,
31 DEFAULT_LMAX,
32 DEFAULT_SEED,
33 DEFAULT_STEPS,
34 DEFAULT_WARMUP,
35 type RunSpec,
36} from '../src/bench/runSpec.ts';
37import { presets } from '../src/mgpu/registry.ts';
38import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
39import { writeFileSync } from 'node:fs';
41const BENCH_SHT_COMMAND = 'npm run bench:sht --';
43const USAGE = `usage: npm run bench:sht -- [options]
45 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
46 --steps <n> timed round trips (default ${DEFAULT_STEPS})
47 --warmup <n> untimed round trips first (default ${DEFAULT_WARMUP})
48 --seed <n> seed of the initial spectrum (default ${DEFAULT_SEED})
49 --batch <n> round trips per submit for the throughput number (default 16)
50 --preset <key> only for its dealiasing degree, so the grid matches the
51 solver benchmark's: ${presets.map((p) => p.key).join(' | ')}
52 (default ${presets[0].key})
53 --fourier <mode> auto | fft | dft (default auto)
54 --digest after timing, re-run exactly --steps round trips from the
55 seed and print a digest of the final spectrum
56 --dump-state <f> like --digest, and write the spectrum to <f> as JSON, for
57 scripts/compare-native.mjs to diff
58 --json machine-readable output
59 --help
61The native counterpart is
62 bench/shtns/shtbench --mode transform --lmax <n> --steps <n> (CPU, fp64)
63 bench/shtns/shtbench_gpu --mode transform --lmax <n> --steps <n> (CUDA, fp32)`;
65function fail(msg: string, code = 1): never {
66 console.error(`bench:sht: ${msg}`);
67 process.exit(code);
68}
70// ---------------------------------------------------------------- arguments
71const argv = process.argv.slice(2);
72if (argv.includes('--help') || argv.includes('-h')) {
73 console.log(USAGE);
74 process.exit(0);
75}
76const wantJson = argv.includes('--json');
77let batch = 16;
78let fourier: 'auto' | 'fft' | 'dft' = 'auto';
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 f = valued('fourier');
100 if (f !== null) {
101 if (f !== 'auto' && f !== 'fft' && f !== 'dft') fail(`--fourier must be auto|fft|dft`, 2);
102 fourier = f;
103 continue;
104 }
105 const d = valued('dump-state');
106 if (d !== null) {
107 dumpState = d;
108 wantDigest = true;
109 continue;
110 }
111 rest.push(a);
112}
113if (!Number.isInteger(batch) || batch < 1) fail('--batch must be an integer >= 1', 2);
115let spec: RunSpec;
116try {
117 spec = parseArgs(rest);
118} catch (e) {
119 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
120}
121const cfg = configForSpec(spec);
123// -------------------------------------------------------------- the spectrum
124/**
125 * A seeded starting spectrum, uniform in [-1, 1). Deliberately the plainest
126 * thing both sides can agree on bit for bit: mulberry32 only, no transcendental
127 * functions, so a difference in the result is a difference in the transforms and
128 * not in the input. The m = 0 imaginary parts are zeroed, since a real field has
129 * none and the two libraries need not treat a coefficient that cannot occur
130 * alike. Mirrors shtb_seeded_spectrum() in bench/shtns/spec.h.
131 */
132function seededSpectrum(lmax: number, mmax: number, nlm: number, seed: number): Float32Array {
133 const rand = makeRand(seed);
134 const qlm = new Float32Array(2 * nlm);
135 for (let m = 0; m <= mmax; m++) {
136 for (let l = m; l <= lmax; l++) {
137 const lm = lmIndex(lmax, l, m);
138 qlm[2 * lm] = 2 * rand() - 1;
139 const im = 2 * rand() - 1;
140 qlm[2 * lm + 1] = m === 0 ? 0 : im;
141 }
142 }
143 return qlm;
144}
146// ---------------------------------------------------------------------- run
147let device: GPUDevice | null = null;
148let plan: ShtPlan | null = null;
150try {
151 const runtime = await installWebGpu();
152 device = await requestShtDevice().catch((e: unknown) => {
153 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
154 });
155 const adapter = await describeAdapter(device);
156 plan = await ShtPlan.create(device, cfg, { fourier });
157 const nlm = plan.nlm;
158 const npts = cfg.nlat * cfg.nphi;
160 // Two spectral buffers and one spatial one, so a round trip needs no copy:
161 // round trips alternate direction, A -> spat -> B then B -> spat -> A.
162 const mk = (label: string, size: number) =>
163 device!.createBuffer({
164 label,
165 size,
166 usage:
167 GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
168 });
169 const qlm: [GPUBuffer, GPUBuffer] = [mk('sht-bench-qa', 8 * nlm), mk('sht-bench-qb', 8 * nlm)];
170 const spat = mk('sht-bench-spat', 4 * npts);
171 const readback = device.createBuffer({
172 label: 'sht-bench-readback',
173 size: 8 * nlm,
174 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
175 });
176 // Built once, at plan time — a bind group per round trip would be measuring
177 // bind-group creation.
178 const synth: [ShtBinding, ShtBinding] = [
179 plan.createSynthBinding(qlm[0], spat),
180 plan.createSynthBinding(qlm[1], spat),
181 ];
182 const analys: [ShtBinding, ShtBinding] = [
183 plan.createAnalysBinding(spat, qlm[1]),
184 plan.createAnalysBinding(spat, qlm[0]),
185 ];
187 let cur = 0;
188 /** Record `n` round trips into one submission. Returns nothing; the result is
189 * in qlm[cur] once the queue has drained. */
190 const submit = (n: number): void => {
191 const enc = device!.createCommandEncoder({ label: 'sht-bench' });
192 const pass = enc.beginComputePass({ label: 'sht-bench' });
193 for (let i = 0; i < n; i++) {
194 plan!.encodeSynthInto(pass, synth[cur]);
195 plan!.encodeAnalysInto(pass, analys[cur]);
196 cur ^= 1;
197 }
198 pass.end();
199 device!.queue.submit([enc.finish()]);
200 };
201 /** `submit` in chunks, so an arbitrary round-trip count does not build one
202 * command buffer with tens of thousands of dispatches in it. */
203 const submitAll = (n: number, chunk = batch): void => {
204 for (let done = 0; done < n; done += chunk) submit(Math.min(chunk, n - done));
205 };
206 const seed = (): void => {
207 cur = 0;
208 const q0 = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
209 device!.queue.writeBuffer(qlm[0], 0, q0 as Float32Array<ArrayBuffer>);
210 };
211 const readSpectrum = async (): Promise<Float32Array> => {
212 const enc = device!.createCommandEncoder({ label: 'sht-bench-read' });
213 enc.copyBufferToBuffer(qlm[cur], 0, readback, 0, 8 * nlm);
214 device!.queue.submit([enc.finish()]);
215 await readback.mapAsync(GPUMapMode.READ);
216 const out = new Float32Array(readback.getMappedRange().slice(0));
217 readback.unmap();
218 return out;
219 };
220 const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
222 // What every report of this run says about itself, whether it succeeded, failed
223 // early, or is being written to a state file for compare-native.mjs to diff.
224 const identity = {
225 mode: 'transform',
226 spec: {
227 preset: spec.preset,
228 lmax: spec.lmax,
229 seed: spec.seed,
230 steps: spec.steps,
231 warmup: spec.warmup,
232 },
233 backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },
234 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },
235 fourier: plan.fourierMode,
236 };
238 if (!wantJson) {
239 console.log('turing-sphere bench:sht — transforms only, no solver, no rendering\n');
240 console.log(
241 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${nlm.toLocaleString()}` +
242 ` (dealiased for ${modelForSpec(spec).key}, pdeg ${modelForSpec(spec).pdeg})`,
243 );
244 console.log(` step 1 synthesis + 1 analysis (one round trip)`);
245 console.log(` fourier ${plan.fourierMode.toUpperCase()} stage`);
246 console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
247 console.log(
248 ` run ${spec.warmup} warmup + ${spec.steps} timed round trips, seed ${spec.seed}\n`,
249 );
250 }
252 /*
253 * Before timing anything: does one round trip work on this device?
254 *
255 * analys(synth(q)) is the identity for a band-limited q — exact Gauss
256 * quadrature, nphi past the aliasing limit — so a single round trip should
257 * return the input to fp32 round-off, and this is the sharpest check the
258 * transforms are working at all here. Doing it separately means a bad result
259 * says *which* it is: broken from the first transform, or drifted over the
260 * thousands of iterations the timing run does. Otherwise that is a manual
261 * bisection on --steps.
262 */
263 seed();
264 const input = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
265 submit(1);
266 await done();
267 const afterOne = await readSpectrum();
268 const firstFinite = afterOne.every((v) => Number.isFinite(v));
269 const firstRelL2 = firstFinite ? relL2(afterOne, input) : NaN;
270 if (!wantJson) {
271 console.log(
272 ` one round trip: ${
273 firstFinite
274 ? `back to the input to ${firstRelL2.toExponential(2)} relative L2`
275 : 'NOT FINITE'
276 }`,
277 );
278 }
279 if (!firstFinite || !(firstRelL2 < 1e-3)) {
280 // Report it the same way a good run reports itself, so a caller reading
281 // --json learns what went wrong and on which device rather than having to
282 // scrape stderr. Then say it in prose and stop: timing a transform that does
283 // not transform is a waste of minutes.
284 if (wantJson) {
285 console.log(
286 JSON.stringify(
287 {
288 ...identity,
289 firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },
290 throughput: null,
291 latency: null,
292 digest: null,
293 input: null,
294 state: { min: null, max: null, finite: firstFinite },
295 },
296 null,
297 2,
298 ),
299 );
300 }
301 const detail = firstFinite
302 ? `it came back ${firstRelL2.toExponential(3)} away from the input, which is far\n` +
303 ` outside fp32 round-off (~1e-7)`
304 : `it came back with no finite values at all`;
305 fail(
306 `a single spectral -> grid -> spectral round trip does not round-trip on this\n` +
307 ` device: ${detail}.\n\n` +
308 ` That is a correctness problem in the transforms here, not a benchmarking one,\n` +
309 ` so there is nothing worth timing yet. Adapter: ${adapter || '(unknown)'};\n` +
310 ` Fourier stage: ${plan.fourierMode.toUpperCase()}.\n\n` +
311 ` Worth trying, in order:\n` +
312 ` npm run test:node the repo's own transform check against\n` +
313 ` its f64 CPU twin, on this device\n` +
314 ` ${BENCH_SHT_COMMAND} --fourier dft the other Fourier stage; if this works,\n` +
315 ` the WGSL FFT is the problem\n` +
316 ` ${BENCH_SHT_COMMAND} --lmax 15 does it depend on the grid size?`,
317 );
318 }
320 seed();
321 submitAll(spec.warmup);
322 await done();
324 // --- throughput: batches submitted together, awaited once each ---
325 const batches = Math.max(1, Math.ceil(spec.steps / batch));
326 const tp0 = performance.now();
327 let stepsRun = 0;
328 let encodeMs = 0;
329 for (let b = 0; b < batches; b++) {
330 const n = Math.min(batch, spec.steps - stepsRun);
331 const e0 = performance.now();
332 submit(n);
333 encodeMs += performance.now() - e0;
334 await done();
335 stepsRun += n;
336 }
337 const throughputMs = (performance.now() - tp0) / stepsRun;
338 const encodePerStep = encodeMs / stepsRun;
340 // --- latency: one round trip per submit ---
341 const latencySteps = Math.min(spec.steps, 200);
342 const samples = new Float64Array(latencySteps);
343 for (let s = 0; s < latencySteps; s++) {
344 const t0 = performance.now();
345 submit(1);
346 await done();
347 samples[s] = performance.now() - t0;
348 }
349 const sorted = Float64Array.from(samples).sort();
350 const q = (p: number): number => sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
351 let latTotal = 0;
352 for (const v of samples) latTotal += v;
353 const latency = {
354 meanMs: latTotal / samples.length,
355 medianMs: q(0.5),
356 p05Ms: q(0.05),
357 p95Ms: q(0.95),
358 minMs: sorted[0],
359 };
361 // --- a reproducible spectrum to compare across implementations ---
362 let digest = null;
363 let inputDigest = null;
364 let state: Float32Array | null = null;
365 if (wantDigest) {
366 inputDigest = digestOf(input, plan.fourierMode, adapter);
367 seed();
368 submitAll(spec.steps);
369 await done();
370 state = await readSpectrum();
371 digest = digestOf(state, plan.fourierMode, adapter);
372 }
374 const current = await readSpectrum();
375 let finite = true;
376 let min = Infinity;
377 let max = -Infinity;
378 for (const v of current) {
379 if (!Number.isFinite(v)) finite = false;
380 if (v < min) min = v;
381 if (v > max) max = v;
382 }
384 if (wantJson) {
385 console.log(
386 JSON.stringify(
387 {
388 ...identity,
389 firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },
390 throughput: {
391 batch,
392 msPerStep: throughputMs,
393 stepsPerSec: 1000 / throughputMs,
394 encodeMsPerStep: encodePerStep,
395 },
396 latency,
397 digest,
398 input: inputDigest,
399 state: { min, max, finite },
400 },
401 null,
402 2,
403 ),
404 );
405 } else {
406 console.log(
407 ` ${throughputMs.toFixed(3)} ms/round trip ` +
408 `${(1000 / throughputMs).toFixed(1)} round trips/s (batches of ${batch})`,
409 );
410 console.log(` i.e. ${(throughputMs / 2).toFixed(3)} ms per single transform`);
411 console.log(
412 ` of which CPU command encoding: ${encodePerStep.toFixed(3)} ms/round trip ` +
413 `(${((100 * encodePerStep) / throughputMs).toFixed(0)}% — the rest is the GPU)`,
414 );
415 console.log(
416 ` one round trip per submit: ${latency.meanMs.toFixed(3)} ms mean · ` +
417 `median ${latency.medianMs.toFixed(3)} · p05 ${latency.p05Ms.toFixed(3)} · ` +
418 `p95 ${latency.p95Ms.toFixed(3)} · min ${latency.minMs.toFixed(3)}`,
419 );
420 if (!finite) console.log(' — NOT FINITE');
421 if (digest) {
422 console.log(`\n spectrum after ${spec.steps} round trips from seed ${spec.seed}:`);
423 console.log(` ${formatDigest(digest)}`);
424 }
425 console.log(
426 `\n The native counterpart is bench/shtns/shtbench{,_gpu} --mode transform;\n` +
427 ` scripts/compare-native.mjs runs both and lines the numbers up.`,
428 );
429 }
431 if (dumpState && state && digest) {
432 writeFileSync(
433 dumpState,
434 JSON.stringify({
435 ...identity,
436 digest,
437 input: inputDigest,
438 state: [...state],
439 }),
440 );
441 if (!wantJson) console.log(`\n wrote ${dumpState}`);
442 }
444 for (const b of [qlm[0], qlm[1], spat, readback]) b.destroy();
445 plan.destroy();
446 device.destroy();
447 process.exit(finite ? 0 : 1);
448} catch (e) {
449 plan?.destroy();
450 device?.destroy();
451 fail(errMsg(e));
452}