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
99/** Run one side and parse its --json output. `ok: false` with a reason if it is
100 * not available here — a missing binary, no adapter, no CUDA. */
101function run(label, cmd, args, statePath) {
102 const full = statePath
103 ? [...args, '--steps', checkSteps, '--warmup', '0', '--dump-state', statePath]
104 : args;
105 const r = spawnSync(cmd, full, {
106 encoding: 'utf8',
107 cwd: ROOT,
108 maxBuffer: 256 * 1024 * 1024,
109 });
110 if (r.error) return { label, ok: false, why: r.error.message };
111 if (r.status !== 0) {
112 const detail = (r.stderr || r.stdout || '').trim().split('\n').slice(0, 6).join('\n ');
113 return { label, ok: false, why: detail || `exit ${r.status}` };
114 }
115 let json;
116 try {
117 json = JSON.parse(r.stdout);
118 } catch {
119 return { label, ok: false, why: `did not print JSON:\n ${r.stdout.slice(0, 300)}` };
120 }
121 let state = null;
122 if (statePath && existsSync(statePath)) {
123 cleanup.push(statePath);
124 state = JSON.parse(readFileSync(statePath, 'utf8'));
125 }
126 return { label, ok: true, json, state };
127}
129const common = ['--lmax', lmax, '--steps', steps, '--warmup', warmup, '--preset', preset];
130const nativeCommon = [...common, '--batch', batch, '--layout', layout, '--json'];
132const jobs = [];
133if (mode === 'transform') {
134 jobs.push({
135 label: 'webgpu',
136 cmd: 'npx',
137 args: ['vite-node', 'scripts/bench-sht.ts', '--json', ...common, '--batch', batch],
138 });
139} else {
140 jobs.push({
141 label: 'webgpu',
142 cmd: 'npx',
143 args: ['vite-node', 'scripts/bench.ts', '--json', ...common, '--batch', batch],
144 });
145}
146const gpuBin = join(NATIVE, 'shtbench_gpu');
147const cpuBin = join(NATIVE, 'shtbench');
148const nativeMode = ['--mode', mode];
149if (existsSync(gpuBin)) {
150 jobs.push({ label: 'shtns cuda', cmd: gpuBin, args: [...nativeCommon, ...nativeMode] });
151} else {
152 jobs.push({
153 label: 'shtns cuda',
154 missing:
155 `bench/shtns/shtbench_gpu is not built. On a machine with nvcc:\n` +
156 ` cd bench/shtns && ./bootstrap.sh && make`,
157 });
158}
159if (wantCpu) {
160 if (existsSync(cpuBin)) {
161 jobs.push({
162 label: 'shtns cpu',
163 cmd: cpuBin,
164 args: [...nativeCommon, ...nativeMode, '--threads', threads],
165 });
166 } else {
167 jobs.push({
168 label: 'shtns cpu',
169 missing: `bench/shtns/shtbench is not built:\n cd bench/shtns && ./bootstrap.sh && make`,
170 });
171 }
172}
174if (!wantJson) {
175 console.log(
176 `comparing ${mode === 'transform' ? 'transforms' : 'solver timesteps'} — ` +
177 `lmax ${lmax}, ${steps} steps, preset ${preset}` +
178 (mode === 'transform' ? ' (for its grid rule)' : '') +
179 `\n`,
180 );
181}
183const results = [];
184for (const job of jobs) {
185 if (job.missing) {
186 results.push({ label: job.label, ok: false, why: job.missing });
187 continue;
188 }
189 if (progress) process.stderr.write(`\r\x1b[K running ${job.label}...`);
190 results.push(run(job.label, job.cmd, job.args, null));
191 if (progress) process.stderr.write('\r\x1b[K');
192}
194/* The state comparison is a second, short run: the timing wants thousands of
195 * steps and the state comparison wants as few as possible, since fp32 round-off
196 * accumulates and a solver run amplifies it. */
197const states = new Map();
198if (wantCheck) {
199 for (const job of jobs) {
200 if (job.missing || !results.find((r) => r.label === job.label)?.ok) continue;
201 if (progress) process.stderr.write(`\r\x1b[K checking ${job.label}...`);
202 const r = run(job.label, job.cmd, job.args, tmp(job.label.replace(/ /g, '-')));
203 if (r.ok && r.state) states.set(job.label, r.state);
204 if (progress) process.stderr.write('\r\x1b[K');
205 }
206}
208const good = results.filter((r) => r.ok);
209if (!good.length) {
210 console.error('compare-native: nothing ran.\n');
211 for (const r of results) console.error(` ${r.label}: ${r.why}`);
212 process.exit(1);
213}
215// ------------------------------------------------- are these the same problem?
216// The native side keeps its own copy of the presets and the grid rule (C cannot
217// import the TypeScript), so this is the one thing that can silently drift.
218const gridOf = (r) => r.json.grid;
219const ref = good[0];
220const gridProblems = [];
221for (const r of good.slice(1)) {
222 const a = gridOf(ref);
223 const b = gridOf(r);
224 for (const k of ['lmax', 'nlat', 'nphi', 'nlm']) {
225 if (a[k] !== b[k]) gridProblems.push(`${r.label}: ${k} ${b[k]} vs ${ref.label}'s ${a[k]}`);
226 }
227 const pa = ref.json.spec?.params;
228 const pb = r.json.spec?.params;
229 if (pa && pb) {
230 for (const k of Object.keys(pa)) {
231 if (Math.abs(Number(pa[k]) - Number(pb[k])) > 1e-12 * Math.max(1, Math.abs(Number(pa[k])))) {
232 gridProblems.push(`${r.label}: ${k} = ${pb[k]} vs ${ref.label}'s ${pa[k]}`);
233 }
234 }
235 }
236}
237if (gridProblems.length) {
238 console.error(
239 `compare-native: the two sides are not running the same problem, so there is\n` +
240 `nothing to compare. bench/shtns/spec.h has drifted from src/mgpu/registry.ts\n` +
241 `or src/sht/layout.ts:\n`,
242 );
243 for (const p of gridProblems) console.error(` ${p}`);
244 process.exit(1);
245}
247// -------------------------------------------------------------------- report
248const rate = (r) => r.json.throughput.msPerStep;
249const base = rate(good.find((r) => r.label === 'webgpu') ?? good[0]);
251if (wantJson) {
252 console.log(
253 JSON.stringify(
254 {
255 mode,
256 spec: {
257 lmax: Number(lmax),
258 steps: Number(steps),
259 warmup: Number(warmup),
260 preset,
261 batch: Number(batch),
262 layout,
263 threads: Number(threads),
264 },
265 grid: gridOf(ref),
266 runs: results.map((r) =>
267 r.ok
268 ? {
269 label: r.label,
270 msPerStep: rate(r),
271 stepsPerSec: r.json.throughput.stepsPerSec,
272 encodeMsPerStep: r.json.throughput.encodeMsPerStep,
273 ratioToWebgpu: rate(r) / base,
274 precision: r.json.backend.precision,
275 adapter: r.json.backend.adapter,
276 library: r.json.backend.library,
277 digest: r.json.digest ?? null,
278 }
279 : { label: r.label, ok: false, why: r.why },
280 ),
281 },
282 null,
283 2,
284 ),
285 );
286} else {
287 const unit = mode === 'transform' ? 'ms/round trip' : 'ms/step';
288 console.log(
289 ` grid lmax ${gridOf(ref).lmax} · ${gridOf(ref).nlat}×${gridOf(ref).nphi} · ` +
290 `nlm ${gridOf(ref).nlm.toLocaleString()}` +
291 (mode === 'transform' ? ' (one synthesis + one analysis per round trip)' : ''),
292 );
293 console.log();
294 for (const r of results) {
295 if (!r.ok) {
296 console.log(` ${r.label.padEnd(11)} not available — ${r.why}`);
297 continue;
298 }
299 const ms = rate(r);
300 const ratio = ms / base;
301 console.log(
302 ` ${r.label.padEnd(11)} ${ms.toFixed(3)} ${unit} ` +
303 `${r.json.throughput.stepsPerSec.toFixed(0)}/s ` +
304 `${r.label === 'webgpu' ? '(baseline)' : `${ratio.toFixed(2)}x webgpu`} ` +
305 `${r.json.backend.precision}`,
306 );
307 console.log(
308 ` ${''.padEnd(11)} ${r.json.backend.adapter}` +
309 (r.json.throughput.encodeMsPerStep
310 ? ` · CPU-side launching ${r.json.throughput.encodeMsPerStep.toFixed(3)} ms/step`
311 : ''),
312 );
313 }
315 // Same caution compare-perf.mjs takes: a ratio between two different devices
316 // is not a comparison of implementations.
317 const wg = good.find((r) => r.label === 'webgpu');
318 const cuda = good.find((r) => r.label === 'shtns cuda');
319 const software = (a) => /swiftshader|llvmpipe|software|basic render/i.test(a ?? '');
320 if (wg && cuda) {
321 const a = wg.json.backend.adapter ?? '';
322 const b = cuda.json.backend.adapter ?? '';
323 if (software(a)) {
324 console.log(
325 `\n STOP the WGSL side is on a software renderer (${a}), so the ratio above\n` +
326 ` compares a CPU emulation against a real GPU and means nothing. Dawn reaches\n` +
327 ` the GPU through Vulkan; DAWN_FLAGS='backend=vulkan' makes it explain itself.`,
328 );
329 } else if (!sameDevice(a, b)) {
330 console.log(
331 `\n NOTE the two name different devices. If this machine has more than one GPU,\n` +
332 ` they are not comparable — point Dawn and --device at the same one:\n` +
333 ` webgpu: ${a}\n shtns cuda: ${b}`,
334 );
335 }
336 }
337 if (cuda && wg) {
338 const ratio = rate(cuda) / rate(wg);
339 console.log(
340 `\n ${
341 ratio < 1
342 ? `SHTNS' CUDA transforms are ${(1 / ratio).toFixed(2)}x faster than the WGSL ones`
343 : `the WGSL transforms are ${ratio.toFixed(2)}x faster than SHTNS' CUDA ones`
344 } on the same device, at the same precision and grid.`,
345 );
346 console.log(
347 ` Things that are genuinely different, and worth checking before reading much\n` +
348 ` into the number: SHTNS runs its Legendre recurrence in fp64 for lmax <= 128\n` +
349 ` even in fp32 mode (SHTNS_GPU_REC_PREC=1 forces fp32, which is what WebGPU is\n` +
350 ` restricted to); it uses cuFFT or VkFFT for the Fourier stage against a WGSL\n` +
351 ` FFT; and --layout theta is its native layout, phi is the WGSL one.`,
352 );
353 }
354}
356// --------------------------------------------------------------------- check
357let checkFailed = false;
358if (wantCheck) {
359 const labels = [...states.keys()].filter((k) => states.get(k).state?.length);
360 if (labels.length < 2) {
361 if (!wantJson) console.log(`\n --check: fewer than two implementations produced a state.`);
362 } else {
363 if (!wantJson)
364 console.log(
365 `\n --check: the spectral state after exactly ${checkSteps} ` +
366 `${mode === 'transform' ? 'round trips' : 'steps'} from seed 1\n`,
367 );
368 const bl = labels[0];
369 const b = states.get(bl);
370 for (const label of labels) {
371 const s = states.get(label);
372 let note = '(reference)';
373 if (label !== bl) {
374 const rel = relL2(s.state, b.state);
375 checkFailed = checkFailed || !(rel < tolerance);
376 note = `relative L2 vs ${bl}: ${rel.toExponential(3)}`;
377 }
378 if (!wantJson) {
379 console.log(` ${label.padEnd(11)} ${digestLine(s.digest)}`);
380 console.log(` ${''.padEnd(11)} ${note}`);
381 }
382 // The seeded input has to match, or the two states are answers to
383 // different questions and the L2 above says nothing about the transforms.
384 if (label !== bl && s.input && b.input && Math.abs(s.input.rms - b.input.rms) > 1e-9) {
385 console.log(
386 ` ${''.padEnd(11)} MISMATCHED INPUT: seeded spectrum rms ${s.input.rms} vs ` +
387 `${b.input.rms}.\n` +
388 ` ${''.padEnd(11)} The two seeded generators disagree (shtb_seeded_spectrum in\n` +
389 ` ${''.padEnd(11)} bench/shtns/spec.h against seededSpectrum in bench-sht.ts), so the\n` +
390 ` ${''.padEnd(11)} difference above is not about the transforms.`,
391 );
392 checkFailed = true;
393 }
394 }
395 if (!wantJson) {
396 console.log(
397 `\n ${checkFailed ? 'FAIL' : 'PASS'} every implementation agrees to better than ` +
398 `${tolerance.toExponential(1)} relative L2`,
399 );
400 console.log(
401 ` fp32 against fp64 lands near 1e-6 for a single transform and drifts\n` +
402 ` upward with the step count; two fp32 implementations differ in\n` +
403 ` fused-multiply-add and the other latitude fp32 allows. Raise\n` +
404 ` --check-steps to watch the drift accumulate.`,
405 );
406 }
407 }
408}
410for (const p of cleanup) {
411 try {
412 unlinkSync(p);
413 } catch {
414 /* best effort */
415 }
416}
417process.exit(checkFailed ? 1 : 0);
419// ------------------------------------------------------------------- helpers
420function relL2(a, b) {
421 let num = 0;
422 let den = 0;
423 const n = Math.min(a.length, b.length);
424 for (let i = 0; i < n; i++) {
425 const d = a[i] - b[i];
426 num += d * d;
427 den += b[i] * b[i];
428 }
429 return Math.sqrt(num / Math.max(den, 1e-300));
430}
432function digestLine(d) {
433 if (!d) return '(no digest)';
434 const g = (v) => Number(v).toPrecision(9);
435 return `min=${g(d.min)} max=${g(d.max)} mean=${g(d.mean)} rms=${g(d.rms)}`;
436}
438/** Two adapter strings for the same GPU rarely match textually — Dawn says
439 * "NVIDIA GeForce RTX 4090" where CUDA says "NVIDIA GeForce RTX 4090 (sm_89,
440 * 128 SMs)". Compare on the words they have in common instead. */
441function sameDevice(a, b) {
442 const words = (s) =>
443 new Set(
444 (s ?? '')
445 .toLowerCase()
446 .replace(/[^a-z0-9 ]+/g, ' ')
447 .split(/\s+/)
448 .filter((w) => w.length > 2),
449 );
450 const wa = words(a);
451 const wb = words(b);
452 let shared = 0;
453 for (const w of wa) if (wb.has(w)) shared++;
454 return shared >= 2;
455}