2 * How do the WGSL transforms compare with upstream SHTNS?
3 *
4 * node scripts/compare-native.mjs # transforms, lmax 63
5 * node scripts/compare-native.mjs --mode solver # a whole IMEX timestep
6 * node scripts/compare-native.mjs --check # and diff the final state
7 *
8 * Runs the same spec through every implementation available on this machine and
9 * puts the numbers in one table:
10 *
11 * webgpu src/sht, fp32, through Dawn — what the app runs
12 * shtns cuda SHTNS' own CUDA kernels, fp32 — the like-for-like comparison
13 * shtns cpu SHTNS on the CPU, fp64 — the accuracy and "what a CPU does"
14 * reference (SHTNS has no CPU single precision)
15 *
16 * Everything runs in this one invocation, back to back, so a second process
17 * competing for the GPU affects both sides rather than one. Missing
18 * implementations are reported and skipped, not fatal: on a machine without
19 * CUDA you still get webgpu against the CPU.
20 *
21 * With --check it also re-runs each side with --dump-state and diffs the final
22 * spectral state, which is what makes the timing comparison mean anything: the
23 * two are only comparable if they compute the same thing. The native solver is
24 * a transcription of models/<key>.m rather than the .m itself (C cannot run
25 * numbl), so this is the check on that transcription.
26 *
27 * Needs `bench/shtns/bootstrap.sh && make` in bench/shtns first, and desktop
28 * WebGPU (the optional `webgpu` package) for the WGSL side.
29 *
30 * This compares implementations, all in the terminal. For the browser against
31 * the terminal, see scripts/compare-perf.mjs.
32 */
33import { spawnSync } from 'node:child_process';
34import { readFileSync, existsSync, unlinkSync } from 'node:fs';
35import { join } from 'node:path';
36import { tmpdir } from 'node:os';
38const ROOT = new URL('..', import.meta.url).pathname;
39const NATIVE = join(ROOT, 'bench', 'shtns');
41// ---------------------------------------------------------------- arguments
42const argv = process.argv.slice(2);
43const has = (name) => argv.includes(`--${name}`);
44const flag = (name, dflt) => {
45 const i = argv.indexOf(`--${name}`);
46 if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
47 const eq = argv.find((a) => a.startsWith(`--${name}=`));
48 return eq ? eq.slice(name.length + 3) : dflt;
49};
50if (has('help') || argv.includes('-h')) {
51 console.log(`usage: node scripts/compare-native.mjs [options]
53 --mode transform|solver what to compare (default transform)
54 --lmax <n> spherical harmonic truncation (default 63)
55 --steps <n> timed steps / round trips (default 1000)
56 --warmup <n> untimed steps first (default 50)
57 --preset <key> model preset; in transform mode only its dealiasing
58 degree is used (default schnak-spots)
59 --batch <n> steps per synchronization on both sides (default 16)
60 --layout theta|phi spatial layout for the native runs. theta is SHTNS'
61 native and fastest; phi is what the WGSL side uses
62 (default theta)
63 --threads <n> OpenMP threads for the SHTNS CPU run. 1 by default,
64 which is the reproducible per-core reference; 0 lets
65 the library use every core, which at small lmax is
66 often slower than one thread
67 --check also diff the final state between implementations
68 --check-steps <n> steps for that state, kept short on purpose: fp32
69 round-off accumulates, and in solver mode the
70 unstable modes amplify it (default 20)
71 --tolerance <x> --check threshold on the relative L2 (default 2e-3)
72 --no-cpu skip the SHTNS CPU run
73 --json machine-readable output`);
74 process.exit(0);
75}
76const mode = flag('mode', 'transform');
77if (mode !== 'transform' && mode !== 'solver') {
78 console.error(`compare-native: --mode must be 'transform' or 'solver'`);
79 process.exit(2);
80}
81const lmax = flag('lmax', '63');
82const steps = flag('steps', '1000');
83const warmup = flag('warmup', '50');
84const preset = flag('preset', 'schnak-spots');
85const batch = flag('batch', '16');
86const layout = flag('layout', 'theta');
87const threads = flag('threads', '1');
88const wantCheck = has('check');
89const checkSteps = flag('check-steps', '20');
90const wantJson = has('json');
91const tolerance = Number(flag('tolerance', '2e-3'));
92const wantCpu = !has('no-cpu');
93const progress = !wantJson && process.stderr.isTTY;
95const tmp = (tag) => join(tmpdir(), `turing-sphere-native-${tag}-${process.pid}.json`);
96const cleanup = [];
98// ------------------------------------------------------------------- runners
100 * Pull our JSON object out of stdout.
101 *
102 * A linked library shares the process's stdout, and SHTns writes to it
103 * unconditionally — it announces the GPU it found, which FFT layout it chose and
104 * which VkFFT it linked, none of it gated by shtns_verbose(). So the object may
105 * not start at byte 0. Every producer here prints it with `{` alone on a line and
106 * `}` alone on the last, which is enough to find it.
107 */
108function extractJson(out) {
109 try {
110 return JSON.parse(out);
111 } catch {
112 /* there is something else on stdout; find where the object starts */
113 }
114 const start = out.search(/^\{$/m);
115 const ends = [...out.matchAll(/^\}$/gm)];
116 if (start < 0 || !ends.length) return null;
117 const last = ends[ends.length - 1];
118 try {
119 return JSON.parse(out.slice(start, last.index + 1));
120 } catch {
121 return null;
122 }
123}
125/** Dawn reports two of these on every start-up; they are not the failure. */
126const NOISE = /^Warning: max(Dynamic|Compute|Storage)/;
128/** The tail of a failed run's output, which is where the actual error is. */
129function failureDetail(r, cmd, args) {
130 const lines = `${r.stderr ?? ''}\n${r.stdout ?? ''}`
131 .split('\n')
132 .map((s) => s.trimEnd())
133 .filter((s) => s && !NOISE.test(s));
134 const tail = lines.slice(-14).map((l) => ` ${l}`);
135 return [
136 `exit ${r.status}`,
137 ...tail,
138 ` re-run it alone to see everything:`,
139 ` ${cmd} ${args.join(' ')}`,
140 ].join('\n');
141}
144 * Why a run that reported itself is still not usable. Every producer records the
145 * device, the Fourier stage and whether its result stayed finite, so a wrong
146 * answer can be described rather than dumped.
147 */
148function badResult(json) {
149 const b = json.backend ?? {};
150 const lines = [];
151 const first = json.firstRoundTrip;
152 if (first && first.finite === false) {
153 lines.push('a single spectral -> grid -> spectral round trip returns no finite values,');
154 lines.push('so the transforms are wrong on this device — nothing here is worth timing.');
155 } else if (first && !(first.relL2 < 1e-3)) {
156 lines.push(
157 `a single round trip comes back ${Number(first.relL2).toExponential(2)} away from its`,
158 );
159 lines.push('input, far outside fp32 round-off (~1e-7). The transforms are wrong here.');
160 } else if (json.state && json.state.finite === false) {
161 lines.push('it ran, but the final state has no finite values — one round trip is fine,');
162 lines.push('so something diverges over the length of the run.');
163 } else {
164 lines.push('it reported a failure without saying why; run it alone.');
165 }
166 lines.push(`device: ${b.adapter || '(unknown)'}`);
167 if (json.fourier) lines.push(`Fourier stage: ${String(json.fourier).toUpperCase()}`);
168 return lines.map((l, i) => (i === 0 ? l : ` ${l}`)).join('\n');
169}
b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 171/** Run one side and parse its --json output. `ok: false` with a reason if it is
172 * not available here — a missing binary, no adapter, no CUDA. */
173function run(label, cmd, args, statePath) {
174 const full = statePath
175 ? [...args, '--steps', checkSteps, '--warmup', '0', '--dump-state', statePath]
176 : args;
177 const r = spawnSync(cmd, full, {
178 encoding: 'utf8',
179 cwd: ROOT,
180 maxBuffer: 256 * 1024 * 1024,
181 });
182 if (r.error) return { label, ok: false, why: r.error.message };
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 183 const json = extractJson(r.stdout);
721b8dfbench:sht: check one round trip before timing thousandsJeremy Magland 184 if (r.status !== 0) {
185 // A run that failed but still reported itself is the interesting case: it
186 // computed something wrong rather than failing to start, and it already said
187 // what and on which device. Use that instead of dumping its output.
188 return { label, ok: false, why: json ? badResult(json) : failureDetail(r, cmd, full) };
189 }
191 return {
192 label,
193 ok: false,
194 why: `printed no JSON object:\n${failureDetail(r, cmd, full)}`,
195 };
197 let state = null;
198 if (statePath && existsSync(statePath)) {
199 cleanup.push(statePath);
200 state = JSON.parse(readFileSync(statePath, 'utf8'));
201 }
202 return { label, ok: true, json, state };
203}
205const common = ['--lmax', lmax, '--steps', steps, '--warmup', warmup, '--preset', preset];
206const nativeCommon = [...common, '--batch', batch, '--layout', layout, '--json'];
208const jobs = [];
209if (mode === 'transform') {
210 jobs.push({
211 label: 'webgpu',
212 cmd: 'npx',
213 args: ['vite-node', 'scripts/bench-sht.ts', '--json', ...common, '--batch', batch],
214 });
215} else {
216 jobs.push({
217 label: 'webgpu',
218 cmd: 'npx',
219 args: ['vite-node', 'scripts/bench.ts', '--json', ...common, '--batch', batch],
220 });
221}
222const gpuBin = join(NATIVE, 'shtbench_gpu');
223const cpuBin = join(NATIVE, 'shtbench');
224const nativeMode = ['--mode', mode];
225if (existsSync(gpuBin)) {
226 jobs.push({ label: 'shtns cuda', cmd: gpuBin, args: [...nativeCommon, ...nativeMode] });
227} else {
228 jobs.push({
229 label: 'shtns cuda',
230 missing:
231 `bench/shtns/shtbench_gpu is not built. On a machine with nvcc:\n` +
232 ` cd bench/shtns && ./bootstrap.sh && make`,
233 });
234}
235if (wantCpu) {
236 if (existsSync(cpuBin)) {
237 jobs.push({
238 label: 'shtns cpu',
239 cmd: cpuBin,
240 args: [...nativeCommon, ...nativeMode, '--threads', threads],
241 });
242 } else {
243 jobs.push({
244 label: 'shtns cpu',
245 missing: `bench/shtns/shtbench is not built:\n cd bench/shtns && ./bootstrap.sh && make`,
246 });
247 }
248}
250if (!wantJson) {
251 console.log(
252 `comparing ${mode === 'transform' ? 'transforms' : 'solver timesteps'} — ` +
253 `lmax ${lmax}, ${steps} steps, preset ${preset}` +
254 (mode === 'transform' ? ' (for its grid rule)' : '') +
255 `\n`,
256 );
257}
259const results = [];
260for (const job of jobs) {
261 if (job.missing) {
262 results.push({ label: job.label, ok: false, why: job.missing });
263 continue;
264 }
265 if (progress) process.stderr.write(`\r\x1b[K running ${job.label}...`);
266 results.push(run(job.label, job.cmd, job.args, null));
267 if (progress) process.stderr.write('\r\x1b[K');
268}
270/* The state comparison is a second, short run: the timing wants thousands of
271 * steps and the state comparison wants as few as possible, since fp32 round-off
272 * accumulates and a solver run amplifies it. */
273const states = new Map();
274if (wantCheck) {
275 for (const job of jobs) {
276 if (job.missing || !results.find((r) => r.label === job.label)?.ok) continue;
277 if (progress) process.stderr.write(`\r\x1b[K checking ${job.label}...`);
278 const r = run(job.label, job.cmd, job.args, tmp(job.label.replace(/ /g, '-')));
279 if (r.ok && r.state) states.set(job.label, r.state);
280 if (progress) process.stderr.write('\r\x1b[K');
281 }
282}
284const good = results.filter((r) => r.ok);
285if (!good.length) {
286 console.error('compare-native: nothing ran.\n');
287 for (const r of results) console.error(` ${r.label}: ${r.why}`);
288 process.exit(1);
289}
291// ------------------------------------------------- are these the same problem?
292// The native side keeps its own copy of the presets and the grid rule (C cannot
293// import the TypeScript), so this is the one thing that can silently drift.
294const gridOf = (r) => r.json.grid;
295const ref = good[0];
296const gridProblems = [];
297for (const r of good.slice(1)) {
298 const a = gridOf(ref);
299 const b = gridOf(r);
300 for (const k of ['lmax', 'nlat', 'nphi', 'nlm']) {
301 if (a[k] !== b[k]) gridProblems.push(`${r.label}: ${k} ${b[k]} vs ${ref.label}'s ${a[k]}`);
302 }
303 const pa = ref.json.spec?.params;
304 const pb = r.json.spec?.params;
305 if (pa && pb) {
306 for (const k of Object.keys(pa)) {
307 if (Math.abs(Number(pa[k]) - Number(pb[k])) > 1e-12 * Math.max(1, Math.abs(Number(pa[k])))) {
308 gridProblems.push(`${r.label}: ${k} = ${pb[k]} vs ${ref.label}'s ${pa[k]}`);
309 }
310 }
311 }
312}
313if (gridProblems.length) {
314 console.error(
315 `compare-native: the two sides are not running the same problem, so there is\n` +
316 `nothing to compare. bench/shtns/spec.h has drifted from src/mgpu/registry.ts\n` +
317 `or src/sht/layout.ts:\n`,
318 );
319 for (const p of gridProblems) console.error(` ${p}`);
320 process.exit(1);
321}
323// -------------------------------------------------------------------- report
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 324// Ratios are against the WGSL run, which is the point of the comparison. If that
325// is the side that failed, fall back to whatever did run and say so, rather than
326// printing "1.00x webgpu" for a run webgpu had no part in.
b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 327const rate = (r) => r.json.throughput.msPerStep;
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 328const baseRun = good.find((r) => r.label === 'webgpu') ?? good[0];
329const base = rate(baseRun);
331if (wantJson) {
332 console.log(
333 JSON.stringify(
334 {
335 mode,
336 spec: {
337 lmax: Number(lmax),
338 steps: Number(steps),
339 warmup: Number(warmup),
340 preset,
341 batch: Number(batch),
342 layout,
343 threads: Number(threads),
344 },
345 grid: gridOf(ref),
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 346 baseline: baseRun.label,
b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 347 runs: results.map((r) =>
348 r.ok
349 ? {
350 label: r.label,
351 msPerStep: rate(r),
352 stepsPerSec: r.json.throughput.stepsPerSec,
353 encodeMsPerStep: r.json.throughput.encodeMsPerStep,
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 354 ratioToBaseline: rate(r) / base,
b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 355 precision: r.json.backend.precision,
356 adapter: r.json.backend.adapter,
357 library: r.json.backend.library,
358 digest: r.json.digest ?? null,
359 }
360 : { label: r.label, ok: false, why: r.why },
361 ),
362 },
363 null,
364 2,
365 ),
366 );
367} else {
368 const unit = mode === 'transform' ? 'ms/round trip' : 'ms/step';
369 console.log(
370 ` grid lmax ${gridOf(ref).lmax} · ${gridOf(ref).nlat}×${gridOf(ref).nphi} · ` +
371 `nlm ${gridOf(ref).nlm.toLocaleString()}` +
372 (mode === 'transform' ? ' (one synthesis + one analysis per round trip)' : ''),
373 );
374 console.log();
375 for (const r of results) {
376 if (!r.ok) {
377 console.log(` ${r.label.padEnd(11)} not available — ${r.why}`);
378 continue;
379 }
380 const ms = rate(r);
381 const ratio = ms / base;
382 console.log(
383 ` ${r.label.padEnd(11)} ${ms.toFixed(3)} ${unit} ` +
384 `${r.json.throughput.stepsPerSec.toFixed(0)}/s ` +
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 385 `${r === baseRun ? '(baseline)' : `${ratio.toFixed(2)}x ${baseRun.label}`} ` +
b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 386 `${r.json.backend.precision}`,
387 );
388 console.log(
389 ` ${''.padEnd(11)} ${r.json.backend.adapter}` +
390 (r.json.throughput.encodeMsPerStep
391 ? ` · CPU-side launching ${r.json.throughput.encodeMsPerStep.toFixed(3)} ms/step`
392 : ''),
393 );
e302a6eDo not assume a linked library leaves stdout aloneJeremy Magland 394 if (r === baseRun && baseRun.label !== 'webgpu') {
395 console.log(
396 ` ${''.padEnd(11)} webgpu did not run, so this is the baseline instead — which is` +
397 `\n ${''.padEnd(11)} not the comparison you wanted. Fix that side first.`,
398 );
399 }
402 // Same caution compare-perf.mjs takes: a ratio between two different devices
403 // is not a comparison of implementations.
404 const wg = good.find((r) => r.label === 'webgpu');
405 const cuda = good.find((r) => r.label === 'shtns cuda');
406 const software = (a) => /swiftshader|llvmpipe|software|basic render/i.test(a ?? '');
407 if (wg && cuda) {
408 const a = wg.json.backend.adapter ?? '';
409 const b = cuda.json.backend.adapter ?? '';
410 if (software(a)) {
411 console.log(
412 `\n STOP the WGSL side is on a software renderer (${a}), so the ratio above\n` +
413 ` compares a CPU emulation against a real GPU and means nothing. Dawn reaches\n` +
414 ` the GPU through Vulkan; DAWN_FLAGS='backend=vulkan' makes it explain itself.`,
415 );
416 } else if (!sameDevice(a, b)) {
417 console.log(
418 `\n NOTE the two name different devices. If this machine has more than one GPU,\n` +
419 ` they are not comparable — point Dawn and --device at the same one:\n` +
420 ` webgpu: ${a}\n shtns cuda: ${b}`,
421 );
422 }
423 }
3bc607ecompare-native: say when a row is measuring launches, not the GPUJeremy Magland 424 // A run whose CPU-side share is most of its wall time is not measuring the GPU
425 // at all — it is measuring how long the host takes to queue the work. That is a
426 // real cost, but it puts a floor under the number that has nothing to do with
427 // the transform, and it means a ratio against it understates the gap in GPU
428 // work. Say so rather than letting the headline ratio be read as compute.
429 const LAUNCH_BOUND = 0.5;
430 for (const r of good) {
431 const share = (r.json.throughput.encodeMsPerStep ?? 0) / rate(r);
432 if (share > LAUNCH_BOUND) {
433 console.log(
434 `\n NOTE ${r.label} spends ${(100 * share).toFixed(0)}% of its time on the CPU queueing\n` +
435 ` work, so ${rate(r).toFixed(3)} ms is roughly what it costs to *submit* a round trip\n` +
436 ` here, not what the GPU spends on one — the real GPU time is below that and\n` +
437 ` this measurement cannot see it. Raise --lmax until the GPU dominates, or read\n` +
438 ` any ratio involving this row as a lower bound on the difference in GPU work.`,
439 );
440 }
441 }
443 const ratio = rate(cuda) / rate(wg);
444 console.log(
445 `\n ${
446 ratio < 1
447 ? `SHTNS' CUDA transforms are ${(1 / ratio).toFixed(2)}x faster than the WGSL ones`
448 : `the WGSL transforms are ${ratio.toFixed(2)}x faster than SHTNS' CUDA ones`
449 } on the same device, at the same precision and grid.`,
450 );
451 console.log(
452 ` Things that are genuinely different, and worth checking before reading much\n` +
453 ` into the number: SHTNS runs its Legendre recurrence in fp64 for lmax <= 128\n` +
454 ` even in fp32 mode (SHTNS_GPU_REC_PREC=1 forces fp32, which is what WebGPU is\n` +
455 ` restricted to); it uses cuFFT or VkFFT for the Fourier stage against a WGSL\n` +
456 ` FFT; and --layout theta is its native layout, phi is the WGSL one.`,
457 );
458 }
459}
461// --------------------------------------------------------------------- check
462let checkFailed = false;
463if (wantCheck) {
464 const labels = [...states.keys()].filter((k) => states.get(k).state?.length);
465 if (labels.length < 2) {
466 if (!wantJson) console.log(`\n --check: fewer than two implementations produced a state.`);
467 } else {
468 if (!wantJson)
469 console.log(
470 `\n --check: the spectral state after exactly ${checkSteps} ` +
471 `${mode === 'transform' ? 'round trips' : 'steps'} from seed 1\n`,
472 );
473 const bl = labels[0];
474 const b = states.get(bl);
475 for (const label of labels) {
476 const s = states.get(label);
477 let note = '(reference)';
478 if (label !== bl) {
479 const rel = relL2(s.state, b.state);
480 checkFailed = checkFailed || !(rel < tolerance);
481 note = `relative L2 vs ${bl}: ${rel.toExponential(3)}`;
482 }
483 if (!wantJson) {
484 console.log(` ${label.padEnd(11)} ${digestLine(s.digest)}`);
485 console.log(` ${''.padEnd(11)} ${note}`);
486 }
487 // The seeded input has to match, or the two states are answers to
488 // different questions and the L2 above says nothing about the transforms.
489 if (label !== bl && s.input && b.input && Math.abs(s.input.rms - b.input.rms) > 1e-9) {
490 console.log(
491 ` ${''.padEnd(11)} MISMATCHED INPUT: seeded spectrum rms ${s.input.rms} vs ` +
492 `${b.input.rms}.\n` +
493 ` ${''.padEnd(11)} The two seeded generators disagree (shtb_seeded_spectrum in\n` +
494 ` ${''.padEnd(11)} bench/shtns/spec.h against seededSpectrum in bench-sht.ts), so the\n` +
495 ` ${''.padEnd(11)} difference above is not about the transforms.`,
496 );
497 checkFailed = true;
498 }
499 }
500 if (!wantJson) {
501 console.log(
502 `\n ${checkFailed ? 'FAIL' : 'PASS'} every implementation agrees to better than ` +
503 `${tolerance.toExponential(1)} relative L2`,
504 );
505 console.log(
506 ` fp32 against fp64 lands near 1e-6 for a single transform and drifts\n` +
507 ` upward with the step count; two fp32 implementations differ in\n` +
508 ` fused-multiply-add and the other latitude fp32 allows. Raise\n` +
509 ` --check-steps to watch the drift accumulate.`,
510 );
511 }
512 }
513}
515for (const p of cleanup) {
516 try {
517 unlinkSync(p);
518 } catch {
519 /* best effort */
520 }
521}
522process.exit(checkFailed ? 1 : 0);
524// ------------------------------------------------------------------- helpers
525function relL2(a, b) {
526 let num = 0;
527 let den = 0;
528 const n = Math.min(a.length, b.length);
529 for (let i = 0; i < n; i++) {
530 const d = a[i] - b[i];
531 num += d * d;
532 den += b[i] * b[i];
533 }
534 return Math.sqrt(num / Math.max(den, 1e-300));
535}
537function digestLine(d) {
538 if (!d) return '(no digest)';
539 const g = (v) => Number(v).toPrecision(9);
540 return `min=${g(d.min)} max=${g(d.max)} mean=${g(d.mean)} rms=${g(d.rms)}`;
541}
543/** Two adapter strings for the same GPU rarely match textually — Dawn says
544 * "NVIDIA GeForce RTX 4090" where CUDA says "NVIDIA GeForce RTX 4090 (sm_89,
545 * 128 SMs)". Compare on the words they have in common instead. */
546function sameDevice(a, b) {
547 const words = (s) =>
548 new Set(
549 (s ?? '')
550 .toLowerCase()
551 .replace(/[^a-z0-9 ]+/g, ' ')
552 .split(/\s+/)
553 .filter((w) => w.length > 2),
554 );
555 const wa = words(a);
556 const wb = words(b);
557 let shared = 0;
558 for (const w of wa) if (wb.has(w)) shared++;
559 return shared >= 2;
560}