Tidy up after the SHTNS comparison - concept-collection/turing-sphere

concept-collection / turing-sphere

Tidy up after the SHTNS comparison

Nothing behavioural. Collected while re-reading what the last several commits left behind: - bench-sht.ts said what run it was three times over — once for the success report, once for the early-exit report, once for the state file. One `identity` object now, spread into each. - record the Blackwell leg_synth bug in the README, next to the canvas resizing note, since it is the same kind of thing: a failure that took real work to localize and whose fix looks arbitrary without the story. - the README's example table showed "shtns cuda not available", which was honest when there was no Nvidia card to hand and is now just stale. Replaced with the RTX PRO 6000 numbers, and with what makes them readable: SHTNS' row is 86% CPU launch overhead, so 4x is a lower bound on the gap in GPU work rather than a measurement of it, and one CPU core keeping pace with a 188-SM GPU at lmax 63 is about occupancy. The "run a sweep" advice was there twice; it is there once now. - list diagnose-leg.ts among the commands, and have diagnose-sht.ts print the diagnose-leg invocation to run next. - shtns.mk records the GPU Fourier stage and nothing read it; `make` now reports it, and warns that cuFFT rules out --layout phi. - drop an unused GridSizes import in src/mgpu/plan.ts, unrelated but it is the only thing --noUnusedLocals finds in the tree.

Jeremy Magland <jmagland@flatironinstitute.org> committed 20 days ago commit 4166b4848eda parent 3bc607e Browse files

5 changed files+85−25

README.mdmodified+59−9View file

@@ -340,19 +340,36 @@ node scripts/compare-native.mjs --check # then, from the repo root

340340 you the CPU comparison anywhere and the GPU one on a machine with an Nvidia card.

341341 `compare-native.mjs` runs every implementation present, back to back in one

342342 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:

343+than one, and prints them in one table. On an RTX PRO 6000 Blackwell, at the

344+app's default lmax:

345345

346346 ```

347347 grid lmax 63 · 128×256 · nlm 2,080 (one synthesis + one analysis per round trip)

348348

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

349+ webgpu 0.084 ms/round trip 11848/s (baseline) fp32

350+ NVIDIA (blackwell), via Dawn · CPU-side launching 0.012 ms/step

351+ shtns cuda 0.021 ms/round trip 48009/s 0.25x webgpu fp32

352+ NVIDIA RTX PRO 6000 (sm_120, 188 SMs) · CPU-side launching 0.018 ms/step

353+ shtns cpu 0.072 ms/round trip 13982/s 0.85x webgpu fp64

353354 CPU, 1 thread

354355 ```

355356

357+Read that carefully rather than as "4x". The two GPU rows are limited by different

358+things: the WGSL row spends 14% of its time on the CPU and is genuinely GPU-bound,

359+while SHTNS spends **86%** — 0.018 ms of 0.021 — queueing its six-or-so kernels, so

360+its number is close to what it costs to *submit* a round trip on that host and its

361+actual GPU time is below that and unresolved. The 4x is a lower bound on the gap in

362+GPU work, not a measurement of it. `compare-native.mjs` flags any row above 50%

363+for this reason.

364+

365+The other number worth noticing is the third row: one CPU core in fp64 is about

366+level with the WGSL transforms on a 188-SM datacentre GPU. At lmax 63 there are

367+2,080 coefficients on a 128×256 grid — far too little work to occupy that card, so

368+this says more about occupancy than about the shaders. Sweep lmax before drawing

369+conclusions, and stop at 511: above that `16*nphi` exceeds the workgroup-storage

370+limit, the FFT stage falls back to the O(nphi·mmax) DFT, and the comparison stops

371+being about the FFT.

372+

356373 Two things are measured, because they answer different questions:

357374

358375 - **transforms** (`npm run bench:sht` here, `--mode transform` there) — one

@@ -383,9 +400,8 @@ recurrence in fp64 even in fp32 mode for `lmax <= 128` (WebGPU has no fp64 at

383400 all), the Fourier stages are cuFFT/VkFFT/FFTW against a WGSL FFT, and SHTNS'

384401 polar optimization is off by default here because we have none.

385402

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

403+How much the grid size matters is easiest to see on a weak GPU, where there is no

404+launch-overhead floor to hide behind. On an Intel Xe iGPU against one core of the

389405 same laptop, one round trip costs:

390406

391407 | lmax | grid | WGSL (fp32) | SHTNS, 1 CPU core (fp64) |

@@ -454,6 +470,10 @@ Other commands:

454470 say *which* stage is wrong. It reads the intermediate `fm` back out and scores

455471 the Legendre and Fourier stages of each direction separately against the f64

456472 reference, then breaks the error down by order `m` and by latitude.

473+- `npx vite-node scripts/diagnose-leg.ts [--m 0]` — the follow-up to that: read the

474+ Legendre recurrence out of the production shader term by term, by synthesizing a

475+ spectrum that is 1 at a single coefficient, and compare each `ỹ_l^m` with the f64

476+ reference. The first term that disagrees names the culprit.

457477 - `npx vite-node scripts/longrun-node.ts [lmax]` — run to t = 100 and confirm the

458478 pattern saturates into O(1)-contrast spots rather than decaying or diverging.

459479 - `node scripts/soak.mjs [steps] [lmax]` — drive the demo for many steps,

@@ -484,6 +504,36 @@ reflowed the panel, fired the `ResizeObserver`, and called

484504 same bug caused visible flicker. Fixed by giving the colorbar column a fixed

485505 width and making `SphereScene.resize()` return early on no-op resizes.

486506

507+### A note on the Legendre recurrence on Blackwell

508+

509+The first run on an Nvidia GPU — an RTX PRO 6000, driver 590.48, reached through

