b689087Benchmark the WGSL transforms against upstream SHTNSJeremy Magland 1# The WGSL transforms against upstream SHTNS
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
6check it against its own f64 CPU twin, which says it is self-consistent — not
7how it compares with the library it is modeled on.
9This directory answers that. It builds upstream SHTNS and runs the same spec
10through it, so the numbers line up against `npm run bench`:
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 |
18`shtbench_gpu` is the like-for-like comparison: same GPU, same precision, same
19grid, same spectral conventions — the only difference is who computes the
20transform. `shtbench` is fp64 because SHTNS' single precision exists only on the
21GPU, so it serves as the accuracy reference (how far has fp32 drifted?) and as
22the "what does a well-optimized CPU do" reference.
24## Build
26```
27./bootstrap.sh # clone SHTNS at a pinned commit, configure, build
28make # shtbench, plus shtbench_gpu if SHTNS got CUDA support
29```
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
34the `Makefile` includes — so a host that needed MKL or a different FFTW still
35links without editing anything here. Both the checkout and `shtns.mk` are
36gitignored; `make distclean` throws them away.
38Needs: a C++ compiler, FFTW3 headers (`libfftw3-dev`), and for the GPU half a
39CUDA toolkit. SHTNS' `configure` wants `CUDA_PATH` set; `bootstrap.sh` derives it
40from `nvcc`'s location if it is not.
42## Run
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```
50Both 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`.
54Two things are measured, and they answer different questions:
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.
64Both report throughput (a batch launched together, waited for once — what
65`--batch` controls, matching `npm run bench -- --batch`) and a per-step
66distribution from one synchronization per step.
68## Compare
70```
71node scripts/compare-native.mjs # transforms, lmax 63
72node scripts/compare-native.mjs --mode solver
73node scripts/compare-native.mjs --check # and diff the final state
74```
76from the repo root (or `npm run bench:native --`). It runs every implementation
77present on the machine, back to back in one invocation so a second process
78competing for the GPU affects both sides rather than one, and prints them in one
79table. Missing implementations are reported and skipped, so this is still useful
80on a machine with no CUDA.
82`--check` adds a short second pass that diffs the final spectral state across
83implementations. That is what makes the timing mean anything: two numbers are
84only comparable if they are the cost of the same computation.
86## What is and is not the same on the two sides
88The comparison is exact where it can be:
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.
104And explicit where it cannot be:
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.
140## Reading the result
142Small grids flatter the CPU: at `lmax 31` there are 528 coefficients, and a GPU
143spends most of a transform on launch latency rather than arithmetic. The gap
144closes with `lmax`, so run a sweep before concluding anything:
146```
147for l in 31 63 127 255; do node scripts/compare-native.mjs --lmax $l --steps 500; done
148```
150If the WGSL side lands on a software adapter, `compare-native.mjs` says so and
151stops you — the ratio then compares a CPU emulation against a real GPU and means
152nothing.
154## Editing `shtbench_gpu.cu` without a GPU
156`nvcc` is needed to build it, but not to typecheck it. Rewriting the launch
157syntax as calls, against a handful of stub declarations, gets g++ to check
158everything else:
160```
161sed -e 's/<<</\/*/g; s/>>>/*\//g' shtbench_gpu.cu > /tmp/check.cpp
162g++ -fsyntax-only -std=c++14 -I. -Ishtns -I/path/to/cuda-stubs /tmp/check.cpp
163```
165where 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.
169## License
171CECILL-2.1, as the rest of this repo — the same license SHTNS itself is under.