src/main.ts at cde22ea37ca3b1ab2cd016bd80054ad1db1e8727

concept-collection / turing-sphere

482 lines · 14.9 KBBlame History Raw

1import { requestShtDevice, describeAdapter } from './sht/sht.ts';

2import { ModelSession } from './mgpu/session.ts';

3import { mModelByKey, presets, type MModel, type Params } from './mgpu/registry.ts';

4import { ModelCompileError, formatFailure } from './mgpu/errors.ts';

5import { EXTERNAL_OPS } from './mgpu/externals.ts';

6import { CodeEditor } from './editor/codeEditor.ts';

7import {

8 formatCommand,

9 resolvePreset,

10 DEFAULT_STEPS,

11 DEFAULT_WARMUP,

12 type RunSpec,

13} from './bench/runSpec.ts';

14import {

15 buildTopology,

16 fillFieldValues,

17 fillColors,

18 type SphereMeshTopology,

19} from './render/sphereMesh.ts';

20import { SphereScene } from './render/SphereScene.ts';

21import { Colorbar } from './render/colorbar.ts';

22import { colormaps, colormapNames } from './render/colormaps.ts';

24const $ = <T extends HTMLElement>(id: string): T =>

25 document.getElementById(id) as T;

27const elModel = $<HTMLSelectElement>('model');

28const elLmax = $<HTMLSelectElement>('lmax');

29const elColormap = $<HTMLSelectElement>('colormap');

30const elRunPause = $<HTMLButtonElement>('runpause');

31const elReseed = $<HTMLButtonElement>('reseed');

32const elResetView = $<HTMLButtonElement>('resetview');

33const elParams = $('params');

34const elPanels = $('panels');

35const elStats = $('stats');

36const elCmd = $('cmd');

37const elCopyCmd = $<HTMLButtonElement>('copycmd');

38const elBlurb = $('blurb');

39const elErr = $('err');

40const elSource = $<HTMLTextAreaElement>('source');

41const elHighlight = $('highlight');

42const elCompiled = $('compiled');

43const elEditorTitle = $('editor-title');

44const elRecompile = $<HTMLButtonElement>('recompile');

45const elRevert = $<HTMLButtonElement>('revert');

47for (const p of presets) {

48 const o = document.createElement('option');

49 o.value = p.key;

50 o.textContent = p.label;

51 elModel.append(o);

52}

53for (const name of colormapNames) {

54 const o = document.createElement('option');

55 o.value = name;

56 o.textContent = name;

57 elColormap.append(o);

58}

59elColormap.value = 'jet';

61/** The model source, with MATLAB highlighting. The host-provided operations are

62 * marked so the boundary between the model and what it is given is visible. */

63const editor = new CodeEditor({

64 textarea: elSource,

65 overlay: elHighlight,

66 external: EXTERNAL_OPS,

67 onInput: (value) => {

68 editedSource = value;

69 elRecompile.textContent = 'Recompile *';

70 },

71});

73/** Timesteps submitted per rendered frame. Nothing is read back between them,

74 * so the batch costs one submit and one readback regardless of size. */

75const STEPS_PER_FRAME = 4;

77/**

78 * Steps in a solver-timing burst, and how often to run one.

79 *

80 * Timing the solver needs a `queue.onSubmittedWorkDone()` to know the work

81 * finished, and in a browser that is an IPC round trip into the GPU process — a

82 * fixed cost of a few milliseconds. Spread over one frame's four steps it would

83 * swamp them on a fast GPU and make the solver look far slower than it is. So the

84 * rate is measured in an occasional larger batch, where the single sync is

85 * amortized the way the desktop benchmark amortizes its own. These are ordinary

86 * steps: the simulation advances by them like any others.

87 */

88const MEASURE_BURST = 32;

89const MEASURE_EVERY_MS = 2000;

91// ---------------------------------------------------------------- state

92let device: GPUDevice | null = null;

93let session: ModelSession | null = null;

94let topo: SphereMeshTopology | null = null;

95let scenes: SphereScene[] = [];

96let colorbars: Colorbar[] = [];

97let valueBufs: Float32Array[] = [];

98let colorBufs: Float32Array[] = [];

99let ranges: { lo: number; hi: number }[] = [];