510+Dawn's Vulkan backend — failed 11 of the tests. `synth` was off by 5.5e+3 while

511+`analys` was accurate to 7.3e-7, and the solver produced NaN within 40 steps.

512+

513+The two diagnostic scripts above were written for it and localized it in two

514+steps: `leg_synth` was the only wrong shader, and within it the recurrence was

515+right at `l = m` and `l = m+1` and then returned *exactly zero* at `l = m+2`, at

516+every latitude. That is not a precision failure. It is

517+

518+```wgsl

519+let c0 = ab[base + (l + 2u - m)];

520+y0 = c0.x * ct * y1 + c0.y * y0; // c0 reads as (0, 0) on the first iteration

521+```

522+

523+with the `ab` read two lines later working fine. The buffer was not at fault:

524+`leg_analys` reads the same array correctly on the same device, and `m = 62, 63`

525+— the only orders whose loop breaks before that line — were the only correct

526+ones. Nothing about that WGSL is invalid, so it was a miscompiled load.

527+

528+Fixed by giving the advance the shape `leg_analys` already used, which that

529+driver compiles correctly: both coefficients fetched unconditionally, and the new

530+`y0` carried in a temporary rather than assigned and then read back by the `y1`

531+update. Two shaders doing the same recurrence should have agreed on form anyway.

532+

533+Worth knowing for what it says about the transforms in general: nothing had

534+exercised them on Nvidia hardware before, and the existing test caught it

535+immediately — it just could not say where. That is what the diagnostics are for.

536+

487537 ## Development

488538

489539 ```

bench/shtns/Makefilemodified+3−1View file

@@ -34,7 +34,9 @@ endif

3434 all: $(TARGETS)

3535 @echo

3636 @echo "built: $(TARGETS)"

37-ifneq ($(SHTNS_HAS_GPU),1)

37+ifeq ($(SHTNS_HAS_GPU),1)

38+ @echo "SHTNS GPU Fourier stage: $(SHTNS_GPU_FFT)$(if $(filter cufft,$(SHTNS_GPU_FFT)), — theta-contiguous only; --layout phi needs VkFFT,)"

39+else

3840 @echo "note: SHTNS was built without GPU support, so there is no shtbench_gpu."

3941 @echo " Re-run ./bootstrap.sh on a machine with nvcc to get it."

4042 endif

scripts/bench-sht.tsmodified+19−13View file

@@ -219,6 +219,22 @@ try {

219219 };

220220 const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();

221221

222+ // What every report of this run says about itself, whether it succeeded, failed

223+ // early, or is being written to a state file for compare-native.mjs to diff.

224+ const identity = {

225+ mode: 'transform',

226+ spec: {

227+ preset: spec.preset,

228+ lmax: spec.lmax,

229+ seed: spec.seed,

230+ steps: spec.steps,

231+ warmup: spec.warmup,

232+ },

233+ backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },

234+ grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },

235+ fourier: plan.fourierMode,

236+ };

237+

222238 if (!wantJson) {

223239 console.log('turing-sphere bench:sht — transforms only, no solver, no rendering\n');

224240 console.log(

@@ -269,11 +285,7 @@ try {

269285 console.log(

270286 JSON.stringify(

271287 {

272- mode: 'transform',

273- spec: { preset: spec.preset, lmax: spec.lmax, seed: spec.seed, steps: spec.steps, warmup: spec.warmup },

274- backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },

275- grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },

276- fourier: plan.fourierMode,

288+ ...identity,

277289 firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },

278290 throughput: null,

279291 latency: null,

@@ -373,11 +385,7 @@ try {

373385 console.log(

374386 JSON.stringify(

375387 {

376- mode: 'transform',

377- spec: { preset: spec.preset, lmax: spec.lmax, seed: spec.seed, steps: spec.steps, warmup: spec.warmup },

378- backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },

379- grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },

380- fourier: plan.fourierMode,

388+ ...identity,

381389 firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },

382390 throughput: {

383391 batch,

@@ -424,9 +432,7 @@ try {

424432 writeFileSync(

425433 dumpState,

426434 JSON.stringify({

427- mode: 'transform',

428- spec: { preset: spec.preset, lmax: spec.lmax, seed: spec.seed, steps: spec.steps, warmup: spec.warmup },

429- backend: { library: 'shtns-webgpu (src/sht)', adapter, precision: 'fp32' },

435+ ...identity,

430436 digest,

431437 input: inputDigest,

432438 state: [...state],

scripts/diagnose-sht.tsmodified+3−1View file

@@ -283,7 +283,9 @@ try {

283283 ` If only high m are wrong, or only latitudes near the poles where\n` +

284284 ` sin(theta) is small, the rescaled seed (sinpow_rescaled in\n` +

285285 ` src/sht/wgsl/common.ts) is the place to look. If every m is wrong by\n` +

286- ` a similar amount, it is indexing or the dispatch, not the recurrence.`,

286+ ` a similar amount, it is indexing or the dispatch, not the recurrence.\n` +

287+ ` Either way, follow it term by term from here:\n` +

288+ ` npx vite-node scripts/diagnose-leg.ts --lmax ${lmax} --m 0`,

287289 );

288290 }

289291 console.log();

src/mgpu/plan.tsmodified+1−1View file

@@ -13,7 +13,7 @@ import type { Assign, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/i

1313 import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';

1414 import { ShtPlan, type ShtBinding } from '../sht/sht.ts';

1515 import type { CompiledFunction } from './compile.ts';

16-import { EXTERNAL_OPS, type GridSizes } from './externals.ts';

16+import { EXTERNAL_OPS } from './externals.ts';

1717 import {

1818 buildKernel,

1919 UnsupportedOnGpu,