bench:sht: check one round trip before timing thousands - concept-collection/turing-sphere

concept-collection / turing-sphere

bench:sht: check one round trip before timing thousands

A round trip is a projection — analys(synth(q)) == q for a band-limited q, with exact Gauss quadrature and nphi past the aliasing limit — so a single one should return the input to fp32 round-off. Check that first, and report the relative L2. It makes a bad result say which kind it is. On an RTX PRO 6000 the timing run finished and then reported a final state with no finite values in it, which could equally have been transforms that never worked on that device or something that diverged over the 2250 iterations the run does. Those need different fixes and the difference was a manual bisection on --steps. Now the run says so and stops, instead of spending minutes timing a transform that does not transform. Also here: - submit the warmup and the digest run in chunks rather than putting an arbitrary --steps worth of dispatches into one command buffer. - compare-native.mjs: when a run exits non-zero but still printed its JSON, describe what it reported — which round trip failed, the device, the Fourier stage — instead of dumping the tail of its output. It had already said everything needed.

Jeremy Magland <jmagland@flatironinstitute.org> committed 20 days ago commit 721b8df99ca2 parent e302a6e Browse files

2 changed files+118−7

scripts/bench-sht.tsmodified+84−6View file

@@ -23,7 +23,7 @@

2323 import { ShtPlan, requestShtDevice, describeAdapter, type ShtBinding } from '../src/sht/sht.ts';

2424 import { lmIndex } from '../src/sht/layout.ts';

2525 import { makeRand } from '../src/mgpu/noise.ts';

26-import { digestOf, formatDigest } from '../src/mgpu/digest.ts';

26+import { digestOf, formatDigest, relL2 } from '../src/mgpu/digest.ts';

2727 import {

2828 parseArgs,

2929 configForSpec,

@@ -38,6 +38,8 @@ import { presets } from '../src/mgpu/registry.ts';

3838 import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';

3939 import { writeFileSync } from 'node:fs';

4040

41+const BENCH_SHT_COMMAND = 'npm run bench:sht --';

42+

4143 const USAGE = `usage: npm run bench:sht -- [options]

4244

4345 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})

@@ -196,6 +198,11 @@ try {

196198 pass.end();

197199 device!.queue.submit([enc.finish()]);

198200 };

201+ /** `submit` in chunks, so an arbitrary round-trip count does not build one

202+ * command buffer with tens of thousands of dispatches in it. */

203+ const submitAll = (n: number, chunk = batch): void => {

204+ for (let done = 0; done < n; done += chunk) submit(Math.min(chunk, n - done));

205+ };

199206 const seed = (): void => {

200207 cur = 0;

201208 const q0 = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);

@@ -226,8 +233,80 @@ try {

226233 );

227234 }

228235

236+ /*

237+ * Before timing anything: does one round trip work on this device?

238+ *

239+ * analys(synth(q)) is the identity for a band-limited q — exact Gauss

240+ * quadrature, nphi past the aliasing limit — so a single round trip should

241+ * return the input to fp32 round-off, and this is the sharpest check the

242+ * transforms are working at all here. Doing it separately means a bad result

243+ * says *which* it is: broken from the first transform, or drifted over the

244+ * thousands of iterations the timing run does. Otherwise that is a manual

245+ * bisection on --steps.

246+ */

247+ seed();

248+ const input = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);

249+ submit(1);

250+ await done();

251+ const afterOne = await readSpectrum();

252+ const firstFinite = afterOne.every((v) => Number.isFinite(v));

253+ const firstRelL2 = firstFinite ? relL2(afterOne, input) : NaN;

254+ if (!wantJson) {

255+ console.log(

256+ ` one round trip: ${

257+ firstFinite

258+ ? `back to the input to ${firstRelL2.toExponential(2)} relative L2`

259+ : 'NOT FINITE'

260+ }`,

261+ );

262+ }

263+ if (!firstFinite || !(firstRelL2 < 1e-3)) {

264+ // Report it the same way a good run reports itself, so a caller reading

265+ // --json learns what went wrong and on which device rather than having to

266+ // scrape stderr. Then say it in prose and stop: timing a transform that does

267+ // not transform is a waste of minutes.

268+ if (wantJson) {

269+ console.log(

270+ JSON.stringify(

271+ {

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,

277+ firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },

278+ throughput: null,

279+ latency: null,

280+ digest: null,

281+ input: null,

282+ state: { min: null, max: null, finite: firstFinite },

283+ },

284+ null,

285+ 2,

286+ ),

287+ );

288+ }

289+ const detail = firstFinite

290+ ? `it came back ${firstRelL2.toExponential(3)} away from the input, which is far\n` +

291+ ` outside fp32 round-off (~1e-7)`

292+ : `it came back with no finite values at all`;

293+ fail(

294+ `a single spectral -> grid -> spectral round trip does not round-trip on this\n` +

295+ ` device: ${detail}.\n\n` +

296+ ` That is a correctness problem in the transforms here, not a benchmarking one,\n` +

297+ ` so there is nothing worth timing yet. Adapter: ${adapter || '(unknown)'};\n` +

298+ ` Fourier stage: ${plan.fourierMode.toUpperCase()}.\n\n` +

299+ ` Worth trying, in order:\n` +

300+ ` npm run test:node the repo's own transform check against\n` +

301+ ` its f64 CPU twin, on this device\n` +

302+ ` ${BENCH_SHT_COMMAND} --fourier dft the other Fourier stage; if this works,\n` +

303+ ` the WGSL FFT is the problem\n` +

304+ ` ${BENCH_SHT_COMMAND} --lmax 15 does it depend on the grid size?`,

305+ );