100let resizeObs: ResizeObserver | null = null;

101

102const initial = resolvePreset(presets[0].key);

103let model: MModel = mModelByKey(initial.model.key)!;

104let params: Params = initial.params;

105/** The .m as edited in the page; `null` while it matches the file. */

106let editedSource: string | null = null;

107let seed = 1;

108let running = false;

109let adapterName = '';

110let pumping = false;

111let solverMs = 0;

112let frameMs = 0;

113let lastMeasure = 0;

114let generation = 0; // bumped on every rebuild to cancel stale pumps

115

116const source = (): string => editedSource ?? model.source;

117

118// ---------------------------------------------------------------- UI wiring

119function buildParamInputs(): void {

120 elParams.replaceChildren();

121 for (const spec of model.params) {

122 const label = document.createElement('label');

123 label.textContent = `${spec.label} `;

124 const input = document.createElement('input');

125 input.type = 'number';

126 input.min = String(spec.min);

127 input.max = String(spec.max);

128 input.step = String(spec.step);

129 input.value = String(params[spec.key]);

130 input.addEventListener('change', () => {

131 const v = Number(input.value);

132 if (Number.isFinite(v)) params[spec.key] = v;

133 // Parameters are uniforms, not constants baked into the kernels, so a

134 // change costs an upload rather than a recompile.

135 session?.setParams(params);

136 updateCommand();

137 });

138 label.append(input);

139 elParams.append(label);

140 }

141}

142

143function applyPreset(presetKey: string): void {

144 const resolved = resolvePreset(presetKey);

145 const next = mModelByKey(resolved.model.key);

146 if (!next) {

147 elErr.textContent = `No .m model for '${resolved.model.key}'`;

148 return;

149 }

150 model = next;

151 params = resolved.params;

152 editedSource = null;

153 editor.value = model.source;

154 elEditorTitle.textContent =

155 `models/${model.key}.m — init() and step(), compiled to WebGPU`;

156 buildParamInputs();

157 elBlurb.textContent = model.blurb;

158 updateCommand();

159}

160

161/** The run currently on screen, as the benchmark's RunSpec. */

162function currentSpec(): RunSpec {

163 return {

164 preset: elModel.value,

165 lmax: Number(elLmax.value),

166 seed,

167 steps: DEFAULT_STEPS,

168 warmup: DEFAULT_WARMUP,

169 params,

170 };

171}

172

173function updateCommand(): void {

174 elCmd.textContent = formatCommand(currentSpec());

175}

176

177elModel.addEventListener('change', () => {

178 applyPreset(elModel.value);

179 void rebuild();

180});

181elLmax.addEventListener('change', () => void rebuild());

182elColormap.addEventListener('change', () => void draw());

183

184function setRunning(next: boolean): void {

185 running = next;

186 elRunPause.textContent = running ? 'Pause' : 'Run';

187 if (running) void pump();

188}

189

190elRunPause.addEventListener('click', () => setRunning(!running));

191elReseed.addEventListener('click', () => {

192 seed = (Math.random() * 2 ** 31) >>> 0;

193 setRunning(false);

194 updateCommand();

195 void reseed();

196});

197elResetView.addEventListener('click', () => {

198 for (const s of scenes) s.resetCamera();

199});

200

201elRecompile.addEventListener('click', () => {

202 editedSource = editor.value;

203 void rebuild();

204});

205elRevert.addEventListener('click', () => {

206 editedSource = null;

207 editor.value = model.source;

208 void rebuild();

209});

210

211// The command reproduces this run's parameters on the desktop; keep it

212// selectable even where the clipboard API is unavailable.

213elCopyCmd.addEventListener('click', () => {

214 const text = elCmd.textContent ?? '';

215 const flash = (msg: string): void => {

216 elCopyCmd.textContent = msg;

217 setTimeout(() => (elCopyCmd.textContent = 'Copy'), 1200);

218 };

219 const selectCommand = (): void => {

220 const range = document.createRange();

221 range.selectNodeContents(elCmd);

222 const sel = getSelection();

223 sel?.removeAllRanges();

224 sel?.addRange(range);

225 flash('Selected');

226 };

227 if (!navigator.clipboard) return selectCommand();

228 navigator.clipboard.writeText(text).then(() => flash('Copied'), selectCommand);

229});

