/ concept-collection / turing-sphere
Sign in
concept-collection / turing-sphere
Add scripts/diagnose-sht.ts: which stage of the transform is wrong?
npm run test:node on an RTX PRO 6000 reports transforms: WGSL fp32 vs f64 CPU reference synth 5.53e+3, analys 7.27e-7 so analysis is exactly as accurate as fp32 should be and synthesis is producing garbage, on a GPU these transforms have not run on before. Both directions are two stages and they share code, so that pair of numbers is as far as the existing test goes. This goes further by reading the intermediate `fm` back out and scoring each stage on its own: fm wrong -> the Legendre stage fm right but spat wrong -> the Fourier stage both fine in DFT, wrong in FFT -> the WGSL FFT specifically The Fourier stage is judged against a host Fourier transform of the fm the GPU actually produced, not against the reference's, so a Legendre error does not get counted twice. It runs both Fourier stages in one invocation, and when the Legendre stage is the culprit it breaks the error down by order m and by latitude — "only high m" or "only near the poles" implicates the rescaled seed in wgsl/common.ts, while a uniform error implicates indexing or the dispatch. fmBuf gains COPY_SRC so the stage boundary is observable at all. It is a usage flag on a buffer that is never mapped in the hot path, so it costs nothing; test:node is unchanged by it.
Jeremy Magland <jmagland@flatironinstitute.org> committed commit 30ed90fe30a3 parent 721b8df Browse files
3 changed files+317−2
README.mdmodified+4−0View file
@@ -450,6 +450,10 @@ Other commands:
450450 [Desktop vs browser](#desktop-vs-browser)).
451451 - `npm run bench:sht -- --help` — the transforms alone, with no solver around
452452 them, for comparing against upstream SHTNS.
453+- `npx vite-node scripts/diagnose-sht.ts` — when the transform tests fail on a GPU,
454+ say *which* stage is wrong. It reads the intermediate `fm` back out and scores
455+ the Legendre and Fourier stages of each direction separately against the f64
456+ reference, then breaks the error down by order `m` and by latitude.
453457 - `npx vite-node scripts/longrun-node.ts [lmax]` — run to t = 100 and confirm the
454458 pattern saturates into O(1)-contrast spots rather than decaying or diverging.
455459 - `node scripts/soak.mjs [steps] [lmax]` — drive the demo for many steps,
scripts/diagnose-sht.tsadded+309−0View file
@@ -0,0 +1,309 @@
1+/**
2+ * Which half of a transform is wrong on this device?
3+ *
4+ * npx vite-node scripts/diagnose-sht.ts [--lmax 63] [--seed 12345]
5+ *
6+ * A transform is two stages, and both directions share code, so a single
7+ * pass/fail says very little:
8+ *
9+ * synthesis: qlm --[leg_synth]--> fm --[fft_synth | dft_synth]--> spat
10+ * analysis: spat --[fft_analys | dft_analys]--> fm --[leg_analys]--> qlm
11+ *
12+ * This reads the intermediate `fm` back out and compares each stage against
13+ * src/sht/reference.ts (f64, direct summation) on its own:
14+ *
15+ * - fm wrong -> the Legendre stage
16+ * - fm right but spat wrong -> the Fourier stage
17+ * - both right in DFT, wrong in FFT -> the WGSL FFT specifically
18+ *
19+ * and breaks the error down by m and by latitude, because "only high m" or "only
20+ * near the poles" points straight at the rescaled recurrence, while "every m
21+ * equally" points at indexing.
22+ *
23+ * Written for a report of `npm run test:node` failing on a GPU the transforms
24+ * have not run on before. Nothing here is a benchmark.
25+ */
26+import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
27+import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
28+import { gridForLmax, isPowerOfTwo, type ShtConfig } from '../src/sht/layout.ts';
29+import { fftThreads } from '../src/sht/wgsl/fourier.ts';
30+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
31+
32+const argv = process.argv.slice(2);
33+if (argv.includes('--help') || argv.includes('-h')) {
34+ console.log(`usage: npx vite-node scripts/diagnose-sht.ts [options]
35+
36+ --lmax <n> spherical harmonic truncation (default 63, the app's)
37+ --seed <n> seed of the test spectrum (default 12345)
38+ --fourier <m> only test this stage: fft | dft (default: both)
39+ --help
40+
41+Compares each stage of each direction against the f64 CPU reference and says
42+which one is wrong. See the header of this file.`);
43+ process.exit(0);
44+}
45+const flag = (name: string, dflt: string): string => {
46+ const i = argv.indexOf(`--${name}`);
47+ if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
48+ const eq = argv.find((a) => a.startsWith(`--${name}=`));
49+ return eq ? eq.slice(name.length + 3) : dflt;
50+};
51+const lmax = Number(flag('lmax', '63'));
52+const seed = Number(flag('seed', '12345'));
53+const only = flag('fourier', '');
54+
55+/** Relative L2 of a against b, both flat. */
56+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
57+ let num = 0;
58+ let den = 0;
59+ for (let i = 0; i < b.length; i++) {
60+ const d = (a[i] ?? NaN) - b[i];
61+ num += d * d;
62+ den += b[i] * b[i];
63+ }
64+ return Math.sqrt(num / Math.max(den, 1e-300));
65+}
66+
67+function anyNonFinite(a: ArrayLike<number>): boolean {
68+ for (let i = 0; i < a.length; i++) if (!Number.isFinite(a[i])) return true;
69+ return false;
70+}
71+
72+/** The Fourier half of a synthesis, on the host, from whatever fm it is given.
73+ * Mirrors ShtReference.synth's inner loop — so feeding it the GPU's own fm says
74+ * what the Fourier stage should have produced from the input it actually had. */
75+function fourierSynth(cfg: ShtConfig, fm: ArrayLike<number>): Float64Array {
76+ const { mmax, nlat, nphi } = cfg;
77+ const spat = new Float64Array(nlat * nphi);
78+ for (let i = 0; i < nlat; i++) {
79+ for (let j = 0; j < nphi; j++) {
80+ const phi = (2 * Math.PI * j) / nphi;
81+ let v = fm[2 * i];
82+ for (let m = 1; m <= mmax; m++) {
83+ const o = 2 * (m * nlat + i);
84+ v += 2 * (fm[o] * Math.cos(m * phi) - fm[o + 1] * Math.sin(m * phi));
85+ }
86+ spat[i * nphi + j] = v;
87+ }
88+ }
89+ return spat;
90+}
91+
92+/** Worst offender along one axis of the [m][ilat] complex fm array. */
93+function fmBreakdown(
94+ cfg: ShtConfig,
95+ got: ArrayLike<number>,
96+ want: ArrayLike<number>,
97+): { byM: { m: number; rel: number }[]; worstLat: { ilat: number; rel: number } } {
98+ const { mmax, nlat } = cfg;
99+ const byM: { m: number; rel: number }[] = [];
100+ const latNum = new Float64Array(nlat);
101+ const latDen = new Float64Array(nlat);
102+ for (let m = 0; m <= mmax; m++) {
103+ let num = 0;
104+ let den = 0;
105+ for (let i = 0; i < nlat; i++) {
106+ for (let c = 0; c < 2; c++) {
107+ const k = 2 * (m * nlat + i) + c;
108+ const d = (got[k] ?? NaN) - want[k];
109+ num += d * d;
110+ den += want[k] * want[k];
111+ latNum[i] += d * d;
112+ latDen[i] += want[k] * want[k];
113+ }
114+ }
115+ byM.push({ m, rel: Math.sqrt(num / Math.max(den, 1e-300)) });
116+ }
117+ let worstLat = { ilat: 0, rel: 0 };
118+ for (let i = 0; i < nlat; i++) {
119+ const rel = Math.sqrt(latNum[i] / Math.max(latDen[i], 1e-300));
120+ if (rel > worstLat.rel) worstLat = { ilat: i, rel };
121+ }
122+ return { byM, worstLat };
123+}
124+
125+const OK = 1e-4; // fp32 through these transforms lands near 1e-6; 1e-4 is generous
126+const verdict = (rel: number): string => (rel < OK ? 'ok ' : 'WRONG');
127+
128+let device: GPUDevice | null = null;
129+try {
130+ const runtime = await installWebGpu();
131+ device = await requestShtDevice().catch((e: unknown) => {
132+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
133+ });
134+ const adapter = await describeAdapter(device);
135+
136+ const { nlat, nphi } = gridForLmax(lmax, 3);
137+ const cfg: ShtConfig = { lmax, mmax: lmax, nlat, nphi };
138+ const ref = new ShtReference(cfg);
139+ const qlm = randomSpectrum(cfg, seed);
140+
141+ console.log('turing-sphere — which stage of the transform is wrong?\n');
142+ console.log(` device ${adapter || '(unknown)'}\n ${runtime}`);
143+ console.log(` grid lmax ${cfg.lmax} · ${nlat}×${nphi} · nlm ${ref.nlm}`);
144+ console.log(
145+ ` limits maxComputeWorkgroupStorageSize ${device.limits.maxComputeWorkgroupStorageSize}` +
146+ `, maxComputeInvocationsPerWorkgroup ${device.limits.maxComputeInvocationsPerWorkgroup}`,
147+ );
148+ console.log(
149+ ` the FFT stage needs a power-of-two nphi (${isPowerOfTwo(nphi)}), ` +
150+ `16*nphi = ${16 * nphi} bytes of\n workgroup storage and ` +
151+ `${fftThreads(nphi)} invocations per workgroup\n`,
152+ );
153+
154+ // reference values, computed once
155+ const fmRef = ref.legendreSynth(qlm);
156+ const spatRef = ref.synth(qlm);
157+ const qlmRef = ref.analys(spatRef);
158+
159+ const modes: ('fft' | 'dft')[] =
160+ only === 'fft' || only === 'dft' ? [only] : ['fft', 'dft'];
161+ const summary: string[] = [];
162+
163+ for (const mode of modes) {
164+ let plan: ShtPlan | null = null;
165+ try {
166+ plan = await ShtPlan.create(device, cfg, { fourier: mode });
167+ } catch (e) {
168+ console.log(`${mode.toUpperCase()} stage: unavailable — ${errMsg(e)}\n`);
169+ continue;
170+ }
171+
172+ const stageFm = device.createBuffer({
173+ size: 8 * (cfg.mmax + 1) * nlat,
174+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
175+ });
176+ const stageSpat = device.createBuffer({
177+ size: 4 * nlat * nphi,
178+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
179+ });
180+ const read = async (buf: GPUBuffer): Promise<Float32Array> => {
181+ await buf.mapAsync(GPUMapMode.READ);
182+ const out = new Float32Array(buf.getMappedRange().slice(0));
183+ buf.unmap();
184+ return out;
185+ };
186+
187+ // --- synthesis, stopping to look at fm on the way through ---
188+ device.queue.writeBuffer(plan.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
189+ const enc = device.createCommandEncoder();
190+ plan.encodeSynth(enc);
191+ enc.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
192+ enc.copyBufferToBuffer(plan.spatBuf, 0, stageSpat, 0, 4 * nlat * nphi);
193+ device.queue.submit([enc.finish()]);
194+ const fmGpu = await read(stageFm);
195+ const spatGpu = await read(stageSpat);
196+
197+ const legSynthRel = relL2(fmGpu, fmRef);
198+ const synthRel = relL2(spatGpu, spatRef);
199+ // The Fourier stage judged on its own input, not on the reference's: if the
200+ // Legendre stage is already wrong, comparing spat with spatRef only repeats
201+ // that. This asks whether the Fourier stage did the right thing with the fm
202+ // it was actually handed.
203+ const fourierRel = relL2(spatGpu, fourierSynth(cfg, fmGpu));
204+
205+ // --- analysis, for contrast: same two stages, opposite order ---
206+ const spatIn = Float32Array.from(spatRef);
207+ device.queue.writeBuffer(plan.spatBuf, 0, spatIn as Float32Array<ArrayBuffer>);
208+ const enc2 = device.createCommandEncoder();
209+ plan.encodeAnalys(enc2);
210+ enc2.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
211+ device.queue.submit([enc2.finish()]);
212+ const fmAnalysGpu = await read(stageFm);
213+ const qlmGpu = await plan.analys(spatIn);
214+ // forward Fourier of the reference field, in the reference's own normalization
215+ const gmRef = new Float64Array(2 * (cfg.mmax + 1) * nlat);
216+ for (let i = 0; i < nlat; i++) {
217+ for (let m = 0; m <= cfg.mmax; m++) {
218+ let re = 0;
219+ let im = 0;
220+ for (let j = 0; j < nphi; j++) {
221+ const phi = (2 * Math.PI * j) / nphi;
222+ re += spatRef[i * nphi + j] * Math.cos(m * phi);
223+ im -= spatRef[i * nphi + j] * Math.sin(m * phi);
224+ }
225+ gmRef[2 * (m * nlat + i)] = re;
226+ gmRef[2 * (m * nlat + i) + 1] = im;
227+ }
228+ }
229+ const analysFourierRel = relL2(fmAnalysGpu, gmRef);
230+ const analysRel = relL2(qlmGpu, qlmRef);
231+
232+ console.log(`${mode.toUpperCase()} stage — plan chose ${plan.fourierMode.toUpperCase()}\n`);
233+ console.log(` synthesis qlm -> fm -> spat`);
234+ console.log(
235+ ` ${verdict(legSynthRel)} leg_synth fm vs f64 reference ` +
236+ `${legSynthRel.toExponential(2)}${anyNonFinite(fmGpu) ? ' (has NaN/Inf)' : ''}`,
237+ );
238+ console.log(
239+ ` ${verdict(fourierRel)} ${mode}_synth spat vs host Fourier of that fm ` +
240+ `${fourierRel.toExponential(2)}${anyNonFinite(spatGpu) ? ' (has NaN/Inf)' : ''}`,
241+ );
242+ console.log(
243+ ` ${verdict(synthRel)} end to end spat vs f64 reference ` +
244+ `${synthRel.toExponential(2)}`,
245+ );
246+ console.log(`\n analysis spat -> fm -> qlm`);
247+ console.log(
248+ ` ${verdict(analysFourierRel)} ${mode}_analys fm vs f64 reference ` +
249+ `${analysFourierRel.toExponential(2)}`,
250+ );
251+ console.log(
252+ ` ${verdict(analysRel)} end to end qlm vs f64 reference ` +
253+ `${analysRel.toExponential(2)}`,
254+ );
255+
256+ if (legSynthRel >= OK) {
257+ const { byM, worstLat } = fmBreakdown(cfg, fmGpu, fmRef);
258+ const bad = byM.filter((e) => e.rel >= OK);
259+ const good = byM.filter((e) => e.rel < OK);
260+ console.log(`\n leg_synth is wrong. Where:`);
261+ console.log(
262+ ` ${bad.length} of ${byM.length} orders m are wrong` +
263+ (good.length
264+ ? `; the ones that are right are m = ${good.slice(0, 12).map((e) => e.m).join(', ')}` +
265+ (good.length > 12 ? ', ...' : '')
266+ : ' (all of them)'),
267+ );
268+ if (bad.length) {
269+ const first = bad[0];
270+ const worst = bad.reduce((a, b) => (b.rel > a.rel ? b : a));
271+ console.log(
272+ ` lowest wrong m = ${first.m} (${first.rel.toExponential(2)}), ` +
273+ `worst m = ${worst.m} (${worst.rel.toExponential(2)})`,
274+ );
275+ }
276+ const theta = (Math.acos(ref.ct[worstLat.ilat]) * 180) / Math.PI;
277+ console.log(
278+ ` worst latitude ilat = ${worstLat.ilat} of ${nlat} ` +
279+ `(theta = ${theta.toFixed(1)}°, sin(theta) = ` +
280+ `${ref.st[worstLat.ilat].toExponential(2)}), rel ${worstLat.rel.toExponential(2)}`,
281+ );
282+ console.log(
283+ ` If only high m are wrong, or only latitudes near the poles where\n` +
284+ ` sin(theta) is small, the rescaled seed (sinpow_rescaled in\n` +
285+ ` 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.`,
287+ );
288+ }
289+ console.log();
290+
291+ summary.push(
292+ `${mode}: leg_synth ${verdict(legSynthRel).trim()}, ${mode}_synth ` +
293+ `${verdict(fourierRel).trim()}, ${mode}_analys ${verdict(analysFourierRel).trim()}, ` +
294+ `leg_analys ${verdict(analysRel).trim()}`,
295+ );
296+
297+ stageFm.destroy();
298+ stageSpat.destroy();
299+ plan.destroy();
300+ }
301+
302+ console.log('summary');
303+ for (const s of summary) console.log(` ${s}`);
304+ device.destroy();
305+} catch (e) {
306+ device?.destroy();
307+ console.error(`diagnose-sht: ${errMsg(e)}`);
308+ process.exit(1);
309+}
src/sht/sht.tsmodified+4−2View file
@@ -80,7 +80,9 @@ export class ShtPlan {
8080 readonly qlmIn!: GPUBuffer;
8181 /** Spectral output (analysis). */
8282 readonly qlmOut!: GPUBuffer;
83- /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. */
83+ /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
84+ * stage boundary is observable: a transform is Legendre-then-Fourier, and
85+ * scripts/diagnose-sht.ts tells the two apart by reading this. */
8486 readonly fmBuf!: GPUBuffer;
8587 /** Spatial field [ilat*nphi + iphi], f32. */
8688 readonly spatBuf!: GPUBuffer;
@@ -157,7 +159,7 @@ export class ShtPlan {
157159 this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
158160 self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
159161 self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
160- self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE);
162+ self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
161163 self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
162164 this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
163165 this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
moveopenescclose