Benchmark the WGSL transforms against upstream SHTNS
The transform tests compare src/sht against its own f64 CPU twin, which
shows it is self-consistent but says nothing about how it compares with
the library it is a translation of. bench/shtns/ builds real SHTNS and
runs the same spec through it.
bootstrap.sh clones SHTNS at a pinned commit and configures it, adding
CUDA when nvcc is present, so the same tree gives the CPU comparison
anywhere and shtns_cuda's own fp32 kernels on a machine with an Nvidia
card. spec.h carries the presets, grid rule and seeding from the
TypeScript side, and transcribes the reaction and IMEX update from
models/*.m — the one second implementation of the loop in this repo,
which exists only to be compared against.
Two measurements: --mode transform is one spectral -> grid -> spectral
round trip and nothing else, the library-against-library number;
--mode solver is a whole IMEX step, comparable to `npm run bench`.
npm run bench:sht is the WebGPU counterpart of the first.
compare-native.mjs runs every implementation present back to back in one
invocation and tabulates them. --check diffs the final spectral state,
which is both what makes the timing mean anything and the guard on that
transcription: fp32 WGSL against fp64 SHTNS agrees to ~1e-6 over 20
steps, for every model. It also refuses to compare two runs whose
resolved grid or parameters disagree, since spec.h duplicating
registry.ts is the one thing that can silently drift.
The comparison lands exactly because SHTNS' spectral conventions are
already ours — LM(l,m) agrees index for index, same orthonormal +
Condon-Shortley normalization — so states diff element by element with no
reindexing. What is not identical is written down in bench/shtns/README.md.
shtbench_gpu.cu has not been compiled or run: no nvcc here. It is
syntax-checked against stub CUDA headers only.
13 changed files+3040−8
.gitignoremodified+8−0View file
@@ -2,3 +2,11 @@ node_modules/
22 dist/
33 *.log
44 *.png
5+
6+# bench/shtns: the upstream SHTns checkout, what configure decided, and the
7+# binaries. ./bootstrap.sh rebuilds all of it.
8+bench/shtns/shtns/
9+bench/shtns/shtns.mk
10+bench/shtns/shtbench
11+bench/shtns/shtbench_gpu
12+bench/shtns/*.o
README.mdmodified+89−3View file
@@ -107,7 +107,9 @@ compute, so this port swaps in:
107107 fp32 spherical harmonic transforms in WGSL compute shaders, modeled on
108108 [SHTNS](https://nschaeff.bitbucket.io/shtns/). Its source is vendored under
109109 [`src/sht/`](src/sht/) (CECILL-2.1), including the f64 CPU reference
110- transform used for testing.
110+ transform used for testing. [`bench/shtns/`](bench/shtns/) builds the real
111+ SHTNS and measures ours against it — see
112+ [Against upstream SHTNS](#against-upstream-shtns).
111113 - **Rendering:** three.js spheres with per-vertex colormaps, adapted from the
112114 `SphereEmbedding` view in
113115 [figpack](https://github.com/flatironinstitute/figpack)'s experimental
@@ -319,6 +321,82 @@ genuinely different algorithms that round differently, so a mismatch there
319321 explains a difference in the values rather than being a symptom of one. The app's
320322 stats line and the benchmark both print the chosen stage for the same reason.
321323
324+## Against upstream SHTNS
325+
326+The transforms are a WGSL translation of
327+[SHTNS](https://nschaeff.bitbucket.io/shtns/), and the tests check them against
328+their own f64 CPU twin — which shows they are self-consistent, not how they
329+compare with the library they are modeled on. SHTNS itself runs on the CPU with
330+hand-tuned SIMD codelets, and on Nvidia GPUs with its own CUDA kernels, including
331+a single-precision mode. That is a direct comparison, and
332+[`bench/shtns/`](bench/shtns/) makes it:
333+
334+```
335+cd bench/shtns && ./bootstrap.sh && make # clone SHTns at a pinned commit, build
336+node scripts/compare-native.mjs --check # then, from the repo root
337+```
338+
339+`bootstrap.sh` adds CUDA support when `nvcc` is on `PATH`, so the same tree gives
340+you the CPU comparison anywhere and the GPU one on a machine with an Nvidia card.
341+`compare-native.mjs` runs every implementation present, back to back in one
342+invocation so a second process competing for the GPU affects both sides rather
343+than one, and prints them in one table — here on a machine with no Nvidia card,
344+so the CUDA row is missing rather than invented:
345+
346+```
347+ grid lmax 63 · 128×256 · nlm 2,080 (one synthesis + one analysis per round trip)
348+
349+ webgpu 0.250 ms/round trip 4000/s (baseline) fp32
350+ Intel open-source Mesa driver: Mesa 25.0.7 (gen-12lp)
351+ shtns cuda not available — bench/shtns/shtbench_gpu is not built
352+ shtns cpu 0.110 ms/round trip 9068/s 0.44x webgpu fp64
353+ CPU, 1 thread
354+```
355+
356+Two things are measured, because they answer different questions:
357+
358+- **transforms** (`npm run bench:sht` here, `--mode transform` there) — one
359+ spectral → grid → spectral round trip and nothing else. This is the
360+ library-against-library number, and since the transforms are ~96% of the
361+ solver's compute it is what decides how fast the solver can be.
362+- **solver** (`npm run bench` here, `--mode solver` there) — a whole IMEX Euler
363+ timestep, which is what the app's `solver` line reports.
364+
365+`--check` diffs the final spectral state across implementations, which is what
366+makes the timing mean anything: two numbers are only comparable if they are the
367+cost of the same computation. That check is possible at all because the spectral
368+layout and normalization are SHTNS's own — orthonormal with Condon–Shortley,
369+coefficients grouped by `m`, `LM(l,m)` agreeing index for index — so a state can
370+be diffed element by element with no reindexing. Over 20 steps, fp32 WGSL against
371+fp64 SHTNS agrees to **~1e-6** relative L2, for every model.
372+
373+It is also the check on the one second implementation this repo has. The native
374+solver cannot run `models/<key>.m` — C has no numbl — so `bench/shtns/spec.h`
375+restates the same arithmetic, one line per line of MATLAB. `--check` is what
376+keeps that transcription honest, and `compare-native.mjs` refuses to compare two
377+runs whose resolved grid or parameters disagree, which is the other way the two
378+sides could drift.
379+
380+[`bench/shtns/README.md`](bench/shtns/README.md) lists what is *not* identical and
381+should be kept in mind when reading the ratio — SHTNS runs its Legendre
382+recurrence in fp64 even in fp32 mode for `lmax <= 128` (WebGPU has no fp64 at
383+all), the Fourier stages are cuFFT/VkFFT/FFTW against a WGSL FFT, and SHTNS'
384+polar optimization is off by default here because we have none.
385+
386+One thing worth knowing before reading much into a single number: small grids
387+flatter the CPU, because a GPU spends most of a small transform on launch latency
388+rather than arithmetic. Run a sweep. On an Intel Xe iGPU against one core of the
389+same laptop, one round trip costs:
390+
391+| lmax | grid | WGSL (fp32) | SHTNS, 1 CPU core (fp64) |
392+|---|---|---|---|
393+| 31 | 64×128 | 0.183 ms | 0.017 ms |
394+| 63 | 128×256 | 0.250 ms | 0.110 ms |
395+| 127 | 256×512 | 0.733 ms | 0.602 ms |
396+
397+10x behind at lmax 31, 1.2x at lmax 127 — the same comparison, on the same two
398+chips. Whatever a single number says, it is saying it about one grid size.
399+
322400 ## Tests
323401
324402 There is no second implementation of the solver to diff against, so the `.m` path
@@ -352,8 +430,10 @@ invisible in the numbers, so it is asserted directly. (It has already caught one
352430 regression.)
353431
354432 [`test/transformChecks.ts`](test/transformChecks.ts) is the one remaining
355-implementation-vs-implementation check, comparing the WGSL transforms against
356-shtns-webgpu's f64 CPU twin.
433+implementation-vs-implementation check inside the suite, comparing the WGSL
434+transforms against shtns-webgpu's f64 CPU twin. Comparing them against *upstream*
435+SHTNS is a separate, opt-in step, because it needs a native toolchain — see
436+[Against upstream SHTNS](#against-upstream-shtns).
357437
358438 All three modules run in **both** environments, so the two GPU stacks get the same
359439 guarantees:
@@ -368,6 +448,8 @@ Other commands:
368448
369449 - `npm run bench -- --help` — the desktop benchmark (see
370450 [Desktop vs browser](#desktop-vs-browser)).
451+- `npm run bench:sht -- --help` — the transforms alone, with no solver around
452+ them, for comparing against upstream SHTNS.
371453 - `npx vite-node scripts/longrun-node.ts [lmax]` — run to t = 100 and confirm the
372454 pattern saturates into O(1)-contrast spots rather than decaying or diverging.
373455 - `node scripts/soak.mjs [steps] [lmax]` — drive the demo for many steps,
@@ -382,6 +464,10 @@ Other commands:
382464 - `node scripts/compare-perf.mjs` — measure the same solver work in both and split
383465 the difference (see
384466 [Why the browser is slower](#why-the-browser-is-slower-and-how-to-find-out-by-how-much)).
467+- `node scripts/compare-native.mjs` — run one spec through the WGSL transforms and
468+ through upstream SHTNS, and line the numbers up (see
469+ [Against upstream SHTNS](#against-upstream-shtns)). Needs
470+ [`bench/shtns/`](bench/shtns/) built first.
385471 - `test.html?soak=<steps>&lmax=<n>` — solver-only soak with no rendering.
386472
387473 ### A note on canvas resizing
bench/shtns/Makefileadded+58−0View file
@@ -0,0 +1,58 @@
1+# Native SHTNS benchmarks, to compare against the WGSL transforms in ../../src/sht.
2+#
3+# ./bootstrap.sh # fetch and build upstream SHTNS, write shtns.mk
4+# make # shtbench, plus shtbench_gpu if SHTNS has CUDA support
5+#
6+# -ffast-math is deliberately NOT used: the fp32 GPU run is compared element by
7+# element against the WebGPU one, and reassociation would make that a comparison
8+# of two different computations. -O3 without it is what SHTNS' own codelets use.
9+
10+-include shtns.mk
11+
12+ifeq ($(SHTNS_LIB),)
13+$(error run ./bootstrap.sh first — it fetches SHTNS and writes shtns.mk)
14+endif
15+
16+CXX ?= g++
17+NVCC ?= nvcc
18+# Compile for the GPUs in this machine, which is both the fastest build and the
19+# right answer for a benchmark that only ever runs here. Needs CUDA 11.5+; on an
20+# older toolkit, or to build elsewhere and run here, say
21+# `make NVCC_ARCH=-arch=sm_80`.
22+NVCC_ARCH ?= -arch=native
23+
24+CXXFLAGS ?= -O3 -march=native -std=c++14 -fopenmp -Wall -Wextra -Wno-unused-parameter
25+NVCCFLAGS ?= -O3 -std=c++14 -lineinfo
26+INCLUDES = -I$(SHTNS_DIR)
27+LDLIBS = $(SHTNS_LIB) $(SHTNS_LDLIBS)
28+
29+TARGETS = shtbench
30+ifeq ($(SHTNS_HAS_GPU),1)
31+TARGETS += shtbench_gpu
32+endif
33+
34+all: $(TARGETS)
35+ @echo
36+ @echo "built: $(TARGETS)"
37+ifneq ($(SHTNS_HAS_GPU),1)
38+ @echo "note: SHTNS was built without GPU support, so there is no shtbench_gpu."
39+ @echo " Re-run ./bootstrap.sh on a machine with nvcc to get it."
40+endif
41+
42+shtbench: shtbench.cpp spec.h
43+ $(CXX) $(CXXFLAGS) $(INCLUDES) $< -o $@ $(LDLIBS)
44+
45+# nvcc drives the link so the CUDA runtime and the C++ standard library (which
46+# SHTNS' .cu objects need) come along without naming them by hand.
47+shtbench_gpu: shtbench_gpu.cu spec.h
48+ $(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) $(INCLUDES) $< -o $@ \
49+ -Xcompiler -fopenmp $(LDLIBS)
50+
51+clean:
52+ rm -f shtbench shtbench_gpu
53+
54+# Also throws away the SHTNS checkout and shtns.mk; ./bootstrap.sh rebuilds them.
55+distclean: clean
56+ rm -rf $(SHTNS_DIR) shtns.mk
57+
58+.PHONY: all clean distclean
bench/shtns/README.mdadded+171−0View file
@@ -0,0 +1,171 @@
1+# The WGSL transforms against upstream SHTNS
2+
3+[`src/sht/`](../../src/sht/) is a WGSL translation of
4+[SHTNS](https://nschaeff.bitbucket.io/shtns/), vendored from
5+[shtns-webgpu](https://github.com/concept-collection/shtns-webgpu). Its tests
6+check it against its own f64 CPU twin, which says it is self-consistent — not
7+how it compares with the library it is modeled on.
8+
9+This directory answers that. It builds upstream SHTNS and runs the same spec
10+through it, so the numbers line up against `npm run bench`:
11+
12+| | transforms | precision | where |
13+|---|---|---|---|
14+| `npm run bench:sht` | WGSL, via Dawn | fp32 | this repo |
15+| `./shtbench_gpu` | SHTNS' own CUDA kernels | fp32 | upstream |
16+| `./shtbench` | SHTNS on the CPU | fp64 | upstream |
17+
18+`shtbench_gpu` is the like-for-like comparison: same GPU, same precision, same
19+grid, same spectral conventions — the only difference is who computes the
20+transform. `shtbench` is fp64 because SHTNS' single precision exists only on the
21+GPU, so it serves as the accuracy reference (how far has fp32 drifted?) and as
22+the "what does a well-optimized CPU do" reference.
23+
24+## Build
25+
26+```
27+./bootstrap.sh # clone SHTNS at a pinned commit, configure, build
28+make # shtbench, plus shtbench_gpu if SHTNS got CUDA support
29+```
30+
31+`bootstrap.sh` adds `--enable-cuda` when `nvcc` is on `PATH` and passes
32+`--enable-openmp` always; `--no-cuda` and `--cuda=ampere` override it. It writes
33+`shtns.mk` with the library name and the link flags `configure` decided on, which
34+the `Makefile` includes — so a host that needed MKL or a different FFTW still
35+links without editing anything here. Both the checkout and `shtns.mk` are
36+gitignored; `make distclean` throws them away.
37+
38+Needs: a C++ compiler, FFTW3 headers (`libfftw3-dev`), and for the GPU half a
39+CUDA toolkit. SHTNS' `configure` wants `CUDA_PATH` set; `bootstrap.sh` derives it
40+from `nvcc`'s location if it is not.
41+
42+## Run
43+
44+```
45+./shtbench_gpu --mode transform --lmax 63 --steps 2000
46+./shtbench_gpu --mode solver --lmax 63 --steps 2000 --preset schnak-spots
47+./shtbench --help
48+```
49+
50+Both binaries take the same options as `npm run bench` — `--preset`, `--lmax`,
51+`--steps`, `--warmup`, `--seed`, `--batch`, any model parameter by name — plus
52+`--mode`, `--layout`, `--polar-eps`, `--json`, `--digest`, `--dump-state`.
53+
54+Two things are measured, and they answer different questions:
55+
56+- **`--mode transform`** is one spectral → grid → spectral round trip and nothing
57+ else. This is the library-against-library number. Profiling of the reference
58+ implementation puts the transforms at ~96% of the solver's compute, so this is
59+ what decides how fast the solver can be.
60+- **`--mode solver`** is one IMEX Euler timestep of `models/<key>.m`:
61+ 2·species transforms, the reaction on the grid, the spectral update. This is
62+ the number `npm run bench` and the app's `solver` line report.
63+
64+Both report throughput (a batch launched together, waited for once — what
65+`--batch` controls, matching `npm run bench -- --batch`) and a per-step
66+distribution from one synchronization per step.
67+
68+## Compare
69+
70+```
71+node scripts/compare-native.mjs # transforms, lmax 63
72+node scripts/compare-native.mjs --mode solver
73+node scripts/compare-native.mjs --check # and diff the final state
74+```
75+
76+from the repo root (or `npm run bench:native --`). It runs every implementation
77+present on the machine, back to back in one invocation so a second process
78+competing for the GPU affects both sides rather than one, and prints them in one
79+table. Missing implementations are reported and skipped, so this is still useful
80+on a machine with no CUDA.
81+
82+`--check` adds a short second pass that diffs the final spectral state across
83+implementations. That is what makes the timing mean anything: two numbers are
84+only comparable if they are the cost of the same computation.
85+
86+## What is and is not the same on the two sides
87+
88+The comparison is exact where it can be:
89+
90+- **Spectral layout and normalization are identical.** SHTNS' default
91+ `sht_orthonormal` with the Condon–Shortley phase, `mres = 1`, coefficients
92+ grouped by `m` — which is what [`src/sht/layout.ts`](../../src/sht/layout.ts)
93+ implements, down to `LM(l,m)` agreeing index for index. So a spectral state can
94+ be diffed element by element with no reindexing.
95+- **The grid is identical.** `shtb_grid_for_lmax` in [`spec.h`](spec.h) is
96+ `gridForLmax` from `src/sht/layout.ts`, including rounding `nphi` up to a power
97+ of two — which SHTNS does not need, but the grids have to match. Both sides
98+ assert they got the grid they asked for.
99+- **The seed is identical.** `shtb_seeded_noise` and `shtb_seeded_spectrum` are
100+ transcriptions of `src/mgpu/noise.ts`. The transform check deliberately uses a
101+ spectrum drawn with integer arithmetic only, so it is bit-identical on both
102+ sides and a difference in the result is a difference in the transforms.
103+
104+And explicit where it cannot be:
105+
106+- **The solver step is transcribed, not shared.** The app compiles
107+ `models/<key>.m` through numbl; C cannot, so `shtb_react` and `shtb_imex` in
108+ [`spec.h`](spec.h) restate the same arithmetic, one line per line of MATLAB
109+ (including writing `u.^3` as `u*u*u`, which is what the WGSL backend emits for
110+ it). This is the repo's one second implementation of the loop, and it exists
111+ only to be compared against — `compare-native.mjs --check` is what keeps it
112+ honest. It agrees with the real `.m` path to ~1e-6 over 20 steps, for every
113+ model.
114+- **The presets are duplicated.** `spec.h` copies the tables from
115+ `src/mgpu/registry.ts`. This is the one thing that could silently drift, so
116+ `compare-native.mjs` compares both sides' resolved grid and parameters and
117+ refuses to compare two runs that disagree.
118+- **SHTNS runs the Legendre recurrence in fp64 even in fp32 mode**, for
119+ `lmax <= 128` on a GPU with usable fp64 (`SHT_L_RESCALE_FLY_FLOAT` in its
120+ `sht_private.h`). WebGPU has no fp64 at all, so ours cannot. Set
121+ `SHTNS_GPU_REC_PREC=1` to force SHTNS' recurrence into fp32 for the closer
122+ comparison; the run prints which it used.
123+- **Different FFTs.** SHTNS uses cuFFT or VkFFT on the GPU and FFTW on the CPU;
124+ ours is a WGSL FFT (or a DFT when the device's workgroup limits do not fit one
125+ — the run says which). These are different algorithms with different cost and
126+ different rounding.
127+- **Polar optimization is off by default here** (`--polar-eps 0`), because the
128+ WGSL transforms do not have it. SHTNS' own default is `1e-10` and is worth a
129+ few percent; `--polar-eps 1e-10` turns it on.
130+- **The spatial layout defaults to theta-contiguous**, SHTNS' native and fastest.
131+ `--layout phi` is what the WGSL side uses. The reaction is pointwise and the
132+ spectral layout is unaffected, so a state comparison is valid either way — this
133+ only moves the cost. On the GPU, `--layout phi` needs SHTNS' VkFFT backend,
134+ which it uses whenever `vkfft/vkFFT.h` is in its tree (it normally is;
135+ `bootstrap.sh` says which Fourier stage you got). Its cuFFT fallback handles
136+ only theta-contiguous. SHTNS' own accuracy check runs at `shtns_set_grid` time
137+ and aborts on a mismatch, so this fails loudly rather than quietly — but run
138+ `compare-native.mjs --check` after changing the layout anyway.
139+
140+## Reading the result
141+
142+Small grids flatter the CPU: at `lmax 31` there are 528 coefficients, and a GPU
143+spends most of a transform on launch latency rather than arithmetic. The gap
144+closes with `lmax`, so run a sweep before concluding anything:
145+
146+```
147+for l in 31 63 127 255; do node scripts/compare-native.mjs --lmax $l --steps 500; done
148+```
149+
150+If the WGSL side lands on a software adapter, `compare-native.mjs` says so and
151+stops you — the ratio then compares a CPU emulation against a real GPU and means
152+nothing.
153+
154+## Editing `shtbench_gpu.cu` without a GPU
155+
156+`nvcc` is needed to build it, but not to typecheck it. Rewriting the launch
157+syntax as calls, against a handful of stub declarations, gets g++ to check
158+everything else:
159+
160+```
161+sed -e 's/<<</\/*/g; s/>>>/*\//g' shtbench_gpu.cu > /tmp/check.cpp
162+g++ -fsyntax-only -std=c++14 -I. -Ishtns -I/path/to/cuda-stubs /tmp/check.cpp
163+```
164+
165+where the stub directory holds a `cuda_runtime.h` defining `__global__`,
166+`dim3`, `blockIdx`/`threadIdx`/`blockDim`, `cudaStream_t`, and the dozen
167+`cuda*` functions used here.
168+
169+## License
170+
171+CECILL-2.1, as the rest of this repo — the same license SHTNS itself is under.
bench/shtns/bootstrap.shadded+127−0View file
@@ -0,0 +1,127 @@
1+#!/usr/bin/env bash
2+#
3+# Fetch and build upstream SHTNS, then record how to link against it.
4+#
5+# ./bootstrap.sh # CPU, with OpenMP; adds CUDA if nvcc is on PATH
6+# ./bootstrap.sh --no-cuda # CPU only, even if nvcc is there
7+# ./bootstrap.sh --rev master # a different revision than the pinned one
8+#
9+# Writes shtns.mk, which the Makefile includes: the library name, the link
10+# flags configure decided on, and whether GPU transforms were compiled in.
11+# Re-run it after changing --rev or the CUDA choice; it is idempotent otherwise.
12+set -euo pipefail
13+
14+cd "$(dirname "$0")"
15+
16+# SHTNS 3.7.5, master as of 2026-05-07. Pinned so a change upstream shows up as
17+# a deliberate bump here rather than as a benchmark that quietly moved.
18+REV=4e69ceb3fc5f19475d55b03ff0fdec210dc3075e
19+URL=https://bitbucket.org/nschaeff/shtns.git
20+DIR=shtns
21+WANT_CUDA=auto
22+CUDA_ARCH=
23+
24+while [ $# -gt 0 ]; do
25+ case "$1" in
26+ --no-cuda) WANT_CUDA=no ;;
27+ --cuda) WANT_CUDA=yes ;;
28+ --cuda=*) WANT_CUDA=yes; CUDA_ARCH="${1#*=}" ;;
29+ --rev) shift; REV="$1" ;;
30+ --rev=*) REV="${1#*=}" ;;
31+ -h|--help)
32+ sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//'
33+ exit 0 ;;
34+ *) echo "bootstrap: unknown option '$1'" >&2; exit 2 ;;
35+ esac
36+ shift
37+done
38+
39+if [ "$WANT_CUDA" = auto ]; then
40+ if command -v nvcc >/dev/null 2>&1; then WANT_CUDA=yes; else WANT_CUDA=no; fi
41+fi
42+
43+# ---------------------------------------------------------------------- clone
44+if [ ! -d "$DIR/.git" ]; then
45+ echo "==> cloning $URL"
46+ git clone "$URL" "$DIR"
47+fi
48+echo "==> checking out $REV"
49+# --unshallow so an existing shallow clone can still reach a --rev other than the
50+# one it was cloned at; it errors on a full clone, hence the fallback.
51+git -C "$DIR" fetch --tags --unshallow origin 2>/dev/null || git -C "$DIR" fetch --tags origin
52+git -C "$DIR" checkout --detach "$REV"
53+
54+# ------------------------------------------------------------------ configure
55+CONF=(--enable-openmp)
56+if [ "$WANT_CUDA" = yes ]; then
57+ if [ -n "$CUDA_ARCH" ]; then
58+ CONF+=("--enable-cuda=$CUDA_ARCH")
59+ else
60+ CONF+=(--enable-cuda)
61+ fi
62+ # configure looks for these where CUDA_PATH points; say so if it is unset and
63+ # nvcc is somewhere obvious, since the error it produces otherwise is opaque.
64+ if [ -z "${CUDA_PATH:-}" ]; then
65+ nvcc_dir=$(dirname "$(command -v nvcc)")
66+ export CUDA_PATH="${nvcc_dir%/bin}"
67+ echo "==> CUDA_PATH was unset; using $CUDA_PATH"
68+ fi
69+fi
70+
71+echo "==> configure ${CONF[*]}"
72+( cd "$DIR" && ./configure "${CONF[@]}" )
73+
74+echo "==> make"
75+( cd "$DIR" && make -j"$(getconf _NPROCESSORS_ONLN)" )
76+
77+# --------------------------------------------------------------- record how to
78+# link. configure has already worked all of this out; read it back rather than
79+# guessing, so a machine that needed MKL or a different FFTW still links.
80+libname=$(sed -n 's/^libname *= *//p' "$DIR/Makefile" | head -1)
81+libs=$(sed -n 's/^LIBS=//p' "$DIR/Makefile" | head -1)
82+if [ ! -f "$DIR/$libname" ]; then
83+ echo "bootstrap: expected $DIR/$libname but it is not there" >&2
84+ exit 1
85+fi
86+
87+# Whether the GPU entry points are actually in the archive — asked of the
88+# library rather than of configure's intentions, since --enable-cuda can succeed
89+# at configure time and still leave them out.
90+gpu=0
91+if nm -g "$DIR/$libname" 2>/dev/null | grep -q ' T cu_spat_to_SH_float'; then
92+ gpu=1
93+elif sed -n 's/^objs *= *//p' "$DIR/Makefile" | grep -q 'sht_gpu\.o'; then
94+ gpu=1 # no nm on this host; trust what configure put in the object list
95+fi
96+
97+fft=none
98+[ "$gpu" = 1 ] && fft=vkfft
99+if [ "$gpu" = 1 ] && ! grep -q '^#define VKFFT_BACKEND' "$DIR/sht_config.h"; then
100+ # SHTNS falls back to cuFFT when vkfft/vkFFT.h is absent, and its configure
101+ # does not add -lcufft (the check is commented out), so we have to.
102+ fft=cufft
103+ libs="$libs -lcufft"
104+fi
105+
106+cat > shtns.mk <<EOF
107+# Generated by bootstrap.sh — do not edit; re-run it instead.
108+SHTNS_DIR = $DIR
109+SHTNS_LIB = \$(SHTNS_DIR)/$libname
110+SHTNS_LDLIBS = $libs
111+SHTNS_HAS_GPU = $gpu
112+SHTNS_GPU_FFT = $fft
113+EOF
114+
115+echo
116+echo "==> built $DIR/$libname (GPU transforms: $([ "$gpu" = 1 ] && echo yes || echo no))"
117+if [ "$gpu" = 1 ]; then
118+ echo " GPU Fourier stage: $fft"
119+ if [ "$fft" = cufft ]; then
120+ echo " NOTE: SHTNS' cuFFT path handles only the theta-contiguous layout, so"
121+ echo " --layout phi will not work. vkfft/vkFFT.h is normally in the SHTNS"
122+ echo " tree; if it is missing, that is why."
123+ fi
124+fi
125+echo " link flags: $libs"
126+echo
127+echo "Now: make"
bench/shtns/shtbench.cppadded+413−0View file
@@ -0,0 +1,413 @@
1+/*
2+ * The same run as `npm run bench`, on upstream SHTNS on the CPU (fp64).
3+ *
4+ * ./shtbench --preset schnak-spots --lmax 63 --steps 2000
5+ * ./shtbench --mode transform --lmax 63 --steps 2000
6+ *
7+ * Two things are measured, selected by --mode:
8+ *
9+ * - solver: the IMEX Euler timestep of models/<key>.m — 4 transforms, a
10+ * reaction on the grid, and the spectral update — which is what
11+ * the app's `solver` number and `npm run bench` measure.
12+ * - transform: one spectral -> grid -> spectral round trip and nothing else,
13+ * which is the library-against-library number. The solver does
14+ * one of these per species per step.
15+ *
16+ * This is fp64 throughout, because SHTNS' single precision exists only on the
17+ * GPU. So it is not a like-for-like comparison against the fp32 WGSL
18+ * transforms — it is the accuracy reference (how far has fp32 drifted?) and the
19+ * "what does a well-optimized CPU do" reference. shtbench_gpu is the
20+ * like-for-like one.
21+ *
22+ * The reaction is a transcription of the .m rather than the .m itself; see
23+ * spec.h. It is checked, not trusted: `compare-native.mjs --check` diffs the
24+ * final state against a WebGPU run of the actual .m.
25+ */
26+#include "spec.h"
27+
28+#include <shtns.h>
29+
30+static const char *USAGE =
31+ "usage: ./shtbench [options]\n"
32+ "\n"
33+ " --mode solver|transform what to measure (default solver)\n"
34+ " --preset <key> schnak-spots | schnak-coarse | schnak-fine | brussel |\n"
35+ " allencahn (default schnak-spots)\n"
36+ " --lmax <n> spherical harmonic truncation (default 63)\n"
37+ " --steps <n> timed steps, or round trips (default 2000)\n"
38+ " --warmup <n> untimed steps first (default 100)\n"
39+ " --seed <n> seed of the initial noise / spectrum (default 1)\n"
40+ " --threads <n> OpenMP threads (default: the library's choice)\n"
41+ " --batch <n> accepted for symmetry with shtbench_gpu and ignored:\n"
42+ " every CPU transform call is synchronous, so there is\n"
43+ " nothing to batch\n"
44+ " --layout theta|phi spatial layout: theta-contiguous is SHTNS' native and\n"
45+ " fastest, phi-contiguous is what the WGSL side uses\n"
46+ " (default theta)\n"
47+ " --polar-eps <x> SHTNS polar-optimization threshold. 0 disables it, which\n"
48+ " is what the WGSL transforms do; SHTNS' own default is\n"
49+ " 1e-10 (default 0)\n"
50+ " --digest after timing, re-run exactly --steps steps from the seed\n"
51+ " and print a digest of the final spectral state\n"
52+ " --dump-state <f> like --digest, and write the state to <f> as JSON, for\n"
53+ " scripts/compare-native.mjs to diff\n"
54+ " --<param> <v> any parameter of the preset's model, e.g. --dt 0.05\n"
55+ " --json machine-readable output\n"
56+ " --help\n"
57+ "\n"
58+ "Run the same spec through the WGSL transforms with:\n"
59+ " npm run bench -- --preset <key> --lmax <n> --steps <n> (solver)\n"
60+ " npm run bench:sht -- --lmax <n> --steps <n> (transform)";
61+
62+/*
63+ * The layout arithmetic. SHTNS stores a spatial field either theta-contiguous
64+ * (its native layout) or phi-contiguous (what the WGSL side uses), possibly
65+ * with padding between lines. The reaction is pointwise so it does not care,
66+ * but the seeded perturbation is indexed by (ilat, iphi) and does.
67+ *
68+ * Either way the field is `nlines` contiguous runs of `linelen` doubles,
69+ * `stride` apart.
70+ */
71+struct Grid {
72+ long nlat, nphi, nlat_padded;
73+ int layout;
74+ long nlines, linelen, stride;
75+};
76+
77+static Grid grid_of(shtns_cfg sht, int layout) {
78+ Grid g;
79+ g.nlat = sht->nlat;
80+ g.nphi = sht->nphi;
81+ g.nlat_padded = sht->nlat_padded;
82+ g.layout = layout;
83+ if (layout == SHTB_LAYOUT_PHI) {
84+ g.nlines = g.nlat;
85+ g.linelen = g.nphi;
86+ g.stride = g.nphi;
87+ } else {
88+ g.nlines = g.nphi;
89+ g.linelen = g.nlat;
90+ g.stride = g.nlat_padded;
91+ }
92+ return g;
93+}
94+
95+static inline long spat_index(const Grid &g, long ilat, long iphi) {
96+ return g.layout == SHTB_LAYOUT_PHI ? ilat * g.nphi + iphi : iphi * g.nlat_padded + ilat;
97+}
98+
99+static void fill_uniform(const Grid &g, double *f, double value) {
100+ for (long l = 0; l < g.nlines; l++)
101+ for (long i = 0; i < g.linelen; i++) f[l * g.stride + i] = value;
102+}
103+
104+/* value + the seeded perturbation, which arrives in [ilat*nphi + iphi] order */
105+static void fill_perturbed(const Grid &g, double *f, double value, const float *noise) {
106+ for (long ilat = 0; ilat < g.nlat; ilat++)
107+ for (long iphi = 0; iphi < g.nphi; iphi++)
108+ f[spat_index(g, ilat, iphi)] = value + (double)noise[ilat * g.nphi + iphi];
109+}
110+
111+static void react_field(const Grid &g, const shtb_step_const<double> &c, const double *u,
112+ const double *v, double *r1, double *r2) {
113+#ifdef _OPENMP
114+#pragma omp parallel for schedule(static)
115+#endif
116+ for (long l = 0; l < g.nlines; l++) {
117+ const long o = l * g.stride;
118+ for (long i = 0; i < g.linelen; i++) {
119+ double a = 0, b = 0;
120+ shtb_react<double>(c, u[o + i], v ? v[o + i] : 0.0, &a, &b);
121+ r1[o + i] = a;
122+ if (r2) r2[o + i] = b;
123+ }
124+ }
125+}
126+
127+static void imex_update(const shtb_step_const<double> &c, int nspecies, long nlm, double **U,
128+ const double *const *R, const float *lam) {
129+ const long n2 = 2 * nlm;
130+ for (int k = 0; k < nspecies; k++) {
131+#ifdef _OPENMP
132+#pragma omp parallel for schedule(static)
133+#endif
134+ for (long i = 0; i < n2; i++)
135+ U[k][i] = shtb_imex<double>(c, k, U[k][i], R[k][i], (double)lam[i]);
136+ }
137+}
138+
139+static void field_range(const double *f, long nlines, long linelen, long stride, double *mn,
140+ double *mx, int *finite) {
141+ *mn = INFINITY;
142+ *mx = -INFINITY;
143+ *finite = 1;
144+ for (long l = 0; l < nlines; l++)
145+ for (long i = 0; i < linelen; i++) {
146+ double x = f[l * stride + i];
147+ if (x < *mn) *mn = x;
148+ if (x > *mx) *mx = x;
149+ if (!isfinite(x)) *finite = 0;
150+ }
151+}
152+
153+int main(int argc, char **argv) {
154+ shtb_spec spec;
155+ double polar_eps = 0.0;
156+
157+ /* --polar-eps is ours, not part of the shared spec; take it out first. */
158+ int argc2 = 0;
159+ char **argv2 = (char **)malloc(sizeof(char *) * (size_t)argc);
160+ for (int i = 0; i < argc; i++) {
161+ if (strcmp(argv[i], "--polar-eps") == 0 && i + 1 < argc) {
162+ polar_eps = atof(argv[++i]);
163+ continue;
164+ }
165+ if (strncmp(argv[i], "--polar-eps=", 12) == 0) {
166+ polar_eps = atof(argv[i] + 12);
167+ continue;
168+ }
169+ argv2[argc2++] = argv[i];
170+ }
171+ int rc = shtb_parse_spec(argc2, argv2, &spec, USAGE);
172+ free(argv2);
173+ if (rc) return rc == 1 ? 0 : rc;
174+ /* Every CPU transform call is synchronous, so there is nothing to batch. */
175+ spec.batch = 1;
176+
177+ const int quiet = spec.json;
178+ const int transform_mode = spec.mode == SHTB_MODE_TRANSFORM;
179+ shtns_verbose(0);
180+ const int threads = shtns_use_threads(spec.threads);
181+
182+ shtns_cfg sht = shtns_create(spec.lmax, spec.lmax, 1, sht_orthonormal);
183+ if (!sht) {
184+ fprintf(stderr, "shtbench: shtns_create failed\n");
185+ return 1;
186+ }
187+ const int layout_flag =
188+ spec.layout == SHTB_LAYOUT_PHI ? SHT_PHI_CONTIGUOUS : SHT_THETA_CONTIGUOUS;
189+ if (shtns_set_grid(sht, (enum shtns_type)(sht_gauss | layout_flag | SHT_SCALAR_ONLY), polar_eps,
190+ spec.nlat, spec.nphi) <= 0) {
191+ fprintf(stderr, "shtbench: shtns_set_grid failed for lmax %d on a %dx%d grid\n", spec.lmax,
192+ spec.nlat, spec.nphi);
193+ return 1;
194+ }
195+ if ((int)sht->nlat != spec.nlat || (int)sht->nphi != spec.nphi) {
196+ fprintf(stderr, "shtbench: SHTNS chose a %ux%u grid, not the %dx%d asked for\n", sht->nlat,
197+ sht->nphi, spec.nlat, spec.nphi);
198+ return 1;
199+ }
200+
201+ const Grid g = grid_of(sht, spec.layout);
202+ const long nlm = (long)sht->nlm;
203+ const long nspat = (long)sht->nspat;
204+ const int nsp = spec.model->nspecies;
205+
206+ /* Laplace-Beltrami eigenvalues, 2 x nlm with the value duplicated across the
207+ * real and imaginary halves — the layout eigenvalues() builds in
208+ * src/mgpu/model.ts. Held in float so both sides divide by the same number. */
209+ float *lam = (float *)malloc(sizeof(float) * (size_t)(2 * nlm));
210+ for (long lm = 0; lm < nlm; lm++) {
211+ const int l = sht->li[lm];
212+ lam[2 * lm] = lam[2 * lm + 1] = (float)(l * (l + 1));
213+ }
214+
215+ double *spat[2] = {NULL, NULL};
216+ double *rspat[2] = {NULL, NULL};
217+ cplx *Q[2] = {NULL, NULL};
218+ cplx *R[2] = {NULL, NULL};
219+ for (int k = 0; k < nsp; k++) {
220+ spat[k] = (double *)shtns_malloc(sizeof(double) * (size_t)nspat);
221+ rspat[k] = (double *)shtns_malloc(sizeof(double) * (size_t)nspat);
222+ Q[k] = (cplx *)shtns_malloc(sizeof(cplx) * (size_t)nlm);
223+ R[k] = (cplx *)shtns_malloc(sizeof(cplx) * (size_t)nlm);
224+ memset(spat[k], 0, sizeof(double) * (size_t)nspat);
225+ memset(rspat[k], 0, sizeof(double) * (size_t)nspat);
226+ /* through the double view: a cplx array is [re, im] pairs, and memset on
227+ * std::complex itself is a non-trivial-type warning */
228+ memset((double *)Q[k], 0, sizeof(double) * (size_t)(2 * nlm));
229+ memset((double *)R[k], 0, sizeof(double) * (size_t)(2 * nlm));
230+ }
231+ float *noise = (float *)malloc(sizeof(float) * (size_t)(g.nlat * g.nphi));
232+ float *state32 = (float *)malloc(sizeof(float) * (size_t)(2 * nlm));
233+
234+ const shtb_step_const<double> c = shtb_make_step_const<double>(spec.model, &spec.params);
235+ double base[2];
236+ shtb_background(spec.model, &spec.params, base);
237+
238+ /* --- solver: init and one timestep, from models/<key>.m ------------------ */
239+ auto seed_state = [&]() {
240+ shtb_seeded_noise(g.nlat * g.nphi, spec.model->seed_amp, (uint32_t)spec.seed, noise);
241+ fill_perturbed(g, spat[0], base[0], noise);
242+ spat_to_SH(sht, spat[0], Q[0]);
243+ if (nsp > 1) {
244+ fill_uniform(g, spat[1], base[1]);
245+ spat_to_SH(sht, spat[1], Q[1]);
246+ }
247+ };
248+ auto step = [&]() {
249+ for (int k = 0; k < nsp; k++) SH_to_spat(sht, Q[k], spat[k]);
250+ react_field(g, c, spat[0], nsp > 1 ? spat[1] : NULL, rspat[0], nsp > 1 ? rspat[1] : NULL);
251+ for (int k = 0; k < nsp; k++) spat_to_SH(sht, rspat[k], R[k]);
252+ double *Ud[2] = {(double *)Q[0], (double *)Q[1]};
253+ const double *Rd[2] = {(const double *)R[0], (const double *)R[1]};
254+ imex_update(c, nsp, nlm, Ud, Rd, lam);
255+ };
256+
257+ /* --- transform: one synth + one analys, ping-ponging the two buffers ----- */
258+ cplx *tq[2] = {Q[0], R[0]};
259+ int tcur = 0;
260+ auto seed_spectrum = [&]() {
261+ shtb_seeded_spectrum(spec.lmax, spec.lmax, (uint32_t)spec.seed, state32);
262+ tcur = 0;
263+ double *q = (double *)tq[0];
264+ for (long i = 0; i < 2 * nlm; i++) q[i] = (double)state32[i];
265+ };
266+ auto round_trip = [&]() {
267+ SH_to_spat(sht, tq[tcur], spat[0]);
268+ spat_to_SH(sht, spat[0], tq[tcur ^ 1]);
269+ tcur ^= 1;
270+ };
271+
272+ if (transform_mode)
273+ seed_spectrum();
274+ else
275+ seed_state();
276+
277+ if (!quiet) {
278+ printf("shtbench — upstream SHTNS on the CPU, %s only\n\n",
279+ transform_mode ? "transforms" : "solver");
280+ printf(" mode %s\n", transform_mode
281+ ? "transform (one synth + one analys per step)"
282+ : "solver (one IMEX Euler timestep per step)");
283+ if (!transform_mode) {
284+ printf(" preset %s (models/%s.m: %d species)\n", spec.preset->label, spec.model->key,
285+ nsp);
286+ printf(" params ");
287+ for (int i = 0; i < spec.model->nparams; i++)
288+ printf("%s=%g ", spec.model->params[i].key,
289+ *shtb_field_c(&spec.params, spec.model->params[i].off));
290+ printf("\n");
291+ }
292+ printf(" grid lmax %d · %ldx%ld · nlm %ld\n", spec.lmax, g.nlat, g.nphi, nlm);
293+ printf(" layout %s%s\n",
294+ spec.layout == SHTB_LAYOUT_PHI ? "phi-contiguous" : "theta-contiguous (native)",
295+ (long)sht->nlat_padded != g.nlat ? ", padded" : "");
296+ printf(" backend %s\n fp64, %d thread%s, polar opt %g\n",
297+ shtns_get_build_info(), threads, threads == 1 ? "" : "s", polar_eps);
298+ printf(" run %d warmup + %d timed steps, seed %d\n\n", spec.warmup, spec.steps,
299+ spec.seed);
300+ }
301+
302+ for (int i = 0; i < spec.warmup; i++) {
303+ if (transform_mode)
304+ round_trip();
305+ else
306+ step();
307+ }
308+
309+ double *samples = (double *)malloc(sizeof(double) * (size_t)spec.steps);
310+ const double t0 = shtb_now_ms();
311+ for (int i = 0; i < spec.steps; i++) {
312+ const double a = shtb_now_ms();
313+ if (transform_mode)
314+ round_trip();
315+ else
316+ step();
317+ samples[i] = shtb_now_ms() - a;
318+ }
319+ const double total = shtb_now_ms() - t0;
320+
321+ shtb_report rep;
322+ memset(&rep, 0, sizeof(rep));
323+ char libbuf[192], adapterbuf[64];
324+ rep.library = shtb_json_safe(libbuf, sizeof(libbuf), shtns_get_build_info());
325+ rep.runtime = "cpu";
326+ snprintf(adapterbuf, sizeof(adapterbuf), "CPU, %d thread%s", threads, threads == 1 ? "" : "s");
327+ rep.adapter = adapterbuf;
328+ rep.precision = "fp64";
329+ rep.fourier = "fftw";
330+ rep.nlm = nlm;
331+ rep.ops_per_step = transform_mode ? 2 : 2 * nsp + 2;
332+ rep.ms_per_step = total / spec.steps;
333+ rep.encode_ms_per_step = 0; /* nothing is deferred: every call is synchronous */
334+ rep.latency = shtb_stats(samples, spec.steps);
335+ rep.have_latency = 1;
336+ /* the range reported below is the state as it stands now: warmup included */
337+ rep.steps_run = spec.warmup + spec.steps;
338+ rep.model_t = transform_mode ? 0 : rep.steps_run * spec.params.dt;
339+
340+ /* Did the run stay finite and develop contrast? The app shows the same range
341+ * for the first species under its stats line. */
342+ if (transform_mode) {
343+ field_range((const double *)tq[tcur], 1, 2 * nlm, 0, &rep.field_min, &rep.field_max,
344+ &rep.finite);
345+ } else {
346+ for (int k = 0; k < nsp; k++) SH_to_spat(sht, Q[k], spat[k]);
347+ field_range(spat[0], g.nlines, g.linelen, g.stride, &rep.field_min, &rep.field_max,
348+ &rep.finite);
349+ }
350+
351+ /* A reproducible state to compare against a WebGPU run: exactly --steps steps
352+ * from the seed, separate from the timed run above (which has warmup in it). */
353+ if (spec.digest) {
354+ if (transform_mode) {
355+ seed_spectrum();
356+ rep.input_digest = shtb_digest_of(state32, 2 * nlm);
357+ rep.have_input_digest = 1;
358+ for (int i = 0; i < spec.steps; i++) round_trip();
359+ } else {
360+ seed_state();
361+ for (int i = 0; i < spec.steps; i++) step();
362+ }
363+ const double *q = (const double *)(transform_mode ? tq[tcur] : Q[0]);
364+ for (long i = 0; i < 2 * nlm; i++) state32[i] = (float)q[i];
365+ rep.digest = shtb_digest_of(state32, 2 * nlm);
366+ rep.have_digest = 1;
367+ }
368+
369+ if (spec.json) {
370+ shtb_print_json(&spec, &rep);
371+ } else {
372+ printf(" %.3f ms/step %.1f steps/s", rep.ms_per_step, 1000.0 / rep.ms_per_step);
373+ if (!transform_mode) printf(" %.2f model time/s", spec.params.dt * 1000.0 / rep.ms_per_step);
374+ printf("\n");
375+ printf(" per step: %.3f ms mean · median %.3f · p05 %.3f · p95 %.3f · min %.3f\n",
376+ rep.latency.mean_ms, rep.latency.median_ms, rep.latency.p05_ms, rep.latency.p95_ms,
377+ rep.latency.min_ms);
378+ if (transform_mode)
379+ printf(" i.e. %.3f ms per single transform\n", rep.ms_per_step / 2);
380+ else
381+ printf(" after %d steps: t = %.2f, field ∈ [%.4f, %.4f] (contrast %.4f)%s\n",
382+ rep.steps_run, rep.model_t, rep.field_min, rep.field_max,
383+ rep.field_max - rep.field_min, rep.finite ? "" : " — NOT FINITE");
384+ if (rep.have_digest) {
385+ printf("\n state after %d steps from seed %d:\n", spec.steps, spec.seed);
386+ printf(" n=%ld min=%.9g max=%.9g mean=%.9g rms=%.9g\n", rep.digest.n, rep.digest.min,
387+ rep.digest.max, rep.digest.mean, rep.digest.rms);
388+ }
389+ printf("\n This is fp64. The like-for-like fp32 comparison against the WGSL\n"
390+ " transforms is ./shtbench_gpu; this run is the accuracy reference.\n");
391+ }
392+
393+ if (spec.dump_state && rep.have_digest) {
394+ if (shtb_dump_state(spec.dump_state, &spec, &rep, state32, 2 * nlm) != 0) {
395+ fprintf(stderr, "shtbench: cannot write %s\n", spec.dump_state);
396+ return 1;
397+ }
398+ if (!spec.json) printf("\n wrote %s\n", spec.dump_state);
399+ }
400+
401+ free(samples);
402+ free(noise);
403+ free(state32);
404+ free(lam);
405+ for (int k = 0; k < nsp; k++) {
406+ shtns_free(spat[k]);
407+ shtns_free(rspat[k]);
408+ shtns_free(Q[k]);
409+ shtns_free(R[k]);
410+ }
411+ shtns_destroy(sht);
412+ return rep.finite ? 0 : 1;
413+}
bench/shtns/shtbench_gpu.cuadded+625−0View file
@@ -0,0 +1,625 @@
1+/*
2+ * The same run as `npm run bench`, on upstream SHTNS' own CUDA transforms.
3+ *
4+ * ./shtbench_gpu --preset schnak-spots --lmax 63 --steps 2000
5+ * ./shtbench_gpu --mode transform --lmax 63 --steps 2000
6+ *
7+ * This is the like-for-like comparison the WGSL transforms exist to be measured
8+ * against: single precision, everything resident on the GPU, nothing read back
9+ * inside the loop. The only difference between this and `npm run bench` is what
10+ * runs the transforms — SHTNS' hand-written CUDA kernels and cuFFT/VkFFT here,
11+ * generated WGSL and a WGSL FFT there — on the same device.
12+ *
13+ * It keeps the state on the GPU the same way the WebGPU side does: cu_* are the
14+ * on-device entry points, asynchronous on SHTNS' compute stream, and a batch of
15+ * steps is launched before anything is waited for. `--batch` matches
16+ * `npm run bench --batch`, so both sides can be made to synchronize equally
17+ * often.
18+ *
19+ * Two things are measured, selected by --mode:
20+ *
21+ * - solver: one IMEX Euler timestep of models/<key>.m — 2*nspecies
22+ * transforms, one reaction kernel on the grid, one spectral
23+ * update kernel.
24+ * - transform: one spectral -> grid -> spectral round trip and nothing else.
25+ *
26+ * Two SHTNS details worth knowing when reading the numbers:
27+ *
28+ * - in fp32 mode SHTNS runs the *Legendre recurrence* in fp64 when
29+ * lmax <= 128 and the GPU has usable fp64 (SHT_L_RESCALE_FLY_FLOAT in
30+ * sht_private.h), which WebGPU cannot do at all. Set
31+ * SHTNS_GPU_REC_PREC=1 to force the recurrence into fp32 and get the closer
32+ * comparison; the run prints which it got.
33+ * - the spatial layout defaults to theta-contiguous, which is SHTNS' native
34+ * and fastest. --layout phi matches what the WGSL side uses. The spectral
35+ * layout and normalization are identical either way, so a state comparison
36+ * is valid in both.
37+ */
38+#include "spec.h"
39+
40+#include <cuda_runtime.h>
41+#include <shtns.h>
42+#include <shtns_cuda.h>
43+
44+static const char *USAGE =
45+ "usage: ./shtbench_gpu [options]\n"
46+ "\n"
47+ " --mode solver|transform what to measure (default solver)\n"
48+ " --preset <key> schnak-spots | schnak-coarse | schnak-fine | brussel |\n"
49+ " allencahn (default schnak-spots)\n"
50+ " --lmax <n> spherical harmonic truncation (default 63)\n"
51+ " --steps <n> timed steps, or round trips (default 2000)\n"
52+ " --warmup <n> untimed steps first (default 100)\n"
53+ " --batch <n> steps launched per synchronization (default 16), as in\n"
54+ " `npm run bench -- --batch`\n"
55+ " --seed <n> seed of the initial noise / spectrum (default 1)\n"
56+ " --fp64 use SHTNS' double-precision GPU transforms instead of\n"
57+ " single. Not comparable to WebGPU, which has no fp64;\n"
58+ " useful as an accuracy and cost reference\n"
59+ " --layout theta|phi spatial layout: theta-contiguous is SHTNS' native and\n"
60+ " fastest, phi-contiguous is what the WGSL side uses\n"
61+ " (default theta)\n"
62+ " --polar-eps <x> SHTNS polar-optimization threshold (default 0)\n"
63+ " --device <n> CUDA device index (default 0)\n"
64+ " --digest after timing, re-run exactly --steps steps from the seed\n"
65+ " and print a digest of the final spectral state\n"
66+ " --dump-state <f> like --digest, and write the state to <f> as JSON, for\n"
67+ " scripts/compare-native.mjs to diff\n"
68+ " --<param> <v> any parameter of the preset's model, e.g. --dt 0.05\n"
69+ " --json machine-readable output\n"
70+ " --help\n"
71+ "\n"
72+ "Run the same spec through the WGSL transforms with:\n"
73+ " npm run bench -- --preset <key> --lmax <n> --steps <n> (solver)\n"
74+ " npm run bench:sht -- --lmax <n> --steps <n> (transform)";
75+
76+#define CU_CHECK(call) \
77+ do { \
78+ cudaError_t err_ = (call); \
79+ if (err_ != cudaSuccess) { \
80+ fprintf(stderr, "shtbench_gpu: %s failed at %s:%d: %s\n", #call, __FILE__, __LINE__, \
81+ cudaGetErrorString(err_)); \
82+ exit(1); \
83+ } \
84+ } while (0)
85+
86+/* ------------------------------------------------------------------- kernels
87+ *
88+ * The counterparts of the generated WGSL kernels: one thread per output
89+ * element, reading the same inputs and doing the same arithmetic (see
90+ * shtb_react / shtb_imex in spec.h, transcribed from the .m).
91+ */
92+
93+template <typename real>
94+__global__ void k_react(shtb_step_const<real> c, const real *u, const real *v, real *r1, real *r2,
95+ long npts) {
96+ const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
97+ if (i >= npts) return;
98+ real a = 0, b = 0;
99+ shtb_react<real>(c, u[i], v ? v[i] : (real)0, &a, &b);
100+ r1[i] = a;
101+ if (r2) r2[i] = b;
102+}
103+
104+/* Latitude lines are `stride` apart in the theta-contiguous layout, so the flat
105+ * kernel above would step over padding. This one is used when there is any. */
106+template <typename real>
107+__global__ void k_react_strided(shtb_step_const<real> c, const real *u, const real *v, real *r1,
108+ real *r2, long linelen, long stride) {
109+ const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
110+ if (i >= linelen) return;
111+ const long o = (long)blockIdx.y * stride + i;
112+ real a = 0, b = 0;
113+ shtb_react<real>(c, u[o], v ? v[o] : (real)0, &a, &b);
114+ r1[o] = a;
115+ if (r2) r2[o] = b;
116+}
117+
118+/* One thread per real, over the 2 x nlm spectral layout — `lam` carries l(l+1)
119+ * duplicated across the real and imaginary halves, exactly as on the WGSL side. */
120+template <typename real>
121+__global__ void k_imex(shtb_step_const<real> c, int k, real *U, const real *R, const real *lam,
122+ long n2) {
123+ const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
124+ if (i >= n2) return;
125+ U[i] = shtb_imex<real>(c, k, U[i], R[i], lam[i]);
126+}
127+
128+/* Scatter the seeded perturbation, which is generated in [ilat*nphi + iphi]
129+ * order, into whichever layout SHTNS is using. Also used to fill the uniform
130+ * background. */
131+template <typename real>
132+__global__ void k_fill(real *f, real value, const float *noise, long nlat, long nphi,
133+ long stride_lat, long stride_phi) {
134+ const long ilat = (long)blockIdx.y;
135+ const long iphi = (long)blockIdx.x * blockDim.x + threadIdx.x;
136+ if (iphi >= nphi || ilat >= nlat) return;
137+ const real n = noise ? (real)noise[ilat * nphi + iphi] : (real)0;
138+ f[ilat * stride_lat + iphi * stride_phi] = value + n;
139+}
140+
141+/* --------------------------------------------------------------------- setup */
142+
143+/* SHTNS' own device-buffer sizes, from init_cuda_buffer_fft() in sht_gpu.cu.
144+ * Its kernels write a little past nlm ("one more data per m") and the Fourier
145+ * stage needs room for the R2C form, so allocating exactly nlm or nlat*nphi is
146+ * not enough. WARPSZE is 32 on every CUDA GPU SHTNS supports. */
147+static long spec_alloc_reals(long nlm, int mmax) {
148+ const long nlm2 = nlm + (mmax + 1);
149+ return ((2 * nlm2 + 31) / 32) * 32;
150+}
151+static long spat_alloc_reals(long nlat_padded, long nphi, int mmax) {
152+ const long extra = (nphi / 2 == mmax) ? 1 : 0;
153+ return ((nlat_padded * (nphi + extra) + 31) / 32) * 32;
154+}
155+
156+struct Layout {
157+ long nlat, nphi, nlat_padded;
158+ long stride_lat, stride_phi;
159+ long linelen, nlines, stride; /* contiguous runs, for the reaction kernel */
160+ int padded;
161+};
162+
163+static Layout layout_of(shtns_cfg sht, int which) {
164+ Layout l;
165+ l.nlat = sht->nlat;
166+ l.nphi = sht->nphi;
167+ l.nlat_padded = sht->nlat_padded;
168+ if (which == SHTB_LAYOUT_PHI) {
169+ l.stride_lat = l.nphi;
170+ l.stride_phi = 1;
171+ l.nlines = l.nlat;
172+ l.linelen = l.nphi;
173+ l.stride = l.nphi;
174+ l.padded = 0;
175+ } else {
176+ l.stride_lat = 1;
177+ l.stride_phi = l.nlat_padded;
178+ l.nlines = l.nphi;
179+ l.linelen = l.nlat;
180+ l.stride = l.nlat_padded;
181+ l.padded = l.nlat_padded != l.nlat;
182+ }
183+ return l;
184+}
185+
186+/* --------------------------------------------------------------- the run body
187+ *
188+ * Templated on the transform precision so fp32 and fp64 are the same code. The
189+ * fp32 instantiation is the one that matters; fp64 is there as a reference.
190+ */
191+template <typename real>
192+struct Run {
193+ shtns_cfg sht;
194+ const shtb_spec *spec;
195+ /* The stream SHTNS was told to compute on, so our kernels are ordered against
196+ * its transforms and one synchronization waits for the whole step. */
197+ cudaStream_t stream;
198+ Layout lay;
199+ long nlm, n2, npts;
200+ int nsp;
201+ shtb_step_const<real> c;
202+ double base[2];
203+
204+ real *dU[2], *dR[2], *dspat[2], *drspat[2], *dlam;
205+ real *dTq[2]; /* transform mode ping-pong */
206+ int tcur;
207+ float *dnoise;
208+ float *hnoise;
209+ float *hstate;
210+
211+ void sync() { CU_CHECK(cudaStreamSynchronize(stream)); }
212+
213+ void alloc() {
214+ const long spec_n = spec_alloc_reals(nlm, sht->mmax);
215+ long spat_n = spat_alloc_reals(lay.nlat_padded, lay.nphi, sht->mmax);
216+ if ((long)sht->nspat > spat_n) spat_n = (long)sht->nspat;
217+ for (int k = 0; k < nsp; k++) {
218+ CU_CHECK(cudaMalloc(&dU[k], sizeof(real) * (size_t)spec_n));
219+ CU_CHECK(cudaMalloc(&dR[k], sizeof(real) * (size_t)spec_n));
220+ CU_CHECK(cudaMalloc(&dspat[k], sizeof(real) * (size_t)spat_n));
221+ CU_CHECK(cudaMalloc(&drspat[k], sizeof(real) * (size_t)spat_n));
222+ CU_CHECK(cudaMemset(dU[k], 0, sizeof(real) * (size_t)spec_n));
223+ CU_CHECK(cudaMemset(dR[k], 0, sizeof(real) * (size_t)spec_n));
224+ CU_CHECK(cudaMemset(dspat[k], 0, sizeof(real) * (size_t)spat_n));
225+ CU_CHECK(cudaMemset(drspat[k], 0, sizeof(real) * (size_t)spat_n));
226+ }
227+ dTq[0] = dU[0];
228+ dTq[1] = dR[0];
229+ tcur = 0;
230+ CU_CHECK(cudaMalloc(&dlam, sizeof(real) * (size_t)n2));
231+ CU_CHECK(cudaMalloc(&dnoise, sizeof(float) * (size_t)npts));
232+ hnoise = (float *)malloc(sizeof(float) * (size_t)npts);
233+ hstate = (float *)malloc(sizeof(float) * (size_t)(n2));
234+
235+ real *lam = (real *)malloc(sizeof(real) * (size_t)n2);
236+ for (long lm = 0; lm < nlm; lm++) {
237+ const int l = sht->li[lm];
238+ lam[2 * lm] = lam[2 * lm + 1] = (real)(l * (l + 1));
239+ }
240+ CU_CHECK(cudaMemcpy(dlam, lam, sizeof(real) * (size_t)n2, cudaMemcpyHostToDevice));
241+ free(lam);
242+ }
243+
244+ void free_all() {
245+ for (int k = 0; k < nsp; k++) {
246+ cudaFree(dU[k]);
247+ cudaFree(dR[k]);
248+ cudaFree(dspat[k]);
249+ cudaFree(drspat[k]);
250+ }
251+ cudaFree(dlam);
252+ cudaFree(dnoise);
253+ free(hnoise);
254+ free(hstate);
255+ }
256+
257+ void synth(real *qlm, real *spat);
258+ void analys(real *spat, real *qlm);
259+
260+ void seed_state() {
261+ shtb_seeded_noise(npts, spec->model->seed_amp, (uint32_t)spec->seed, hnoise);
262+ CU_CHECK(cudaMemcpy(dnoise, hnoise, sizeof(float) * (size_t)npts, cudaMemcpyHostToDevice));
263+ const int tpb = 256;
264+ dim3 grid((unsigned)((lay.nphi + tpb - 1) / tpb), (unsigned)lay.nlat);
265+ k_fill<real><<<grid, tpb, 0, stream>>>(dspat[0], (real)base[0], dnoise, lay.nlat, lay.nphi,
266+ lay.stride_lat, lay.stride_phi);
267+ analys(dspat[0], dU[0]);
268+ if (nsp > 1) {
269+ k_fill<real><<<grid, tpb, 0, stream>>>(dspat[1], (real)base[1], NULL, lay.nlat, lay.nphi,
270+ lay.stride_lat, lay.stride_phi);
271+ analys(dspat[1], dU[1]);
272+ }
273+ sync();
274+ }
275+
276+ /* One IMEX Euler timestep. Nothing is synchronized: the calls queue on SHTNS'
277+ * compute stream, which is what makes a batch of steps one submission's worth
278+ * of work, as on the WebGPU side. */
279+ void step() {
280+ const int tpb = 256;
281+ for (int k = 0; k < nsp; k++) synth(dU[k], dspat[k]);
282+ if (lay.padded) {
283+ dim3 grid((unsigned)((lay.linelen + tpb - 1) / tpb), (unsigned)lay.nlines);
284+ k_react_strided<real><<<grid, tpb, 0, stream>>>(c, dspat[0], nsp > 1 ? dspat[1] : NULL,
285+ drspat[0], nsp > 1 ? drspat[1] : NULL,
286+ lay.linelen, lay.stride);
287+ } else {
288+ k_react<real><<<(unsigned)((npts + tpb - 1) / tpb), tpb, 0, stream>>>(
289+ c, dspat[0], nsp > 1 ? dspat[1] : NULL, drspat[0], nsp > 1 ? drspat[1] : NULL, npts);
290+ }
291+ for (int k = 0; k < nsp; k++) analys(drspat[k], dR[k]);
292+ for (int k = 0; k < nsp; k++)
293+ k_imex<real><<<(unsigned)((n2 + tpb - 1) / tpb), tpb, 0, stream>>>(c, k, dU[k], dR[k], dlam,
294+ n2);
295+ }
296+
297+ void seed_spectrum() {
298+ shtb_seeded_spectrum(spec->lmax, spec->lmax, (uint32_t)spec->seed, hstate);
299+ tcur = 0;
300+ if (sizeof(real) == sizeof(float)) {
301+ CU_CHECK(cudaMemcpy(dTq[0], hstate, sizeof(float) * (size_t)n2, cudaMemcpyHostToDevice));
302+ } else {
303+ double *tmp = (double *)malloc(sizeof(double) * (size_t)n2);
304+ for (long i = 0; i < n2; i++) tmp[i] = (double)hstate[i];
305+ CU_CHECK(cudaMemcpy(dTq[0], tmp, sizeof(double) * (size_t)n2, cudaMemcpyHostToDevice));
306+ free(tmp);
307+ }
308+ sync();
309+ }
310+
311+ void round_trip() {
312+ synth(dTq[tcur], dspat[0]);
313+ analys(dspat[0], dTq[tcur ^ 1]);
314+ tcur ^= 1;
315+ }
316+
317+ /* the final spectral state of species 0 (or of the round trip), as float */
318+ const float *read_state() {
319+ sync();
320+ const real *src = spec->mode == SHTB_MODE_TRANSFORM ? dTq[tcur] : dU[0];
321+ if (sizeof(real) == sizeof(float)) {
322+ CU_CHECK(cudaMemcpy(hstate, src, sizeof(float) * (size_t)n2, cudaMemcpyDeviceToHost));
323+ } else {
324+ double *tmp = (double *)malloc(sizeof(double) * (size_t)n2);
325+ CU_CHECK(cudaMemcpy(tmp, src, sizeof(double) * (size_t)n2, cudaMemcpyDeviceToHost));
326+ for (long i = 0; i < n2; i++) hstate[i] = (float)tmp[i];
327+ free(tmp);
328+ }
329+ return hstate;
330+ }
331+
332+ /* species 0 on the grid, for the range check */
333+ void read_field(double *mn, double *mx, int *finite) {
334+ for (int k = 0; k < nsp; k++) synth(dU[k], dspat[k]);
335+ sync();
336+ const long n = lay.nlines * lay.stride;
337+ real *h = (real *)malloc(sizeof(real) * (size_t)n);
338+ CU_CHECK(cudaMemcpy(h, dspat[0], sizeof(real) * (size_t)n, cudaMemcpyDeviceToHost));
339+ *mn = INFINITY;
340+ *mx = -INFINITY;
341+ *finite = 1;
342+ for (long l = 0; l < lay.nlines; l++)
343+ for (long i = 0; i < lay.linelen; i++) {
344+ const double x = (double)h[l * lay.stride + i];
345+ if (x < *mn) *mn = x;
346+ if (x > *mx) *mx = x;
347+ if (!isfinite(x)) *finite = 0;
348+ }
349+ free(h);
350+ }
351+};
352+
353+template <>
354+void Run<float>::synth(float *qlm, float *spat) {
355+ cu_SH_to_spat_float(sht, (cplx_f *)qlm, spat, sht->lmax);
356+}
357+template <>
358+void Run<float>::analys(float *spat, float *qlm) {
359+ cu_spat_to_SH_float(sht, spat, (cplx_f *)qlm, sht->lmax);
360+}
361+template <>
362+void Run<double>::synth(double *qlm, double *spat) {
363+ cu_SH_to_spat(sht, (cplx *)qlm, spat, sht->lmax);
364+}
365+template <>
366+void Run<double>::analys(double *spat, double *qlm) {
367+ cu_spat_to_SH(sht, spat, (cplx *)qlm, sht->lmax);
368+}
369+
370+/* ----------------------------------------------------------------------- main */
371+
372+template <typename real>
373+static int run(shtns_cfg sht, cudaStream_t stream, const shtb_spec &spec, const char *adapter,
374+ const char *runtime, const char *cfg_info) {
375+ Run<real> r;
376+ r.sht = sht;
377+ r.spec = &spec;
378+ r.stream = stream;
379+ r.lay = layout_of(sht, spec.layout);
380+ r.nlm = (long)sht->nlm;
381+ r.n2 = 2 * r.nlm;
382+ r.npts = r.lay.nlat * r.lay.nphi;
383+ r.nsp = spec.model->nspecies;
384+ r.c = shtb_make_step_const<real>(spec.model, &spec.params);
385+ shtb_background(spec.model, &spec.params, r.base);
386+ r.alloc();
387+
388+ const int transform_mode = spec.mode == SHTB_MODE_TRANSFORM;
389+ const char *precision = sizeof(real) == 4 ? "fp32" : "fp64";
390+
391+ if (!spec.json) {
392+ printf("shtbench_gpu — upstream SHTNS on CUDA, %s only\n\n",
393+ transform_mode ? "transforms" : "solver");
394+ printf(" mode %s\n", transform_mode
395+ ? "transform (one synth + one analys per step)"
396+ : "solver (one IMEX Euler timestep per step)");
397+ if (!transform_mode) {
398+ printf(" preset %s (models/%s.m: %d species)\n", spec.preset->label, spec.model->key,
399+ r.nsp);
400+ printf(" params ");
401+ for (int i = 0; i < spec.model->nparams; i++)
402+ printf("%s=%g ", spec.model->params[i].key,
403+ *shtb_field_c(&spec.params, spec.model->params[i].off));
404+ printf("\n");
405+ }
406+ printf(" grid lmax %d · %ldx%ld · nlm %ld\n", spec.lmax, r.lay.nlat, r.lay.nphi, r.nlm);
407+ printf(" layout %s%s\n",
408+ spec.layout == SHTB_LAYOUT_PHI ? "phi-contiguous" : "theta-contiguous (native)",
409+ r.lay.padded ? ", padded" : "");
410+ printf(" backend %s\n %s, %s\n %s\n", adapter, precision, runtime,
411+ cfg_info ? cfg_info : "(no GPU config info)");
412+ printf(" run %d warmup + %d timed steps in batches of %d, seed %d\n\n", spec.warmup,
413+ spec.steps, spec.batch, spec.seed);
414+ }
415+
416+ if (transform_mode)
417+ r.seed_spectrum();
418+ else
419+ r.seed_state();
420+
421+ for (int i = 0; i < spec.warmup; i++) transform_mode ? r.round_trip() : r.step();
422+ r.sync();
423+
424+ /* --- throughput: a batch launched together, waited for once ------------- */
425+ const int batches = (spec.steps + spec.batch - 1) / spec.batch;
426+ double launch_ms = 0;
427+ int done = 0;
428+ const double t0 = shtb_now_ms();
429+ for (int b = 0; b < batches; b++) {
430+ const int n = spec.steps - done < spec.batch ? spec.steps - done : spec.batch;
431+ const double e0 = shtb_now_ms();
432+ for (int i = 0; i < n; i++) transform_mode ? r.round_trip() : r.step();
433+ launch_ms += shtb_now_ms() - e0;
434+ r.sync();
435+ done += n;
436+ }
437+ const double throughput_ms = (shtb_now_ms() - t0) / done;
438+
439+ /* --- latency: one step per synchronization, for the distribution -------- */
440+ const int lat_steps = spec.steps < 200 ? spec.steps : 200;
441+ double *samples = (double *)malloc(sizeof(double) * (size_t)lat_steps);
442+ for (int i = 0; i < lat_steps; i++) {
443+ const double a = shtb_now_ms();
444+ transform_mode ? r.round_trip() : r.step();
445+ r.sync();
446+ samples[i] = shtb_now_ms() - a;
447+ }
448+
449+ shtb_report rep;
450+ memset(&rep, 0, sizeof(rep));
451+ char libbuf[192], adapterbuf[160], runtimebuf[160];
452+ rep.library = shtb_json_safe(libbuf, sizeof(libbuf), shtns_get_build_info());
453+ rep.runtime = shtb_json_safe(runtimebuf, sizeof(runtimebuf), runtime);
454+ rep.adapter = shtb_json_safe(adapterbuf, sizeof(adapterbuf), adapter);
455+ rep.precision = precision;
456+ rep.fourier = "cufft/vkfft";
457+ rep.nlm = r.nlm;
458+ rep.ops_per_step = transform_mode ? 2 : 2 * r.nsp + 1 + r.nsp;
459+ rep.ms_per_step = throughput_ms;
460+ rep.encode_ms_per_step = launch_ms / done;
461+ rep.latency = shtb_stats(samples, lat_steps);
462+ rep.have_latency = 1;
463+ rep.steps_run = spec.warmup + spec.steps + lat_steps;
464+ rep.model_t = transform_mode ? 0 : rep.steps_run * spec.params.dt;
465+
466+ if (transform_mode) {
467+ const float *s = r.read_state();
468+ rep.field_min = INFINITY;
469+ rep.field_max = -INFINITY;
470+ rep.finite = shtb_all_finite(s, r.n2);
471+ for (long i = 0; i < r.n2; i++) {
472+ if (s[i] < rep.field_min) rep.field_min = s[i];
473+ if (s[i] > rep.field_max) rep.field_max = s[i];
474+ }
475+ } else {
476+ r.read_field(&rep.field_min, &rep.field_max, &rep.finite);
477+ }
478+
479+ /* A reproducible state to compare against a WebGPU run: exactly --steps steps
480+ * from the seed, separate from the timed runs above. */
481+ if (spec.digest) {
482+ if (transform_mode) {
483+ r.seed_spectrum();
484+ rep.input_digest = shtb_digest_of(r.hstate, r.n2);
485+ rep.have_input_digest = 1;
486+ for (int i = 0; i < spec.steps; i++) r.round_trip();
487+ } else {
488+ r.seed_state();
489+ for (int i = 0; i < spec.steps; i++) r.step();
490+ }
491+ rep.digest = shtb_digest_of(r.read_state(), r.n2);
492+ rep.have_digest = 1;
493+ }
494+
495+ if (spec.json) {
496+ shtb_print_json(&spec, &rep);
497+ } else {
498+ printf(" %.3f ms/step %.1f steps/s", rep.ms_per_step, 1000.0 / rep.ms_per_step);
499+ if (!transform_mode) printf(" %.2f model time/s", spec.params.dt * 1000.0 / rep.ms_per_step);
500+ printf(" (batches of %d)\n", spec.batch);
501+ printf(" of which CPU kernel launching: %.3f ms/step (%.0f%% — the rest is the GPU)\n",
502+ rep.encode_ms_per_step, 100.0 * rep.encode_ms_per_step / rep.ms_per_step);
503+ printf(" one step per sync: %.3f ms mean · median %.3f · p05 %.3f · p95 %.3f · min %.3f\n",
504+ rep.latency.mean_ms, rep.latency.median_ms, rep.latency.p05_ms, rep.latency.p95_ms,
505+ rep.latency.min_ms);
506+ if (transform_mode)
507+ printf(" i.e. %.3f ms per single transform\n", rep.ms_per_step / 2);
508+ else
509+ printf(" after %d steps: t = %.2f, field ∈ [%.4f, %.4f] (contrast %.4f)%s\n",
510+ rep.steps_run, rep.model_t, rep.field_min, rep.field_max,
511+ rep.field_max - rep.field_min, rep.finite ? "" : " — NOT FINITE");
512+ if (rep.have_digest) {
513+ printf("\n state after %d steps from seed %d:\n", spec.steps, spec.seed);
514+ printf(" n=%ld min=%.9g max=%.9g mean=%.9g rms=%.9g\n", rep.digest.n, rep.digest.min,
515+ rep.digest.max, rep.digest.mean, rep.digest.rms);
516+ }
517+ printf("\n Compare with `npm run bench --json` on this machine: same GPU, same\n"
518+ " precision, same grid — the difference is the transform implementation.\n"
519+ " scripts/compare-native.mjs runs both and lines the numbers up.\n");
520+ }
521+
522+ if (spec.dump_state && rep.have_digest) {
523+ if (shtb_dump_state(spec.dump_state, &spec, &rep, r.hstate, r.n2) != 0) {
524+ fprintf(stderr, "shtbench_gpu: cannot write %s\n", spec.dump_state);
525+ return 1;
526+ }
527+ if (!spec.json) printf("\n wrote %s\n", spec.dump_state);
528+ }
529+
530+ free(samples);
531+ r.free_all();
532+ return rep.finite ? 0 : 1;
533+}
534+
535+int main(int argc, char **argv) {
536+ shtb_spec spec;
537+ double polar_eps = 0.0;
538+ int device = 0;
539+
540+ /* --polar-eps and --device are ours, not part of the shared spec. */
541+ int argc2 = 0;
542+ char **argv2 = (char **)malloc(sizeof(char *) * (size_t)argc);
543+ for (int i = 0; i < argc; i++) {
544+ if (strcmp(argv[i], "--polar-eps") == 0 && i + 1 < argc) {
545+ polar_eps = atof(argv[++i]);
546+ continue;
547+ }
548+ if (strncmp(argv[i], "--polar-eps=", 12) == 0) {
549+ polar_eps = atof(argv[i] + 12);
550+ continue;
551+ }
552+ if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) {
553+ device = atoi(argv[++i]);
554+ continue;
555+ }
556+ if (strncmp(argv[i], "--device=", 9) == 0) {
557+ device = atoi(argv[i] + 9);
558+ continue;
559+ }
560+ argv2[argc2++] = argv[i];
561+ }
562+ int rc = shtb_parse_spec(argc2, argv2, &spec, USAGE);
563+ free(argv2);
564+ if (rc) return rc == 1 ? 0 : rc;
565+
566+ CU_CHECK(cudaSetDevice(device));
567+ cudaDeviceProp prop;
568+ CU_CHECK(cudaGetDeviceProperties(&prop, device));
569+ char adapter[128];
570+ snprintf(adapter, sizeof(adapter), "%s (sm_%d%d, %d SMs)", prop.name, prop.major, prop.minor,
571+ prop.multiProcessorCount);
572+ int rtv = 0, drv = 0;
573+ cudaRuntimeGetVersion(&rtv);
574+ cudaDriverGetVersion(&drv);
575+ char runtime[128];
576+ snprintf(runtime, sizeof(runtime), "CUDA runtime %d.%d, driver %d.%d", rtv / 1000,
577+ (rtv % 1000) / 10, drv / 1000, (drv % 1000) / 10);
578+
579+ shtns_verbose(0);
580+ shtns_cfg sht = shtns_create(spec.lmax, spec.lmax, 1, sht_orthonormal);
581+ if (!sht) {
582+ fprintf(stderr, "shtbench_gpu: shtns_create failed\n");
583+ return 1;
584+ }
585+ /* Our reaction and update kernels have to be ordered against SHTNS'
586+ * transforms, so both go on one stream we own. cushtns_set_streams must come
587+ * before shtns_set_grid, which is where the GPU (and its FFT plan) is set up. */
588+ cudaStream_t stream = 0;
589+ CU_CHECK(cudaStreamCreate(&stream));
590+ cushtns_set_streams(sht, stream, 0);
591+
592+ int flags = sht_gauss | SHT_ALLOW_GPU | SHT_SCALAR_ONLY;
593+ flags |= spec.layout == SHTB_LAYOUT_PHI ? SHT_PHI_CONTIGUOUS : SHT_THETA_CONTIGUOUS;
594+ if (spec.fp32) flags |= SHT_FP32;
595+ if (shtns_set_grid(sht, (enum shtns_type)flags, polar_eps, spec.nlat, spec.nphi) <= 0) {
596+ fprintf(stderr, "shtbench_gpu: shtns_set_grid failed for lmax %d on a %dx%d grid\n", spec.lmax,
597+ spec.nlat, spec.nphi);
598+ return 1;
599+ }
600+ if ((int)sht->nlat != spec.nlat || (int)sht->nphi != spec.nphi) {
601+ fprintf(stderr, "shtbench_gpu: SHTNS chose a %ux%u grid, not the %dx%d asked for\n", sht->nlat,
602+ sht->nphi, spec.nlat, spec.nphi);
603+ return 1;
604+ }
605+
606+ /* cushtns_get_cfg_info() returns NULL when the GPU was never initialized, which
607+ * is how SHT_ALLOW_GPU failing shows up — and with SHT_FP32 the CPU fallback
608+ * would read fp64 out of fp32 buffers, so stop rather than produce a number
609+ * for the wrong thing. */
610+ const char *cfg_info = cushtns_get_cfg_info(sht);
611+ if (!cfg_info) {
612+ fprintf(stderr,
613+ "shtbench_gpu: SHTNS did not initialize the GPU for this grid (lmax %d, %dx%d).\n"
614+ " Was it built with --enable-cuda, and does nlat %% 4 == 0 hold for the\n"
615+ " theta-contiguous layout? Try --layout phi.\n",
616+ spec.lmax, spec.nlat, spec.nphi);
617+ return 1;
618+ }
619+
620+ const int status = spec.fp32 ? run<float>(sht, stream, spec, adapter, runtime, cfg_info)
621+ : run<double>(sht, stream, spec, adapter, runtime, cfg_info);
622+ shtns_destroy(sht);
623+ cudaStreamDestroy(stream);
624+ return status;
625+}
bench/shtns/spec.hadded+708−0View file
@@ -0,0 +1,708 @@
1+/*
2+ * One run, described the same way src/bench/runSpec.ts describes it.
3+ *
4+ * This is the native half of the comparison: the same grid rule, the same
5+ * presets and defaults, the same seeded perturbation, so that
6+ *
7+ * npm run bench -- --preset schnak-spots --lmax 63 --steps 2000
8+ * ./shtbench --preset schnak-spots --lmax 63 --steps 2000
9+ *
10+ * describe the same computation, one through the WGSL transforms and one
11+ * through upstream SHTNS.
12+ *
13+ * The tables below are a hand copy of src/mgpu/registry.ts and
14+ * src/sht/layout.ts, which is unavoidable — C cannot import the TypeScript.
15+ * It is also the one place the two sides could silently drift apart, so both
16+ * emit their resolved spec in --json and scripts/compare-native.mjs refuses to
17+ * compare two runs whose specs do not match.
18+ *
19+ * Compiled as C++ (by g++ for the CPU benchmark, by nvcc for the CUDA one), so
20+ * the reaction can be written once as a template and used at float and double.
21+ */
22+#ifndef SHTB_SPEC_H
23+#define SHTB_SPEC_H
24+
25+#include <math.h>
26+#include <stddef.h>
27+#include <stdint.h>
28+#include <stdio.h>
29+#include <stdlib.h>
30+#include <string.h>
31+#include <time.h>
32+
33+#if defined(__CUDACC__)
34+#define SHTB_HD __host__ __device__
35+#else
36+#define SHTB_HD
37+#endif
38+
39+/* ------------------------------------------------------------------ models */
40+
41+enum shtb_model_id { SHTB_SCHNAKENBERG = 0, SHTB_BRUSSELATOR, SHTB_ALLENCAHN };
42+
43+/* Every tunable scalar of every model. A model uses the subset its .m names as
44+ * arguments; the tables below say which, in the .m's declared order. */
45+struct shtb_params {
46+ double a, b; /* schnakenberg */
47+ double A, B; /* brusselator */
48+ double eps2; /* allen-cahn */
49+ double D1, D2; /* diffusivities (allen-cahn uses eps2 instead) */
50+ double dt;
51+};
52+
53+struct shtb_param_def {
54+ const char *key;
55+ size_t off; /* offset into shtb_params */
56+ double dflt;
57+};
58+
59+#define SHTB_P(field, dflt) \
60+ { #field, offsetof(struct shtb_params, field), dflt }
61+
62+static const struct shtb_param_def SHTB_SCHNAK_PARAMS[] = {
63+ SHTB_P(a, 0.1), SHTB_P(b, 0.9), SHTB_P(D1, 4e-4), SHTB_P(D2, 8e-3), SHTB_P(dt, 0.05),
64+};
65+static const struct shtb_param_def SHTB_BRUSSEL_PARAMS[] = {
66+ SHTB_P(A, 3.0), SHTB_P(B, 9.0), SHTB_P(D1, 3.33e-3), SHTB_P(D2, 1.67e-2), SHTB_P(dt, 0.02),
67+};
68+static const struct shtb_param_def SHTB_ALLENCAHN_PARAMS[] = {
69+ SHTB_P(eps2, 1e-3), SHTB_P(dt, 0.02),
70+};
71+
72+struct shtb_model_def {
73+ int id;
74+ const char *key;
75+ const char *label;
76+ int nspecies;
77+ int pdeg; /* polynomial degree of the reaction, for dealiasing */
78+ double seed_amp; /* amplitude of the seeded perturbation */
79+ const struct shtb_param_def *params;
80+ int nparams;
81+};
82+
83+static const struct shtb_model_def SHTB_MODELS[] = {
84+ {SHTB_SCHNAKENBERG, "schnakenberg", "Schnakenberg", 2, 3, 1e-2, SHTB_SCHNAK_PARAMS, 5},
85+ {SHTB_BRUSSELATOR, "brusselator", "Brusselator", 2, 3, 1e-2, SHTB_BRUSSEL_PARAMS, 5},
86+ {SHTB_ALLENCAHN, "allencahn", "Allen-Cahn", 1, 3, 1e-2, SHTB_ALLENCAHN_PARAMS, 2},
87+};
88+#define SHTB_NMODELS 3
89+
90+struct shtb_override {
91+ const char *key;
92+ double value;
93+};
94+
95+struct shtb_preset {
96+ const char *key;
97+ const char *label;
98+ int model;
99+ struct shtb_override over[2]; /* .key == NULL terminates */
100+};
101+
102+static const struct shtb_preset SHTB_PRESETS[] = {
103+ {"schnak-spots", "Schnakenberg - spots", SHTB_SCHNAKENBERG, {{NULL, 0}, {NULL, 0}}},
104+ {"schnak-coarse", "Schnakenberg - coarse spots", SHTB_SCHNAKENBERG,
105+ {{"D1", 1e-3}, {"D2", 2e-2}}},
106+ {"schnak-fine", "Schnakenberg - fine spots", SHTB_SCHNAKENBERG,
107+ {{"D1", 1.6e-4}, {"D2", 3.2e-3}}},
108+ {"brussel", "Brusselator - stripes & spots", SHTB_BRUSSELATOR, {{NULL, 0}, {NULL, 0}}},
109+ {"allencahn", "Allen-Cahn - coarsening", SHTB_ALLENCAHN, {{NULL, 0}, {NULL, 0}}},
110+};
111+#define SHTB_NPRESETS 5
112+
113+static inline double *shtb_field(struct shtb_params *p, size_t off) {
114+ return (double *)((char *)p + off);
115+}
116+static inline const double *shtb_field_c(const struct shtb_params *p, size_t off) {
117+ return (const double *)((const char *)p + off);
118+}
119+
120+/* --------------------------------------------------- the arithmetic per step */
121+
122+/*
123+ * Everything a timestep needs, in the working precision. Products the .m forms
124+ * from parameters — `(dt * D1)`, `(B + 1)` — are formed here in that same
125+ * precision, so a float run multiplies floats exactly as the WGSL kernel does.
126+ */
127+template <typename real>
128+struct shtb_step_const {
129+ int model;
130+ real p0, p1; /* (a, b) | (A, B) | unused */
131+ real dt;
132+ real dtD[2]; /* dt * D_k, one per species */
133+};
134+
135+/*
136+ * The reaction, transcribed from models/<key>.m. One line of MATLAB per line
137+ * here, and `u.^3` is written out as `u*u*u` because that is what the WGSL
138+ * backend emits for it (pow_i3; see emitPower in src/mgpu/wgsl.ts).
139+ *
140+ * schnakenberg: Un = (U + dt*analys(a - u + uuv)) ./ (1 + (dt*D1)*lam)
141+ * Vn = (V + dt*analys(b - uuv)) ./ (1 + (dt*D2)*lam)
142+ * brusselator: Un = (U + dt*analys(A - (B+1)*u + uuv)) ./ ...
143+ * Vn = (V + dt*analys(B*u - uuv)) ./ ...
144+ * allencahn: Un = (U + dt*analys(u - u.^3)) ./ (1 + (dt*eps2)*lam)
145+ */
146+template <typename real>
147+SHTB_HD inline void shtb_react(const shtb_step_const<real> &c, real u, real v, real *r1,
148+ real *r2) {
149+ if (c.model == SHTB_ALLENCAHN) {
150+ *r1 = u - u * u * u;
151+ return;
152+ }
153+ const real uuv = u * u * v;
154+ if (c.model == SHTB_SCHNAKENBERG) {
155+ *r1 = c.p0 - u + uuv;
156+ *r2 = c.p1 - uuv;
157+ } else {
158+ *r1 = c.p0 - (c.p1 + (real)1) * u + uuv;
159+ *r2 = c.p1 * u - uuv;
160+ }
161+}
162+
163+/* One element of the IMEX update, for species k. `lam` is l(l+1) of that
164+ * coefficient; the array is 2*nlm long with the value duplicated across the
165+ * real and imaginary halves, matching the 2 x nlm spectral layout the .m sees. */
166+template <typename real>
167+SHTB_HD inline real shtb_imex(const shtb_step_const<real> &c, int k, real U, real R, real lam) {
168+ return (U + c.dt * R) / ((real)1 + c.dtD[k] * lam);
169+}
170+
171+/* ------------------------------------------------------------------ the spec */
172+
173+enum shtb_mode { SHTB_MODE_SOLVER = 0, SHTB_MODE_TRANSFORM };
174+enum shtb_layout { SHTB_LAYOUT_THETA = 0, SHTB_LAYOUT_PHI };
175+
176+/* Defaults, from src/bench/runSpec.ts. */
177+#define SHTB_DEFAULT_LMAX 63
178+#define SHTB_DEFAULT_SEED 1
179+#define SHTB_DEFAULT_STEPS 2000
180+#define SHTB_DEFAULT_WARMUP 100
181+#define SHTB_DEFAULT_BATCH 16
182+
183+struct shtb_spec {
184+ const struct shtb_preset *preset;
185+ const struct shtb_model_def *model;
186+ struct shtb_params params;
187+ int lmax, nlat, nphi;
188+ int seed, steps, warmup, batch;
189+ int mode;
190+ int layout;
191+ int fp32; /* GPU only: use SHTNS' single-precision transforms */
192+ int threads; /* CPU only: OpenMP threads, 0 = library default */
193+ int json, digest;
194+ const char *dump_state;
195+};
196+
197+/* Grid sizes for a given lmax, dealiased for a reaction of polynomial degree
198+ * pdeg. Identical to gridForLmax() in src/sht/layout.ts, including rounding
199+ * nphi up to a power of two — which the WGSL side needs for its FFT path and
200+ * which SHTNS does not, but the grids have to match to compare anything. */
201+static inline void shtb_grid_for_lmax(int lmax, int pdeg, int *nlat, int *nphi) {
202+ double min_lat = ((double)(pdeg + 1) * lmax + 1) / 2.0;
203+ if (min_lat < lmax + 1) min_lat = lmax + 1;
204+ *nlat = 2 * (int)ceil(min_lat / 2.0);
205+ int n = 1;
206+ while (n < (pdeg + 1) * lmax + 1) n *= 2;
207+ *nphi = n;
208+}
209+
210+static inline long shtb_nlm_calc(int lmax, int mmax) {
211+ return (long)(mmax + 1) * (lmax + 1) - (long)mmax * (mmax + 1) / 2;
212+}
213+
214+/* -------------------------------------------------------------- the seeding */
215+
216+/*
217+ * mulberry32 + Box-Muller, transcribed from src/mgpu/noise.ts so an integer
218+ * seed means the same perturbation on both sides. The integer state evolves
219+ * bit-identically; the Box-Muller step then goes through log/sqrt/sin/cos, so a
220+ * libm that rounds differently from V8's can differ in the last bit. Both sides
221+ * report the perturbation's RMS for that reason.
222+ */
223+struct shtb_rng {
224+ uint32_t s;
225+ int have_spare;
226+ double spare;
227+};
228+
229+static inline void shtb_rng_init(struct shtb_rng *r, uint32_t seed) {
230+ r->s = seed;
231+ r->have_spare = 0;
232+ r->spare = 0;
233+}
234+
235+/* uniform in [0, 1) */
236+static inline double shtb_rand(struct shtb_rng *r) {
237+ r->s = r->s + 0x6d2b79f5u;
238+ uint32_t t = r->s;
239+ t = (t ^ (t >> 15)) * (t | 1u);
240+ t ^= t + (t ^ (t >> 7)) * (t | 61u);
241+ return (double)(t ^ (t >> 14)) / 4294967296.0;
242+}
243+
244+static inline double shtb_randn(struct shtb_rng *r) {
245+ if (r->have_spare) {
246+ r->have_spare = 0;
247+ return r->spare;
248+ }
249+ double u = 0;
250+ while (u == 0) u = shtb_rand(r);
251+ double rad = sqrt(-2 * log(u));
252+ double th = 2 * M_PI * shtb_rand(r);
253+ r->spare = rad * sin(th);
254+ r->have_spare = 1;
255+ return rad * cos(th);
256+}
257+
258+/* amp-scaled normal deviates, one per grid point, in [ilat*nphi + iphi] order —
259+ * the order src/mgpu/noise.ts produces them in. Rounded to float, because the
260+ * browser stores them in a Float32Array; the fp64 run then differs from the
261+ * fp32 one only in the arithmetic, not in the initial condition. */
262+static inline void shtb_seeded_noise(long npts, double amp, uint32_t seed, float *out) {
263+ struct shtb_rng r;
264+ shtb_rng_init(&r, seed);
265+ for (long i = 0; i < npts; i++) out[i] = (float)(amp * shtb_randn(&r));
266+}
267+
268+/*
269+ * A seeded spectrum for the transform-only benchmark: uniform in [-1, 1), and
270+ * bit-identical to the TypeScript side because it never leaves integer
271+ * arithmetic and exactly-representable doubles. The m = 0 imaginary parts are
272+ * zeroed, since a real field has none and the two libraries need not agree on
273+ * what to do with a coefficient that cannot occur.
274+ *
275+ * qlm is interleaved [re, im] per coefficient, SHTNS LM ordering.
276+ */
277+static inline void shtb_seeded_spectrum(int lmax, int mmax, uint32_t seed, float *qlm) {
278+ struct shtb_rng r;
279+ shtb_rng_init(&r, seed);
280+ long lm = 0;
281+ for (int m = 0; m <= mmax; m++) {
282+ for (int l = m; l <= lmax; l++, lm++) {
283+ qlm[2 * lm] = (float)(2 * shtb_rand(&r) - 1);
284+ float im = (float)(2 * shtb_rand(&r) - 1);
285+ qlm[2 * lm + 1] = (m == 0) ? 0.0f : im;
286+ }
287+ }
288+}
289+
290+/* -------------------------------------------------------------- statistics */
291+
292+static inline double shtb_now_ms(void) {
293+ struct timespec ts;
294+ clock_gettime(CLOCK_MONOTONIC, &ts);
295+ return (double)ts.tv_sec * 1e3 + (double)ts.tv_nsec * 1e-6;
296+}
297+
298+struct shtb_timing {
299+ double mean_ms, median_ms, p05_ms, p95_ms, min_ms;
300+};
301+
302+static int shtb_cmp_double(const void *a, const void *b) {
303+ double x = *(const double *)a, y = *(const double *)b;
304+ return (x > y) - (x < y);
305+}
306+
307+static inline struct shtb_timing shtb_stats(double *samples, int n) {
308+ qsort(samples, (size_t)n, sizeof(double), shtb_cmp_double);
309+ double total = 0;
310+ for (int i = 0; i < n; i++) total += samples[i];
311+ int i50 = (int)(0.50 * n), i05 = (int)(0.05 * n), i95 = (int)(0.95 * n);
312+ if (i50 >= n) i50 = n - 1;
313+ if (i05 >= n) i05 = n - 1;
314+ if (i95 >= n) i95 = n - 1;
315+ struct shtb_timing t;
316+ t.mean_ms = total / n;
317+ t.median_ms = samples[i50];
318+ t.p05_ms = samples[i05];
319+ t.p95_ms = samples[i95];
320+ t.min_ms = samples[0];
321+ return t;
322+}
323+
324+/* The same five numbers digestOf() computes in src/mgpu/digest.ts. */
325+struct shtb_digest {
326+ long n;
327+ double min, max, mean, rms;
328+};
329+
330+static inline struct shtb_digest shtb_digest_of(const float *v, long n) {
331+ struct shtb_digest d;
332+ d.n = n;
333+ d.min = INFINITY;
334+ d.max = -INFINITY;
335+ double sum = 0, sumsq = 0;
336+ for (long i = 0; i < n; i++) {
337+ double x = v[i];
338+ if (x < d.min) d.min = x;
339+ if (x > d.max) d.max = x;
340+ sum += x;
341+ sumsq += x * x;
342+ }
343+ d.mean = sum / (double)n;
344+ d.rms = sqrt(sumsq / (double)n);
345+ return d;
346+}
347+
348+static inline int shtb_all_finite(const float *v, long n) {
349+ for (long i = 0; i < n; i++)
350+ if (!isfinite(v[i])) return 0;
351+ return 1;
352+}
353+
354+/* ----------------------------------------------------------- spec resolution */
355+
356+static inline const struct shtb_model_def *shtb_model_by_id(int id) {
357+ for (int i = 0; i < SHTB_NMODELS; i++)
358+ if (SHTB_MODELS[i].id == id) return &SHTB_MODELS[i];
359+ return NULL;
360+}
361+
362+static inline void shtb_default_params(const struct shtb_model_def *m, struct shtb_params *p) {
363+ memset(p, 0, sizeof(*p));
364+ for (int i = 0; i < m->nparams; i++) *shtb_field(p, m->params[i].off) = m->params[i].dflt;
365+}
366+
367+/* Homogeneous background each species starts from, and the diffusivity each is
368+ * advanced with. From the init/step functions of models/<key>.m; only species 0
369+ * gets the seeded perturbation. */
370+static inline void shtb_background(const struct shtb_model_def *m, const struct shtb_params *p,
371+ double base[2]) {
372+ base[0] = base[1] = 0;
373+ if (m->id == SHTB_SCHNAKENBERG) {
374+ double us = p->a + p->b;
375+ base[0] = us;
376+ base[1] = p->b / (us * us);
377+ } else if (m->id == SHTB_BRUSSELATOR) {
378+ base[0] = p->A;
379+ base[1] = p->B / p->A;
380+ }
381+}
382+
383+static inline void shtb_diffusivity(const struct shtb_model_def *m, const struct shtb_params *p,
384+ double d[2]) {
385+ if (m->id == SHTB_ALLENCAHN) {
386+ d[0] = p->eps2;
387+ d[1] = 0;
388+ } else {
389+ d[0] = p->D1;
390+ d[1] = p->D2;
391+ }
392+}
393+
394+template <typename real>
395+static inline shtb_step_const<real> shtb_make_step_const(const struct shtb_model_def *m,
396+ const struct shtb_params *p) {
397+ double d[2];
398+ shtb_diffusivity(m, p, d);
399+ shtb_step_const<real> c;
400+ c.model = m->id;
401+ c.p0 = (real)(m->id == SHTB_BRUSSELATOR ? p->A : p->a);
402+ c.p1 = (real)(m->id == SHTB_BRUSSELATOR ? p->B : p->b);
403+ c.dt = (real)p->dt;
404+ /* (dt * D_k) as one product in the working precision, as the .m writes it */
405+ c.dtD[0] = (real)p->dt * (real)d[0];
406+ c.dtD[1] = (real)p->dt * (real)d[1];
407+ return c;
408+}
409+
410+/* ------------------------------------------------------------ argument parsing
411+ *
412+ * `--key value` or `--key=value`, in any order — the same grammar parseArgs()
413+ * accepts in src/bench/runSpec.ts, plus the flags only a native run has.
414+ */
415+
416+static inline int shtb_parse_spec(int argc, char **argv, struct shtb_spec *s, const char *usage) {
417+ const struct shtb_preset *preset = &SHTB_PRESETS[0];
418+
419+ /* --preset first: it decides which parameter names are legal. */
420+ for (int i = 1; i < argc; i++) {
421+ const char *a = argv[i];
422+ const char *v = NULL;
423+ if (strcmp(a, "--preset") == 0 && i + 1 < argc)
424+ v = argv[i + 1];
425+ else if (strncmp(a, "--preset=", 9) == 0)
426+ v = a + 9;
427+ if (!v) continue;
428+ const struct shtb_preset *found = NULL;
429+ for (int k = 0; k < SHTB_NPRESETS; k++)
430+ if (strcmp(SHTB_PRESETS[k].key, v) == 0) found = &SHTB_PRESETS[k];
431+ if (!found) {
432+ fprintf(stderr, "shtbench: unknown preset '%s' (have:", v);
433+ for (int k = 0; k < SHTB_NPRESETS; k++) fprintf(stderr, " %s", SHTB_PRESETS[k].key);
434+ fprintf(stderr, ")\n");
435+ return 2;
436+ }
437+ preset = found;
438+ }
439+
440+ memset(s, 0, sizeof(*s));
441+ s->preset = preset;
442+ s->model = shtb_model_by_id(preset->model);
443+ shtb_default_params(s->model, &s->params);
444+ for (int k = 0; k < 2; k++)
445+ if (preset->over[k].key) {
446+ for (int i = 0; i < s->model->nparams; i++)
447+ if (strcmp(s->model->params[i].key, preset->over[k].key) == 0)
448+ *shtb_field(&s->params, s->model->params[i].off) = preset->over[k].value;
449+ }
450+ s->lmax = SHTB_DEFAULT_LMAX;
451+ s->seed = SHTB_DEFAULT_SEED;
452+ s->steps = SHTB_DEFAULT_STEPS;
453+ s->warmup = SHTB_DEFAULT_WARMUP;
454+ s->batch = SHTB_DEFAULT_BATCH;
455+ s->mode = SHTB_MODE_SOLVER;
456+ s->layout = SHTB_LAYOUT_THETA;
457+ s->fp32 = 1;
458+ s->threads = 0;
459+
460+ for (int i = 1; i < argc; i++) {
461+ const char *a = argv[i];
462+ if (strncmp(a, "--", 2) != 0) {
463+ fprintf(stderr, "shtbench: unexpected argument '%s'\n\n%s\n", a, usage);
464+ return 2;
465+ }
466+ if (strcmp(a, "--help") == 0 || strcmp(a, "-h") == 0) {
467+ printf("%s\n", usage);
468+ return 1;
469+ }
470+ if (strcmp(a, "--json") == 0) {
471+ s->json = 1;
472+ continue;
473+ }
474+ if (strcmp(a, "--digest") == 0) {
475+ s->digest = 1;
476+ continue;
477+ }
478+ if (strcmp(a, "--fp64") == 0) {
479+ s->fp32 = 0;
480+ continue;
481+ }
482+
483+ /* split key / value */
484+ char key[64];
485+ const char *val = NULL;
486+ const char *eq = strchr(a, '=');
487+ if (eq) {
488+ size_t n = (size_t)(eq - a - 2);
489+ if (n >= sizeof(key)) n = sizeof(key) - 1;
490+ memcpy(key, a + 2, n);
491+ key[n] = 0;
492+ val = eq + 1;
493+ } else {
494+ snprintf(key, sizeof(key), "%s", a + 2);
495+ if (i + 1 >= argc) {
496+ fprintf(stderr, "shtbench: --%s needs a value\n", key);
497+ return 2;
498+ }
499+ val = argv[++i];
500+ }
501+
502+ if (strcmp(key, "preset") == 0) continue; /* handled above */
503+ if (strcmp(key, "lmax") == 0) {
504+ s->lmax = atoi(val);
505+ continue;
506+ }
507+ if (strcmp(key, "seed") == 0) {
508+ s->seed = atoi(val);
509+ continue;
510+ }
511+ if (strcmp(key, "steps") == 0) {
512+ s->steps = atoi(val);
513+ continue;
514+ }
515+ if (strcmp(key, "warmup") == 0) {
516+ s->warmup = atoi(val);
517+ continue;
518+ }
519+ if (strcmp(key, "batch") == 0) {
520+ s->batch = atoi(val);
521+ continue;
522+ }
523+ if (strcmp(key, "threads") == 0) {
524+ s->threads = atoi(val);
525+ continue;
526+ }
527+ if (strcmp(key, "dump-state") == 0) {
528+ s->dump_state = val;
529+ s->digest = 1;
530+ continue;
531+ }
532+ if (strcmp(key, "mode") == 0) {
533+ if (strcmp(val, "solver") == 0)
534+ s->mode = SHTB_MODE_SOLVER;
535+ else if (strcmp(val, "transform") == 0)
536+ s->mode = SHTB_MODE_TRANSFORM;
537+ else {
538+ fprintf(stderr, "shtbench: --mode must be 'solver' or 'transform' (got '%s')\n", val);
539+ return 2;
540+ }
541+ continue;
542+ }
543+ if (strcmp(key, "layout") == 0) {
544+ if (strcmp(val, "theta") == 0)
545+ s->layout = SHTB_LAYOUT_THETA;
546+ else if (strcmp(val, "phi") == 0)
547+ s->layout = SHTB_LAYOUT_PHI;
548+ else {
549+ fprintf(stderr, "shtbench: --layout must be 'theta' or 'phi' (got '%s')\n", val);
550+ return 2;
551+ }
552+ continue;
553+ }
554+
555+ int matched = 0;
556+ for (int p = 0; p < s->model->nparams; p++)
557+ if (strcmp(s->model->params[p].key, key) == 0) {
558+ *shtb_field(&s->params, s->model->params[p].off) = atof(val);
559+ matched = 1;
560+ }
561+ if (!matched) {
562+ fprintf(stderr, "shtbench: unknown option --%s\nparameters of %s:", key, s->model->label);
563+ for (int p = 0; p < s->model->nparams; p++)
564+ fprintf(stderr, " --%s", s->model->params[p].key);
565+ fprintf(stderr, "\n");
566+ return 2;
567+ }
568+ }
569+
570+ if (s->lmax < 1) {
571+ fprintf(stderr, "shtbench: --lmax must be >= 1\n");
572+ return 2;
573+ }
574+ if (s->steps < 1 || s->warmup < 0 || s->batch < 1) {
575+ fprintf(stderr, "shtbench: --steps and --batch must be >= 1, --warmup >= 0\n");
576+ return 2;
577+ }
578+ shtb_grid_for_lmax(s->lmax, s->model->pdeg, &s->nlat, &s->nphi);
579+ return 0;
580+}
581+
582+/* ------------------------------------------------------------- JSON emission
583+ *
584+ * Deliberately shaped like the object scripts/bench.ts prints with --json, so
585+ * scripts/compare-native.mjs can read a WebGPU run and a native run the same
586+ * way. "step" means one solver timestep in solver mode and one
587+ * spectral->grid->spectral round trip in transform mode.
588+ */
589+
590+/* Strings from the library go into JSON, so strip anything that would break it.
591+ * Returns `dst`, for use inline. */
592+static inline char *shtb_json_safe(char *dst, size_t cap, const char *src) {
593+ size_t j = 0;
594+ for (size_t i = 0; src && src[i] && j + 1 < cap; i++) {
595+ unsigned char c = (unsigned char)src[i];
596+ dst[j++] = (c < 0x20 || c == '"' || c == '\\' || c == 0x7f) ? ' ' : (char)c;
597+ }
598+ dst[j] = 0;
599+ return dst;
600+}
601+
602+struct shtb_report {
603+ const char *library; /* e.g. "SHTNS 3.7.5" */
604+ const char *runtime; /* e.g. "cuda 12.4, vkfft" */
605+ const char *adapter; /* GPU name, or the CPU's thread count */
606+ const char *precision;/* "fp32" | "fp64" */
607+ const char *fourier; /* which FFT the library used */
608+ long nlm;
609+ int ops_per_step;
610+ double ms_per_step;
611+ double encode_ms_per_step;
612+ struct shtb_timing latency;
613+ int have_latency;
614+ struct shtb_digest digest;
615+ int have_digest;
616+ struct shtb_digest input_digest; /* transform mode: the seeded spectrum */
617+ int have_input_digest;
618+ double field_min, field_max;
619+ int finite;
620+ double model_t;
621+ int steps_run;
622+};
623+
624+static inline void shtb_print_json(const struct shtb_spec *s, const struct shtb_report *r) {
625+ printf("{\n");
626+ printf(" \"mode\": \"%s\",\n", s->mode == SHTB_MODE_SOLVER ? "solver" : "transform");
627+ printf(" \"spec\": {\n");
628+ printf(" \"preset\": \"%s\",\n", s->preset->key);
629+ printf(" \"lmax\": %d,\n", s->lmax);
630+ printf(" \"seed\": %d,\n", s->seed);
631+ printf(" \"steps\": %d,\n", s->steps);
632+ printf(" \"warmup\": %d,\n", s->warmup);
633+ printf(" \"params\": {");
634+ for (int i = 0; i < s->model->nparams; i++)
635+ printf("%s\"%s\": %.17g", i ? ", " : "", s->model->params[i].key,
636+ *shtb_field_c(&s->params, s->model->params[i].off));
637+ printf("}\n },\n");
638+ printf(" \"model\": \"%s\",\n", s->model->key);
639+ printf(" \"backend\": {\"library\": \"%s\", \"runtime\": \"%s\", \"adapter\": \"%s\", "
640+ "\"precision\": \"%s\", \"layout\": \"%s\"},\n",
641+ r->library, r->runtime, r->adapter, r->precision,
642+ s->layout == SHTB_LAYOUT_PHI ? "phi-contiguous" : "theta-contiguous");
643+ printf(" \"grid\": {\"lmax\": %d, \"nlat\": %d, \"nphi\": %d, \"nlm\": %ld},\n", s->lmax,
644+ s->nlat, s->nphi, r->nlm);
645+ printf(" \"compiled\": {\"opsPerStep\": %d},\n", r->ops_per_step);
646+ printf(" \"throughput\": {\"batch\": %d, \"msPerStep\": %.17g, \"stepsPerSec\": %.17g, "
647+ "\"encodeMsPerStep\": %.17g},\n",
648+ s->batch, r->ms_per_step, 1000.0 / r->ms_per_step, r->encode_ms_per_step);
649+ if (r->have_latency)
650+ printf(" \"latency\": {\"meanMs\": %.17g, \"medianMs\": %.17g, \"p05Ms\": %.17g, "
651+ "\"p95Ms\": %.17g, \"minMs\": %.17g},\n",
652+ r->latency.mean_ms, r->latency.median_ms, r->latency.p05_ms, r->latency.p95_ms,
653+ r->latency.min_ms);
654+ else
655+ printf(" \"latency\": null,\n");
656+ if (r->have_digest)
657+ printf(" \"digest\": {\"n\": %ld, \"min\": %.17g, \"max\": %.17g, \"mean\": %.17g, "
658+ "\"rms\": %.17g, \"fourier\": \"%s\", \"adapter\": \"%s\"},\n",
659+ r->digest.n, r->digest.min, r->digest.max, r->digest.mean, r->digest.rms, r->fourier,
660+ r->adapter);
661+ else
662+ printf(" \"digest\": null,\n");
663+ if (r->have_input_digest)
664+ printf(" \"input\": {\"n\": %ld, \"min\": %.17g, \"max\": %.17g, \"mean\": %.17g, "
665+ "\"rms\": %.17g},\n",
666+ r->input_digest.n, r->input_digest.min, r->input_digest.max, r->input_digest.mean,
667+ r->input_digest.rms);
668+ else
669+ printf(" \"input\": null,\n");
670+ printf(" \"state\": {\"t\": %.17g, \"steps\": %d, \"min\": %.17g, \"max\": %.17g, "
671+ "\"contrast\": %.17g, \"finite\": %s}\n",
672+ r->model_t, r->steps_run, r->field_min, r->field_max, r->field_max - r->field_min,
673+ r->finite ? "true" : "false");
674+ printf("}\n");
675+}
676+
677+/* The state file scripts/compare-native.mjs diffs, in the same shape
678+ * `npm run bench -- --dump-state` writes. */
679+static inline int shtb_dump_state(const char *path, const struct shtb_spec *s,
680+ const struct shtb_report *r, const float *state, long n) {
681+ FILE *f = fopen(path, "w");
682+ if (!f) return -1;
683+ fprintf(f, "{\"spec\":{\"preset\":\"%s\",\"lmax\":%d,\"seed\":%d,\"steps\":%d,\"warmup\":%d,"
684+ "\"params\":{",
685+ s->preset->key, s->lmax, s->seed, s->steps, s->warmup);
686+ for (int i = 0; i < s->model->nparams; i++)
687+ fprintf(f, "%s\"%s\":%.17g", i ? "," : "", s->model->params[i].key,
688+ *shtb_field_c(&s->params, s->model->params[i].off));
689+ fprintf(f, "}},\"mode\":\"%s\",\"backend\":{\"library\":\"%s\",\"adapter\":\"%s\","
690+ "\"precision\":\"%s\"},",
691+ s->mode == SHTB_MODE_SOLVER ? "solver" : "transform", r->library, r->adapter,
692+ r->precision);
693+ fprintf(f, "\"digest\":{\"n\":%ld,\"min\":%.17g,\"max\":%.17g,\"mean\":%.17g,\"rms\":%.17g,"
694+ "\"fourier\":\"%s\",\"adapter\":\"%s\"},",
695+ r->digest.n, r->digest.min, r->digest.max, r->digest.mean, r->digest.rms, r->fourier,
696+ r->adapter);
697+ if (r->have_input_digest)
698+ fprintf(f, "\"input\":{\"n\":%ld,\"min\":%.17g,\"max\":%.17g,\"mean\":%.17g,\"rms\":%.17g},",
699+ r->input_digest.n, r->input_digest.min, r->input_digest.max, r->input_digest.mean,
700+ r->input_digest.rms);
701+ fprintf(f, "\"state\":[");
702+ for (long i = 0; i < n; i++) fprintf(f, "%s%.9g", i ? "," : "", (double)state[i]);
703+ fprintf(f, "]}\n");
704+ fclose(f);
705+ return 0;
706+}
707+
708+#endif /* SHTB_SPEC_H */
package.jsonmodified+3−1View file
@@ -13,7 +13,9 @@
1313 "test:node": "vite-node scripts/test-node.ts",
1414 "test:gpu": "vite build && node scripts/test-gpu.mjs",
1515 "test": "npm run test:node && npm run test:gpu",
16- "bench": "vite-node scripts/bench.ts"
16+ "bench": "vite-node scripts/bench.ts",
17+ "bench:sht": "vite-node scripts/bench-sht.ts",
18+ "bench:native": "node scripts/compare-native.mjs"
1719 },
1820 "dependencies": {
1921 "numbl": "file:../../numbl",
scripts/bench-sht.tsadded+368−0View file
@@ -0,0 +1,368 @@
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+ */
23+import { ShtPlan, requestShtDevice, describeAdapter, type ShtBinding } from '../src/sht/sht.ts';
24+import { lmIndex } from '../src/sht/layout.ts';
25+import { makeRand } from '../src/mgpu/noise.ts';
26+import { digestOf, formatDigest } from '../src/mgpu/digest.ts';
27+import {
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';
37+import { presets } from '../src/mgpu/registry.ts';
38+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
39+import { writeFileSync } from 'node:fs';
40+
41+const USAGE = `usage: npm run bench:sht -- [options]
42+
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
58+
59+The 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)`;
62+
63+function fail(msg: string, code = 1): never {
64+ console.error(`bench:sht: ${msg}`);
65+ process.exit(code);
66+}
67+
68+// ---------------------------------------------------------------- arguments
69+const argv = process.argv.slice(2);
70+if (argv.includes('--help') || argv.includes('-h')) {
71+ console.log(USAGE);
72+ process.exit(0);
73+}
74+const wantJson = argv.includes('--json');
75+let batch = 16;
76+let fourier: 'auto' | 'fft' | 'dft' = 'auto';
77+let dumpState: string | null = null;
78+let wantDigest = false;
79+const rest: string[] = [];
80+for (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+}
111+if (!Number.isInteger(batch) || batch < 1) fail('--batch must be an integer >= 1', 2);
112+
113+let spec: RunSpec;
114+try {
115+ spec = parseArgs(rest);
116+} catch (e) {
117+ fail(`${errMsg(e)}\n\n${USAGE}`, 2);
118+}
119+const cfg = configForSpec(spec);
120+
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+ */
130+function 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+}
143+
144+// ---------------------------------------------------------------------- run
145+let device: GPUDevice | null = null;
146+let plan: ShtPlan | null = null;
147+
148+try {
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;
157+
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+ ];
184+
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();
214+
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+ }
228+
229+ seed();
230+ submit(spec.warmup);
231+ await done();
232+
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;
248+
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+ };
269+
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+ }
284+
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+ }
294+
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+ }
344+
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+ }
359+
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+}
scripts/bench.tsmodified+1−1View file
@@ -252,7 +252,7 @@ try {
252252 command: formatCommand(spec),
253253 spec,
254254 model: model.key,
255- backend: { adapter, runtime },
255+ backend: { adapter, runtime, precision: 'fp32' },
256256 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm: session.sht.nlm },
257257 compiled: { opsPerStep: plan.step.length, kernels },
258258 digest,
scripts/compare-native.mjsadded+455−0View file
@@ -0,0 +1,455 @@
1+/**
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+ */
33+import { spawnSync } from 'node:child_process';
34+import { readFileSync, existsSync, unlinkSync } from 'node:fs';
35+import { join } from 'node:path';
36+import { tmpdir } from 'node:os';
37+
38+const ROOT = new URL('..', import.meta.url).pathname;
39+const NATIVE = join(ROOT, 'bench', 'shtns');
40+
41+// ---------------------------------------------------------------- arguments
42+const argv = process.argv.slice(2);
43+const has = (name) => argv.includes(`--${name}`);
44+const 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+};
50+if (has('help') || argv.includes('-h')) {
51+ console.log(`usage: node scripts/compare-native.mjs [options]
52+
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+}
76+const mode = flag('mode', 'transform');
77+if (mode !== 'transform' && mode !== 'solver') {
78+ console.error(`compare-native: --mode must be 'transform' or 'solver'`);
79+ process.exit(2);
80+}
81+const lmax = flag('lmax', '63');
82+const steps = flag('steps', '1000');
83+const warmup = flag('warmup', '50');
84+const preset = flag('preset', 'schnak-spots');
85+const batch = flag('batch', '16');
86+const layout = flag('layout', 'theta');
87+const threads = flag('threads', '1');
88+const wantCheck = has('check');
89+const checkSteps = flag('check-steps', '20');
90+const wantJson = has('json');
91+const tolerance = Number(flag('tolerance', '2e-3'));
92+const wantCpu = !has('no-cpu');
93+const progress = !wantJson && process.stderr.isTTY;
94+
95+const tmp = (tag) => join(tmpdir(), `turing-sphere-native-${tag}-${process.pid}.json`);
96+const cleanup = [];
97+
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. */
101+function 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+}
128+
129+const common = ['--lmax', lmax, '--steps', steps, '--warmup', warmup, '--preset', preset];
130+const nativeCommon = [...common, '--batch', batch, '--layout', layout, '--json'];
131+
132+const jobs = [];
133+if (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+}
146+const gpuBin = join(NATIVE, 'shtbench_gpu');
147+const cpuBin = join(NATIVE, 'shtbench');
148+const nativeMode = ['--mode', mode];
149+if (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+}
159+if (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+}
173+
174+if (!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+}
182+
183+const results = [];
184+for (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+}
193+
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. */
197+const states = new Map();
198+if (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+}
207+
208+const good = results.filter((r) => r.ok);
209+if (!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+}
214+
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.
218+const gridOf = (r) => r.json.grid;
219+const ref = good[0];
220+const gridProblems = [];
221+for (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+}
237+if (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+}
246+
247+// -------------------------------------------------------------------- report
248+const rate = (r) => r.json.throughput.msPerStep;
249+const base = rate(good.find((r) => r.label === 'webgpu') ?? good[0]);
250+
251+if (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+ }
314+
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+}
355+
356+// --------------------------------------------------------------------- check
357+let checkFailed = false;
358+if (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+}
409+
410+for (const p of cleanup) {
411+ try {
412+ unlinkSync(p);
413+ } catch {
414+ /* best effort */
415+ }
416+}
417+process.exit(checkFailed ? 1 : 0);
418+
419+// ------------------------------------------------------------------- helpers
420+function 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+}
431+
432+function 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+}
437+
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. */
441+function 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+}
src/mgpu/noise.tsmodified+14−3View file
@@ -5,16 +5,27 @@
55 * seed and the same field can be handed to any model.
66 */
77
8-/** Seeded normal deviates: mulberry32 + Box-Muller. */
9-export function makeRandn(seed: number): () => number {
8+/**
9+ * Seeded uniform deviates in [0, 1): mulberry32.
10+ *
11+ * Integer arithmetic and one division by 2^32, so any faithful port of it
12+ * produces bit-identical values — which is what lets the native benchmark under
13+ * bench/shtns/ seed the same run.
14+ */
15+export function makeRand(seed: number): () => number {
1016 let s = seed >>> 0;
11- const rand = (): number => {
17+ return (): number => {
1218 s = (s + 0x6d2b79f5) >>> 0;
1319 let t = s;
1420 t = Math.imul(t ^ (t >>> 15), t | 1);
1521 t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
1622 return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
1723 };
24+}
25+
26+/** Seeded normal deviates: mulberry32 + Box-Muller. */
27+export function makeRandn(seed: number): () => number {
28+ const rand = makeRand(seed);
1829 let spare: number | null = null;
1930 return () => {
2031 if (spare !== null) {