230

231// ---------------------------------------------------------------- setup

232function disposeView(): void {

233 for (const s of scenes) s.dispose();

234 scenes = [];

235 colorbars = [];

236 resizeObs?.disconnect();

237 resizeObs = null;

238 elPanels.replaceChildren();

239}

240

241/** Report a compile failure, and select the offending text in the editor. */

242function reportCompileError(e: unknown): void {

243 elErr.textContent = formatFailure(e, source());

244 elCompiled.textContent = '';

245 if (e instanceof ModelCompileError && e.start !== undefined) {

246 editor.select(e.start, e.end ?? e.start);

247 }

248}

249

250async function rebuild(): Promise<void> {

251 generation++;

252 const gen = generation;

253 setRunning(false);

254 disposeView();

255 session?.destroy();

256 session = null;

257 solverMs = 0;

258 frameMs = 0;

259 lastMeasure = 0;

260 elErr.textContent = '';

261 updateCommand();

262 if (!device) return;

263

264 try {

265 session = await ModelSession.create({

266 device,

267 model,

268 params,

269 lmax: Number(elLmax.value),

270 source: source(),

271 });

272 } catch (e) {

273 reportCompileError(e);

274 return;

275 }

276 if (gen !== generation) return;

277

278 session.seed(seed);

279

280 const plan = session.describe();

281 elCompiled.textContent =

282 `one step compiled to ${plan.step.length} GPU operations:\n` +

283 plan.step.map((l) => ` ${l}`).join('\n');

284 elRecompile.textContent = 'Recompile';

285

286 // mesh + scenes

287 const { nphi } = session.cfg;

288 const phi = new Float64Array(nphi);

289 for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;

290 topo = buildTopology(session.sht.cosTheta, phi);

291

292 const sphereBg = getComputedStyle(document.documentElement)

293 .getPropertyValue('--sphere-bg')

294 .trim();

295 for (let k = 0; k < model.species.length; k++) {

296 const panel = document.createElement('div');

297 panel.className = 'panel';

298 const box = document.createElement('div');

299 box.className = 'sphere-box';

300 const tag = document.createElement('div');

301 tag.className = 'species-tag';

302 tag.textContent = model.species[k];

303 box.append(tag);

304 const side = document.createElement('div');

305 panel.append(box, side);

306 elPanels.append(panel);

307

308 const scene = new SphereScene(

309 box,

310 topo.numVertices,

311 topo.indices,

312 topo.sphereRef,

313 sphereBg || undefined,

314 );

315 scene.fitCamera();

316 scenes.push(scene);

317 colorbars.push(new Colorbar(side));

318 valueBufs[k] = new Float32Array(topo.numVertices);

319 colorBufs[k] = new Float32Array(topo.numVertices * 3);

320 ranges[k] = { lo: NaN, hi: NaN };

321 }

322 for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);

323

324 resizeObs = new ResizeObserver(() => {

325 const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');

326 boxes.forEach((box, i) => {

327 scenes[i]?.resize(box.clientWidth, box.clientHeight);

328 });

329 });

330 elPanels

331 .querySelectorAll<HTMLElement>('.sphere-box')

332 .forEach((box) => resizeObs!.observe(box));

333

334 await draw();

335 updateStats();

336 void pump();

337}

338

339async function reseed(): Promise<void> {

340 if (!session) return;

341 const gen = generation;

342 session.seed(seed);

343 if (gen !== generation) return;

344 for (const r of ranges) {

345 r.lo = NaN;

346 r.hi = NaN;

347 }

348 await draw();

349 updateStats();

350}

351

352// ---------------------------------------------------------------- drawing

