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 } 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 USAGE = `usage: npm run bench:sht -- [options]
43 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
44 --steps <n> timed round trips (default ${DEFAULT_STEPS})
45 --warmup <n> untimed round trips first (default ${DEFAULT_WARMUP})
46 --seed <n> seed of the initial spectrum (default ${DEFAULT_SEED})
47 --batch <n> round trips per submit for the throughput number (default 16)
48 --preset <key> only for its dealiasing degree, so the grid matches the
49 solver benchmark's: ${presets.map((p) => p.key).join(' | ')}
50 (default ${presets[0].key})
51 --fourier <mode> auto | fft | dft (default auto)
52 --digest after timing, re-run exactly --steps round trips from the
53 seed and print a digest of the final spectrum
54 --dump-state <f> like --digest, and write the spectrum to <f> as JSON, for
55 scripts/compare-native.mjs to diff
56 --json machine-readable output
57 --help
59The native counterpart is
60 bench/shtns/shtbench --mode transform --lmax <n> --steps <n> (CPU, fp64)
61 bench/shtns/shtbench_gpu --mode transform --lmax <n> --steps <n> (CUDA, fp32)`;
63function fail(msg: string, code = 1): never {
64 console.error(`bench:sht: ${msg}`);
65 process.exit(code);
66}
68// ---------------------------------------------------------------- arguments
69const argv = process.argv.slice(2);
70if (argv.includes('--help') || argv.includes('-h')) {
71 console.log(USAGE);
72 process.exit(0);
73}
74const wantJson = argv.includes('--json');
75let batch = 16;
76let fourier: 'auto' | 'fft' | 'dft' = 'auto';
77let dumpState: string | null = null;
78let wantDigest = false;
79const rest: string[] = [];
80for (let i = 0; i < argv.length; i++) {
81 const a = argv[i];
82 if (a === '--json') continue;
83 if (a === '--digest') {
84 wantDigest = true;
85 continue;
86 }
87 const valued = (name: string): string | null => {
88 if (a === `--${name}`) return argv[++i];
89 if (a.startsWith(`--${name}=`)) return a.slice(name.length + 3);
90 return null;
91 };
92 const b = valued('batch');
93 if (b !== null) {
94 batch = Number(b);
95 continue;
96 }
97 const f = valued('fourier');
98 if (f !== null) {
99 if (f !== 'auto' && f !== 'fft' && f !== 'dft') fail(`--fourier must be auto|fft|dft`, 2);
100 fourier = f;
101 continue;
102 }
103 const d = valued('dump-state');
104 if (d !== null) {
105 dumpState = d;
106 wantDigest = true;
107 continue;
108 }
109 rest.push(a);
110}
111if (!Number.isInteger(batch) || batch < 1) fail('--batch must be an integer >= 1', 2);
113let spec: RunSpec;
114try {
115 spec = parseArgs(rest);
116} catch (e) {
117 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
118}
119const cfg = configForSpec(spec);
121// -------------------------------------------------------------- the spectrum
122/**
123 * A seeded starting spectrum, uniform in [-1, 1). Deliberately the plainest
124 * thing both sides can agree on bit for bit: mulberry32 only, no transcendental
125 * functions, so a difference in the result is a difference in the transforms and
126 * not in the input. The m = 0 imaginary parts are zeroed, since a real field has
127 * none and the two libraries need not treat a coefficient that cannot occur
128 * alike. Mirrors shtb_seeded_spectrum() in bench/shtns/spec.h.
129 */
130function seededSpectrum(lmax: number, mmax: number, nlm: number, seed: number): Float32Array {
131 const rand = makeRand(seed);
132 const qlm = new Float32Array(2 * nlm);
133 for (let m = 0; m <= mmax; m++) {
134 for (let l = m; l <= lmax; l++) {
135 const lm = lmIndex(lmax, l, m);
136 qlm[2 * lm] = 2 * rand() - 1;
137 const im = 2 * rand() - 1;
138 qlm[2 * lm + 1] = m === 0 ? 0 : im;
139 }
140 }
141 return qlm;
142}
144// ---------------------------------------------------------------------- run
145let device: GPUDevice | null = null;
146let plan: ShtPlan | null = null;
148try {
149 const runtime = await installWebGpu();
150 device = await requestShtDevice().catch((e: unknown) => {
151 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
152 });
153 const adapter = await describeAdapter(device);
154 plan = await ShtPlan.create(device, cfg, { fourier });
155 const nlm = plan.nlm;
156 const npts = cfg.nlat * cfg.nphi;
158 // Two spectral buffers and one spatial one, so a round trip needs no copy:
159 // round trips alternate direction, A -> spat -> B then B -> spat -> A.
160 const mk = (label: string, size: number) =>
161 device!.createBuffer({
162 label,
163 size,
164 usage:
165 GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
166 });
167 const qlm: [GPUBuffer, GPUBuffer] = [mk('sht-bench-qa', 8 * nlm), mk('sht-bench-qb', 8 * nlm)];
168 const spat = mk('sht-bench-spat', 4 * npts);
169 const readback = device.createBuffer({
170 label: 'sht-bench-readback',
171 size: 8 * nlm,
172 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
173 });
174 // Built once, at plan time — a bind group per round trip would be measuring
175 // bind-group creation.
176 const synth: [ShtBinding, ShtBinding] = [
177 plan.createSynthBinding(qlm[0], spat),
178 plan.createSynthBinding(qlm[1], spat),
179 ];
180 const analys: [ShtBinding, ShtBinding] = [
181 plan.createAnalysBinding(spat, qlm[1]),
182 plan.createAnalysBinding(spat, qlm[0]),
183 ];
185 let cur = 0;
186 /** Record `n` round trips into one submission. Returns nothing; the result is
187 * in qlm[cur] once the queue has drained. */
188 const submit = (n: number): void => {
189 const enc = device!.createCommandEncoder({ label: 'sht-bench' });
190 const pass = enc.beginComputePass({ label: 'sht-bench' });
191 for (let i = 0; i < n; i++) {
192 plan!.encodeSynthInto(pass, synth[cur]);
193 plan!.encodeAnalysInto(pass, analys[cur]);
194 cur ^= 1;
195 }
196 pass.end();
197 device!.queue.submit([enc.finish()]);
198 };
199 const seed = (): void => {
200 cur = 0;
201 const q0 = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
202 device!.queue.writeBuffer(qlm[0], 0, q0 as Float32Array<ArrayBuffer>);
203 };
204 const readSpectrum = async (): Promise<Float32Array> => {
205 const enc = device!.createCommandEncoder({ label: 'sht-bench-read' });
206 enc.copyBufferToBuffer(qlm[cur], 0, readback, 0, 8 * nlm);
207 device!.queue.submit([enc.finish()]);
208 await readback.mapAsync(GPUMapMode.READ);
209 const out = new Float32Array(readback.getMappedRange().slice(0));
210 readback.unmap();
211 return out;
212 };
213 const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
215 if (!wantJson) {
216 console.log('turing-sphere bench:sht — transforms only, no solver, no rendering\n');
217 console.log(
218 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${nlm.toLocaleString()}` +
219 ` (dealiased for ${modelForSpec(spec).key}, pdeg ${modelForSpec(spec).pdeg})`,
220 );
221 console.log(` step 1 synthesis + 1 analysis (one round trip)`);
222 console.log(` fourier ${plan.fourierMode.toUpperCase()} stage`);
223 console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
224 console.log(
225 ` run ${spec.warmup} warmup + ${spec.steps} timed round trips, seed ${spec.seed}\n`,
226 );
227 }
229 seed();
230 submit(spec.warmup);
231 await done();
233 // --- throughput: batches submitted together, awaited once each ---
234 const batches = Math.max(1, Math.ceil(spec.steps / batch));
235 const tp0 = performance.now();
236 let stepsRun = 0;
237 let encodeMs = 0;
238 for (let b = 0; b < batches; b++) {
239 const n = Math.min(batch, spec.steps - stepsRun);
240 const e0 = performance.now();
241 submit(n);
242 encodeMs += performance.now() - e0;
243 await done();
244 stepsRun += n;
245 }
246 const throughputMs = (performance.now() - tp0) / stepsRun;
247 const encodePerStep = encodeMs / stepsRun;
249 // --- latency: one round trip per submit ---
250 const latencySteps = Math.min(spec.steps, 200);
251 const samples = new Float64Array(latencySteps);
252 for (let s = 0; s < latencySteps; s++) {
253 const t0 = performance.now();
254 submit(1);
255 await done();
256 samples[s] = performance.now() - t0;
257 }
258 const sorted = Float64Array.from(samples).sort();
259 const q = (p: number): number => sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
260 let latTotal = 0;
261 for (const v of samples) latTotal += v;
262 const latency = {
263 meanMs: latTotal / samples.length,
264 medianMs: q(0.5),
265 p05Ms: q(0.05),
266 p95Ms: q(0.95),
267 minMs: sorted[0],
268 };
270 // --- a reproducible spectrum to compare across implementations ---
271 let digest = null;
272 let inputDigest = null;
273 let state: Float32Array | null = null;
274 if (wantDigest) {
275 const input = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
276 inputDigest = digestOf(input, plan.fourierMode, adapter);
277 cur = 0;
278 device.queue.writeBuffer(qlm[0], 0, input as Float32Array<ArrayBuffer>);
279 submit(spec.steps);
280 await done();
281 state = await readSpectrum();
282 digest = digestOf(state, plan.fourierMode, adapter);
283 }
285 const current = await readSpectrum();
286 let finite = true;
287 let min = Infinity;
288 let max = -Infinity;
289 for (const v of current) {
290 if (!Number.isFinite(v)) finite = false;
291 if (v < min) min = v;
292 if (v > max) max = v;
293 }
295 if (wantJson) {
296 console.log(
297 JSON.stringify(
298 {
299 mode: 'transform',
300 spec: { preset: spec.preset, lmax: spec.lmax, seed: spec.seed, steps: spec.steps, warmup: spec.warmup },
301 backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },
302 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },
303 fourier: plan.fourierMode,
304 throughput: {
305 batch,
306 msPerStep: throughputMs,
307 stepsPerSec: 1000 / throughputMs,
308 encodeMsPerStep: encodePerStep,
309 },
310 latency,
311 digest,
312 input: inputDigest,
313 state: { min, max, finite },
314 },
315 null,
316 2,
317 ),
318 );
319 } else {
320 console.log(
321 ` ${throughputMs.toFixed(3)} ms/round trip ` +
322 `${(1000 / throughputMs).toFixed(1)} round trips/s (batches of ${batch})`,
323 );
324 console.log(` i.e. ${(throughputMs / 2).toFixed(3)} ms per single transform`);
325 console.log(
326 ` of which CPU command encoding: ${encodePerStep.toFixed(3)} ms/round trip ` +
327 `(${((100 * encodePerStep) / throughputMs).toFixed(0)}% — the rest is the GPU)`,
328 );
329 console.log(
330 ` one round trip per submit: ${latency.meanMs.toFixed(3)} ms mean · ` +
331 `median ${latency.medianMs.toFixed(3)} · p05 ${latency.p05Ms.toFixed(3)} · ` +
332 `p95 ${latency.p95Ms.toFixed(3)} · min ${latency.minMs.toFixed(3)}`,
333 );
334 if (!finite) console.log(' — NOT FINITE');
335 if (digest) {
336 console.log(`\n spectrum after ${spec.steps} round trips from seed ${spec.seed}:`);
337 console.log(` ${formatDigest(digest)}`);
338 }
339 console.log(
340 `\n The native counterpart is bench/shtns/shtbench{,_gpu} --mode transform;\n` +
341 ` scripts/compare-native.mjs runs both and lines the numbers up.`,
342 );
343 }
345 if (dumpState && state && digest) {
346 writeFileSync(
347 dumpState,
348 JSON.stringify({
349 mode: 'transform',
350 spec: { preset: spec.preset, lmax: spec.lmax, seed: spec.seed, steps: spec.steps, warmup: spec.warmup },
351 backend: { library: 'shtns-webgpu (src/sht)', adapter, precision: 'fp32' },
352 digest,
353 input: inputDigest,
354 state: [...state],
355 }),
356 );
357 if (!wantJson) console.log(`\n wrote ${dumpState}`);
358 }
360 for (const b of [qlm[0], qlm[1], spat, readback]) b.destroy();
361 plan.destroy();
362 device.destroy();
363 process.exit(finite ? 0 : 1);
364} catch (e) {
365 plan?.destroy();
366 device?.destroy();
367 fail(errMsg(e));
368}