1import { requestShtDevice, describeAdapter } from './sht/sht.ts';
2import { gridForLmax } from './sht/layout.ts';
3import { ModelSession } from './mgpu/session.ts';
4import { mModelByKey, presets, type MModel, type Params } from './mgpu/registry.ts';
5import { ModelCompileError, formatFailure } from './mgpu/errors.ts';
6import { EXTERNAL_OPS } from './mgpu/externals.ts';
7import { CodeEditor } from './editor/codeEditor.ts';
8import {
9 formatCommand,
10 resolvePreset,
11 DEFAULT_STEPS,
12 DEFAULT_WARMUP,
13 type RunSpec,
14} from './bench/runSpec.ts';
15import {
16 buildTopology,
17 fillFieldValues,
18 fillColors,
19 type SphereMeshTopology,
20} from './render/sphereMesh.ts';
21import { SphereScene } from './render/SphereScene.ts';
22import { Colorbar } from './render/colorbar.ts';
23import { colormaps, colormapNames } from './render/colormaps.ts';
25const $ = <T extends HTMLElement>(id: string): T =>
26 document.getElementById(id) as T;
28const elModel = $<HTMLSelectElement>('model');
29const elLmax = $<HTMLSelectElement>('lmax');
30const elOversample = $<HTMLSelectElement>('oversample');
31const elColormap = $<HTMLSelectElement>('colormap');
32const elRunPause = $<HTMLButtonElement>('runpause');
33const elBenchmark = $<HTMLButtonElement>('benchmark');
34const elReseed = $<HTMLButtonElement>('reseed');
35const elResetView = $<HTMLButtonElement>('resetview');
36const elParams = $('params');
37const elPanels = $('panels');
38const elStats = $('stats');
39const elBenchResult = $('benchresult');
40const elCmd = $('cmd');
41const elCopyCmd = $<HTMLButtonElement>('copycmd');
42const elBlurb = $('blurb');
43const elErr = $('err');
44const elSource = $<HTMLTextAreaElement>('source');
45const elHighlight = $('highlight');
46const elCompiled = $('compiled');
47const elEditorTitle = $('editor-title');
48const elRecompile = $<HTMLButtonElement>('recompile');
49const elRevert = $<HTMLButtonElement>('revert');
51for (const p of presets) {
52 const o = document.createElement('option');
53 o.value = p.key;
54 o.textContent = p.label;
55 elModel.append(o);
56}
57for (const name of colormapNames) {
58 const o = document.createElement('option');
59 o.value = name;
60 o.textContent = name;
61 elColormap.append(o);
62}
63elColormap.value = 'jet';
65/** The model source, with MATLAB highlighting. The host-provided operations are
66 * marked so the boundary between the model and what it is given is visible. */
67const editor = new CodeEditor({
68 textarea: elSource,
69 overlay: elHighlight,
70 external: EXTERNAL_OPS,
71 onInput: (value) => {
72 editedSource = value;
73 elRecompile.textContent = 'Recompile *';
74 },
75});
77/** Timesteps submitted per rendered frame. Nothing is read back between them,
78 * so the batch costs one submit and one readback regardless of size. */
79const STEPS_PER_FRAME = 4;
81/**
82 * Steps in a solver-timing burst, and how often to run one.
83 *
84 * Timing the solver needs a `queue.onSubmittedWorkDone()` to know the work
85 * finished, and in a browser that is an IPC round trip into the GPU process — a
86 * fixed cost of a few milliseconds. Spread over one frame's four steps it would
87 * swamp them on a fast GPU and make the solver look far slower than it is. So the
88 * rate is measured in an occasional larger batch, where the single sync is
89 * amortized the way the desktop benchmark amortizes its own. The state is
90 * snapshotted and restored around the batch, so measuring never advances the
91 * simulation — otherwise the pattern would visibly lurch forward at every
92 * measurement.
93 */
94const MEASURE_BURST = 32;
95const MEASURE_EVERY_MS = 2000;
97/**
98 * 'auto' display oversampling targets this many render latitudes: the factor is
99 * the smallest power of two (up to 4) that reaches it. A solver grid already
100 * this fine gains nothing visually and is not oversampled.
101 */
102const AUTO_RENDER_NLAT = 256;
104/** The display oversampling factor the UI currently asks for. */
105function resolveOversample(): number {
106 if (elOversample.value !== 'auto') return Number(elOversample.value);
107 const { nlat } = gridForLmax(Number(elLmax.value), model.pdeg);
108 let os = 1;
109 while (os < 4 && os * nlat < AUTO_RENDER_NLAT) os *= 2;
110 return os;
111}
113// ---------------------------------------------------------------- state
114let device: GPUDevice | null = null;
115let session: ModelSession | null = null;
116let topo: SphereMeshTopology | null = null;
117let scenes: SphereScene[] = [];
118let colorbars: Colorbar[] = [];
119let valueBufs: Float32Array[] = [];
120let colorBufs: Float32Array[] = [];
121let ranges: { lo: number; hi: number }[] = [];
122let resizeObs: ResizeObserver | null = null;
124const initial = resolvePreset(presets[0].key);
125let model: MModel = mModelByKey(initial.model.key)!;
126let params: Params = initial.params;
127/** The .m as edited in the page; `null` while it matches the file. */
128let editedSource: string | null = null;
129let seed = 1;
130let running = false;
131let adapterName = '';
132let pumping = false;
133let solverMs = 0;
134let frameMs = 0;
135let lastMeasure = 0;
136let generation = 0; // bumped on every rebuild to cancel stale pumps
138const source = (): string => editedSource ?? model.source;
140// ---------------------------------------------------------------- UI wiring
141function buildParamInputs(): void {
142 elParams.replaceChildren();
143 for (const spec of model.params) {
144 const label = document.createElement('label');
145 label.textContent = `${spec.label} `;
146 const input = document.createElement('input');
147 input.type = 'number';
148 input.min = String(spec.min);
149 input.max = String(spec.max);
150 input.step = String(spec.step);
151 input.value = String(params[spec.key]);
152 input.addEventListener('change', () => {
153 const v = Number(input.value);
154 if (Number.isFinite(v)) params[spec.key] = v;
155 // Parameters are uniforms, not constants baked into the kernels, so a
156 // change costs an upload rather than a recompile.
157 session?.setParams(params);
158 updateCommand();
159 });
160 label.append(input);
161 elParams.append(label);
162 }
163}
165function applyPreset(presetKey: string): void {
166 const resolved = resolvePreset(presetKey);
167 const next = mModelByKey(resolved.model.key);
168 if (!next) {
169 elErr.textContent = `No .m model for '${resolved.model.key}'`;
170 return;
171 }
172 model = next;
173 params = resolved.params;
174 editedSource = null;
175 editor.value = model.source;
176 elEditorTitle.textContent =
177 `models/${model.key}.m — init() and step(), compiled to WebGPU`;
178 buildParamInputs();
179 elBlurb.textContent = model.blurb;
180 updateCommand();
181}
183/** The run currently on screen, as the benchmark's RunSpec. */
184function currentSpec(): RunSpec {
185 return {
186 preset: elModel.value,
187 lmax: Number(elLmax.value),
188 seed,
189 steps: DEFAULT_STEPS,
190 warmup: DEFAULT_WARMUP,
191 params,
192 };
193}
195function updateCommand(): void {
196 elCmd.textContent = formatCommand(currentSpec());
197}
199elModel.addEventListener('change', () => {
200 applyPreset(elModel.value);
201 void rebuild();
202});
203elLmax.addEventListener('change', () => void rebuild());
204// Oversampling is display-only, so it swaps the render grid in place rather
205// than rebuilding the run. Serialized: a rapid second change waits its turn.
206let viewChange = Promise.resolve();
207elOversample.addEventListener('change', () => {
208 viewChange = viewChange.then(() => applyOversample());
209});
210elColormap.addEventListener('change', () => void draw());
212function setRunning(next: boolean): void {
213 running = next;
214 elRunPause.textContent = running ? 'Pause' : 'Run';
215 if (running) void pump();
216}
218elRunPause.addEventListener('click', () => setRunning(!running));
219elBenchmark.addEventListener('click', () => void benchmark());
220elReseed.addEventListener('click', () => {
221 seed = (Math.random() * 2 ** 31) >>> 0;
222 setRunning(false);
223 updateCommand();
224 void reseed();
225});
226elResetView.addEventListener('click', () => {
227 for (const s of scenes) s.resetCamera();
228});
230elRecompile.addEventListener('click', () => {
231 editedSource = editor.value;
232 void rebuild();
233});
234elRevert.addEventListener('click', () => {
235 editedSource = null;
236 editor.value = model.source;
237 void rebuild();
238});
240// The command reproduces this run's parameters on the desktop; keep it
241// selectable even where the clipboard API is unavailable.
242elCopyCmd.addEventListener('click', () => {
243 const text = elCmd.textContent ?? '';
244 const flash = (msg: string): void => {
245 elCopyCmd.textContent = msg;
246 setTimeout(() => (elCopyCmd.textContent = 'Copy'), 1200);
247 };
248 const selectCommand = (): void => {
249 const range = document.createRange();
250 range.selectNodeContents(elCmd);
251 const sel = getSelection();
252 sel?.removeAllRanges();
253 sel?.addRange(range);
254 flash('Selected');
255 };
256 if (!navigator.clipboard) return selectCommand();
257 navigator.clipboard.writeText(text).then(() => flash('Copied'), selectCommand);
258});
260// ---------------------------------------------------------------- setup
261function disposeView(): void {
262 for (const s of scenes) s.dispose();
263 scenes = [];
264 colorbars = [];
265 resizeObs?.disconnect();
266 resizeObs = null;
267 elPanels.replaceChildren();
268}
270/**
271 * Build the mesh, scenes, colorbars and per-species buffers on the current
272 * render grid. Call disposeView() first. The color ranges are kept if present,
273 * so a display-only rebuild (an oversampling change) does not pop the shading;
274 * a full rebuild clears `ranges` beforehand.
275 */
276function buildView(): void {
277 if (!session) return;
278 const view = session.viewSht;
279 const { nphi } = view.cfg;
280 const phi = new Float64Array(nphi);
281 for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
282 topo = buildTopology(view.cosTheta, phi);
284 const sphereBg = getComputedStyle(document.documentElement)
285 .getPropertyValue('--sphere-bg')
286 .trim();
287 for (let k = 0; k < model.species.length; k++) {
288 const panel = document.createElement('div');
289 panel.className = 'panel';
290 const box = document.createElement('div');
291 box.className = 'sphere-box';
292 const tag = document.createElement('div');
293 tag.className = 'species-tag';
294 tag.textContent = model.species[k];
295 box.append(tag);
296 const side = document.createElement('div');
297 panel.append(box, side);
298 elPanels.append(panel);
300 const scene = new SphereScene(
301 box,
302 topo.numVertices,
303 topo.indices,
304 topo.sphereRef,
305 sphereBg || undefined,
306 );
307 scene.fitCamera();
308 scenes.push(scene);
309 colorbars.push(new Colorbar(side));
310 valueBufs[k] = new Float32Array(topo.numVertices);
311 colorBufs[k] = new Float32Array(topo.numVertices * 3);
312 if (!ranges[k]) ranges[k] = { lo: NaN, hi: NaN };
313 }
314 for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
316 resizeObs = new ResizeObserver(() => {
317 const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
318 boxes.forEach((box, i) => {
319 scenes[i]?.resize(box.clientWidth, box.clientHeight);
320 });
321 });
322 elPanels
323 .querySelectorAll<HTMLElement>('.sphere-box')
324 .forEach((box) => resizeObs!.observe(box));
325}
327/**
328 * Apply the UI's oversampling choice to the running session. Display-only: the
329 * session and its state survive; only the display plan, mesh and scenes are
330 * rebuilt, keeping the camera pose and color ranges. The pump is drained first
331 * so no readback is in flight on the plan being replaced.
332 */
333async function applyOversample(): Promise<void> {
334 if (!session) return;
335 const gen = generation;
336 const os = resolveOversample();
337 if (os === session.oversample) return;
338 const wasRunning = running;
339 setRunning(false);
340 while (pumping) await nextFrame();
341 if (gen !== generation || !session) return;
342 await session.setOversample(os);
343 if (gen !== generation || !session) return;
344 const cam = scenes[0]?.cameraState();
345 disposeView();
346 buildView();
347 if (cam) for (const s of scenes) s.setCameraState(cam);
348 await draw();
349 updateStats();
350 if (wasRunning) setRunning(true);
351}
353/** Report a compile failure, and select the offending text in the editor. */
354function reportCompileError(e: unknown): void {
355 elErr.textContent = formatFailure(e, source());
356 elCompiled.textContent = '';
357 if (e instanceof ModelCompileError && e.start !== undefined) {
358 editor.select(e.start, e.end ?? e.start);
359 }
360}
362async function rebuild(): Promise<void> {
363 generation++;
364 const gen = generation;
365 setRunning(false);
366 disposeView();
367 session?.destroy();
368 session = null;
369 solverMs = 0;
370 frameMs = 0;
371 lastMeasure = 0;
372 elErr.textContent = '';
373 updateCommand();
374 if (!device) return;
376 try {
377 session = await ModelSession.create({
378 device,
379 model,
380 params,
381 lmax: Number(elLmax.value),
382 source: source(),
383 oversample: resolveOversample(),
384 });
385 } catch (e) {
386 reportCompileError(e);
387 return;
388 }
389 if (gen !== generation) return;
391 session.seed(seed);
393 const plan = session.describe();
394 elCompiled.textContent =
395 `one step compiled to ${plan.step.length} GPU operations:\n` +
396 plan.step.map((l) => ` ${l}`).join('\n');
397 elRecompile.textContent = 'Recompile';
399 ranges = [];
400 buildView();
402 await draw();
403 updateStats();
404 void pump();
405}
407async function reseed(): Promise<void> {
408 if (!session) return;
409 const gen = generation;
410 session.seed(seed);
411 if (gen !== generation) return;
412 for (const r of ranges) {
413 r.lo = NaN;
414 r.hi = NaN;
415 }
416 await draw();
417 updateStats();
418}
420// ---------------------------------------------------------------- drawing
421async function draw(): Promise<void> {
422 if (!session || !topo) return;
423 const gen = generation;
424 const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
425 for (let k = 0; k < model.species.length; k++) {
426 // The one readback per frame — the loop is otherwise entirely on the GPU.
427 // A rebuild can land while this is in flight and destroy the buffer being
428 // mapped, which rejects the map; that result is stale anyway, so drop it.
429 let field: Float32Array;
430 try {
431 field = await session.readSpecies(k);
432 } catch (e) {
433 if (gen !== generation) return;
434 throw e;
435 }
436 if (gen !== generation || !topo) return;
437 fillFieldValues(valueBufs[k], field, topo);
438 let lo = Infinity;
439 let hi = -Infinity;
440 for (const v of valueBufs[k]) {
441 if (v < lo) lo = v;
442 if (v > hi) hi = v;
443 }
444 // smooth the color range in both directions so the shading evolves
445 // gently as the pattern grows (out-of-range values clamp meanwhile)
446 const r = ranges[k];
447 if (!Number.isFinite(r.lo)) {
448 r.lo = lo;
449 r.hi = hi;
450 } else {
451 const a = 0.15;
452 r.lo += a * (lo - r.lo);
453 r.hi += a * (hi - r.hi);
454 }
455 if (r.hi - r.lo < 1e-9) {
456 const mid = (r.hi + r.lo) / 2;
457 r.lo = mid - 5e-10;
458 r.hi = mid + 5e-10;
459 }
460 fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);
461 scenes[k]?.updateColors(colorBufs[k]);
462 colorbars[k]?.update(cmap, r.lo, r.hi);
463 }
464}
466function updateStats(): void {
467 if (!session) return;
468 const { nlat, nphi } = session.cfg;
469 const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
470 const solver =
471 solverMs > 0
472 ? `<b>${solverMs.toFixed(2)} ms/step</b> (${(1000 / solverMs).toFixed(0)} steps/s, ` +
473 `batch of ${MEASURE_BURST}, no readback)`
474 : '—';
475 const frame =
476 frameMs > 0
477 ? `${frameMs.toFixed(1)} ms/frame (${STEPS_PER_FRAME} steps + readback + render)`
478 : '—';
479 const view = session.viewSht.cfg;
480 const render =
481 session.oversample > 1
482 ? ` (display ${view.nlat}×${view.nphi}, ${session.oversample}×)`
483 : '';
484 elStats.innerHTML =
485 `<b>${kind}</b> · grid ${nlat}×${nphi}${render} · nlm ${session.sht.nlm.toLocaleString()} · ` +
486 `${session.sht.fourierMode.toUpperCase()} · solver ${solver} · ${frame} · ` +
487 `t = <b>${session.t.toFixed(2)}</b> (${session.steps} steps)`;
488}
490// ---------------------------------------------------------------- sim loop
491const nextFrame = () => new Promise<number>(requestAnimationFrame);
493async function pump(): Promise<void> {
494 if (pumping) return;
495 pumping = true;
496 const gen = generation;
497 try {
498 while (running && session && gen === generation) {
499 // Occasionally, a burst purely to measure the solver rate: many steps,
500 // one sync, nothing read back — directly comparable to the desktop
501 // benchmark's throughput number. State-preserving: the display and
502 // model time are unaffected.
503 if (performance.now() - lastMeasure > MEASURE_EVERY_MS) {
504 const ms = await session.measure(MEASURE_BURST);
505 if (gen !== generation) break;
506 solverMs = ms;
507 lastMeasure = performance.now();
508 }
510 // The frame itself. No explicit sync here — draw()'s readback already
511 // waits for the steps, so asking twice would only add a round trip.
512 const t0 = performance.now();
513 session.step(STEPS_PER_FRAME);
514 await draw();
515 if (gen !== generation) break;
516 frameMs = frameMs === 0
517 ? performance.now() - t0
518 : frameMs + 0.05 * (performance.now() - t0 - frameMs);
519 updateStats();
520 await nextFrame();
521 }
522 if (gen === generation) {
523 await draw();
524 updateStats();
525 }
526 } finally {
527 pumping = false;
528 }
529}
531/**
532 * Sustained solver benchmark, in the page.
533 *
534 * The same measurement `npm run bench` makes: batches of steps submitted
535 * together, waited for, never read back, with no rendering and no animation
536 * pacing in between. That makes it directly comparable to the terminal number,
537 * which is the only way to tell a genuinely slower browser GPU stack apart from
538 * the costs the app adds on top.
539 *
540 * It also reports the ramp — the first third of the run against the last. GPUs
541 * downclock when idle, and an animation-paced loop leaves them idle most of every
542 * frame, so a large ramp means the app's steady-state number is limited by clocks
543 * rather than by the work.
544 *
545 * These are ordinary steps: the simulation advances by them.
546 */
547async function benchmark(): Promise<void> {
548 if (!session) return;
549 setRunning(false);
550 const BATCH = 32;
551 const DURATION_MS = 2000;
552 elBenchResult.textContent = 'benchmarking…';
553 await nextFrame();
555 const gen = generation;
556 const perStep: number[] = [];
557 const t0 = performance.now();
558 while (performance.now() - t0 < DURATION_MS) {
559 const b0 = performance.now();
560 session.step(BATCH);
561 await session.sync();
562 if (gen !== generation) return;
563 perStep.push((performance.now() - b0) / BATCH);
564 }
566 const mean = (xs: number[]): number => xs.reduce((a, b) => a + b, 0) / xs.length;
567 const all = mean(perStep);
568 const best = Math.min(...perStep);
569 const third = Math.max(1, Math.floor(perStep.length / 3));
570 const first = mean(perStep.slice(0, third));
571 const last = mean(perStep.slice(-third));
572 const steps = perStep.length * BATCH;
574 elBenchResult.innerHTML =
575 `sustained solver: <b>${all.toFixed(2)} ms/step</b> ` +
576 `(${(1000 / all).toFixed(0)} steps/s) · best ${best.toFixed(2)} · ` +
577 `ramp ${(first / last).toFixed(2)}× (${first.toFixed(2)} → ${last.toFixed(2)}) · ` +
578 `${steps} steps in batches of ${BATCH} · ` +
579 `compare with <code>npm run bench -- --lmax ${session.cfg.lmax}</code>`;
580 await draw();
581 updateStats();
582}
584// ---------------------------------------------------------------- boot
585async function boot(): Promise<void> {
586 elModel.value = presets[0].key;
587 applyPreset(presets[0].key);
588 try {
589 device = await requestShtDevice();
590 adapterName = await describeAdapter(device);
591 } catch (e) {
592 device = null;
593 elErr.textContent =
594 `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
595 `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +
596 `needs a WebGPU-capable browser (Chrome/Edge 113+).`;
597 return;
598 }
599 device.lost.then((info) => {
600 if (info.reason !== 'destroyed') {
601 elErr.textContent = `WebGPU device lost: ${info.message}`;
602 }
603 });
604 await rebuild();
605}
607void boot();