353async function draw(): Promise<void> {

354 if (!session || !topo) return;

355 const gen = generation;

356 const cmap = colormaps[elColormap.value] ?? colormaps.viridis;

357 for (let k = 0; k < model.species.length; k++) {

358 // The one readback per frame — the loop is otherwise entirely on the GPU.

359 // A rebuild can land while this is in flight and destroy the buffer being

360 // mapped, which rejects the map; that result is stale anyway, so drop it.

361 let field: Float32Array;

362 try {

363 field = await session.read(model.species[k]);

364 } catch (e) {

365 if (gen !== generation) return;

366 throw e;

367 }

368 if (gen !== generation || !topo) return;

369 fillFieldValues(valueBufs[k], field, topo);

370 let lo = Infinity;

371 let hi = -Infinity;

372 for (const v of valueBufs[k]) {

373 if (v < lo) lo = v;

374 if (v > hi) hi = v;

375 }

376 // smooth the color range in both directions so the shading evolves

377 // gently as the pattern grows (out-of-range values clamp meanwhile)

378 const r = ranges[k];

379 if (!Number.isFinite(r.lo)) {

380 r.lo = lo;

381 r.hi = hi;

382 } else {

383 const a = 0.15;

384 r.lo += a * (lo - r.lo);

385 r.hi += a * (hi - r.hi);

386 }

387 if (r.hi - r.lo < 1e-9) {

388 const mid = (r.hi + r.lo) / 2;

389 r.lo = mid - 5e-10;

390 r.hi = mid + 5e-10;

391 }

392 fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);

393 scenes[k]?.updateColors(colorBufs[k]);

394 colorbars[k]?.update(cmap, r.lo, r.hi);

395 }

396}

397

398function updateStats(): void {

399 if (!session) return;

400 const { nlat, nphi } = session.cfg;

401 const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;

402 const solver =

403 solverMs > 0

404 ? `${solverMs.toFixed(2)} ms/step (${(1000 / solverMs).toFixed(0)} steps/s, ` +

405 `batch of ${MEASURE_BURST}, no readback)`

406 : '—';

407 const frame =

408 frameMs > 0

409 ? `${frameMs.toFixed(1)} ms/frame (${STEPS_PER_FRAME} steps + readback + render)`

410 : '—';

411 elStats.innerHTML =

412 `${kind} · grid ${nlat}×${nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +

413 `${session.sht.fourierMode.toUpperCase()} · solver ${solver} · ${frame} · ` +

414 `t = ${session.t.toFixed(2)} (${session.steps} steps)`;

415}

416

417// ---------------------------------------------------------------- sim loop

418const nextFrame = () => new Promise<number>(requestAnimationFrame);

419

420async function pump(): Promise<void> {

421 if (pumping) return;

422 pumping = true;

423 const gen = generation;

424 try {

425 while (running && session && gen === generation) {

426 // Occasionally, a burst purely to measure the solver rate: many steps,

427 // one sync, nothing read back — directly comparable to the desktop

428 // benchmark's throughput number.

429 if (performance.now() - lastMeasure > MEASURE_EVERY_MS) {

430 const m0 = performance.now();

431 session.step(MEASURE_BURST);

432 await session.sync();

433 if (gen !== generation) break;

434 solverMs = (performance.now() - m0) / MEASURE_BURST;

435 lastMeasure = performance.now();

436 }

437

438 // The frame itself. No explicit sync here — draw()'s readback already

439 // waits for the steps, so asking twice would only add a round trip.

440 const t0 = performance.now();

441 session.step(STEPS_PER_FRAME);

442 await draw();

443 if (gen !== generation) break;

444 frameMs = frameMs === 0

445 ? performance.now() - t0

446 : frameMs + 0.05 * (performance.now() - t0 - frameMs);

447 updateStats();

448 await nextFrame();

449 }

450 if (gen === generation) {

451 await draw();

452 updateStats();

453 }

454 } finally {

455 pumping = false;

456 }

457}

458

459// ---------------------------------------------------------------- boot

460async function boot(): Promise<void> {

461 elModel.value = presets[0].key;

462 applyPreset(presets[0].key);

463 try {

464 device = await requestShtDevice();

465 adapterName = await describeAdapter(device);

466 } catch (e) {

467 device = null;

468 elErr.textContent =

469 `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +

470 `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +

471 `needs a WebGPU-capable browser (Chrome/Edge 113+).`;

472 return;

473 }

474 device.lost.then((info) => {

475 if (info.reason !== 'destroyed') {

476 elErr.textContent = `WebGPU device lost: ${info.message}`;

477 }

478 });

479 await rebuild();

480}

481

482void boot();