306+ }

307+

229308 seed();

230- submit(spec.warmup);

309+ submitAll(spec.warmup);

231310 await done();

232311

233312 // --- throughput: batches submitted together, awaited once each ---

@@ -272,11 +351,9 @@ try {

272351 let inputDigest = null;

273352 let state: Float32Array | null = null;

274353 if (wantDigest) {

275- const input = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);

276354 inputDigest = digestOf(input, plan.fourierMode, adapter);

277- cur = 0;

278- device.queue.writeBuffer(qlm[0], 0, input as Float32Array<ArrayBuffer>);

279- submit(spec.steps);

355+ seed();

356+ submitAll(spec.steps);

280357 await done();

281358 state = await readSpectrum();

282359 digest = digestOf(state, plan.fourierMode, adapter);

@@ -301,6 +378,7 @@ try {

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

302379 grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },

303380 fourier: plan.fourierMode,

381+ firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },

304382 throughput: {

305383 batch,

306384 msPerStep: throughputMs,

scripts/compare-native.mjsmodified+34−1View file

@@ -140,6 +140,34 @@ function failureDetail(r, cmd, args) {

140140 ].join('\n');

141141 }

142142

143+/**

144+ * Why a run that reported itself is still not usable. Every producer records the

145+ * device, the Fourier stage and whether its result stayed finite, so a wrong

146+ * answer can be described rather than dumped.

147+ */

148+function badResult(json) {

149+ const b = json.backend ?? {};

150+ const lines = [];

151+ const first = json.firstRoundTrip;

152+ if (first && first.finite === false) {

153+ lines.push('a single spectral -> grid -> spectral round trip returns no finite values,');

154+ lines.push('so the transforms are wrong on this device — nothing here is worth timing.');

155+ } else if (first && !(first.relL2 < 1e-3)) {

156+ lines.push(

157+ `a single round trip comes back ${Number(first.relL2).toExponential(2)} away from its`,

158+ );

159+ lines.push('input, far outside fp32 round-off (~1e-7). The transforms are wrong here.');

160+ } else if (json.state && json.state.finite === false) {

161+ lines.push('it ran, but the final state has no finite values — one round trip is fine,');

162+ lines.push('so something diverges over the length of the run.');

163+ } else {

164+ lines.push('it reported a failure without saying why; run it alone.');

165+ }

166+ lines.push(`device: ${b.adapter || '(unknown)'}`);

167+ if (json.fourier) lines.push(`Fourier stage: ${String(json.fourier).toUpperCase()}`);

168+ return lines.map((l, i) => (i === 0 ? l : ` ${l}`)).join('\n');

169+}

170+

143171 /** Run one side and parse its --json output. `ok: false` with a reason if it is

144172 * not available here — a missing binary, no adapter, no CUDA. */

145173 function run(label, cmd, args, statePath) {

@@ -152,8 +180,13 @@ function run(label, cmd, args, statePath) {

152180 maxBuffer: 256 * 1024 * 1024,

153181 });

154182 if (r.error) return { label, ok: false, why: r.error.message };

155- if (r.status !== 0) return { label, ok: false, why: failureDetail(r, cmd, full) };

156183 const json = extractJson(r.stdout);

184+ if (r.status !== 0) {

185+ // A run that failed but still reported itself is the interesting case: it

186+ // computed something wrong rather than failing to start, and it already said

187+ // what and on which device. Use that instead of dumping its output.

188+ return { label, ok: false, why: json ? badResult(json) : failureDetail(r, cmd, full) };

189+ }

157190 if (!json) {

158191 return {

159192 label,