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 elBenchmark = $<HTMLButtonElement>('benchmark');
32const elReseed = $<HTMLButtonElement>('reseed');
33const elResetView = $<HTMLButtonElement>('resetview');
34const elParams = $('params');
35const elPanels = $('panels');
36const elStats = $('stats');
37const elBenchResult = $('benchresult');
38const elCmd = $('cmd');
39const elCopyCmd = $<HTMLButtonElement>('copycmd');
40const elBlurb = $('blurb');
41const elErr = $('err');
42const elSource = $<HTMLTextAreaElement>('source');
43const elHighlight = $('highlight');
44const elCompiled = $('compiled');
45const elEditorTitle = $('editor-title');
46const elRecompile = $<HTMLButtonElement>('recompile');
47const elRevert = $<HTMLButtonElement>('revert');
49for (const p of presets) {
50 const o = document.createElement('option');
51 o.value = p.key;
52 o.textContent = p.label;
53 elModel.append(o);
54}
55for (const name of colormapNames) {
56 const o = document.createElement('option');
57 o.value = name;
58 o.textContent = name;
59 elColormap.append(o);
60}
61elColormap.value = 'jet';
63/** The model source, with MATLAB highlighting. The host-provided operations are
64 * marked so the boundary between the model and what it is given is visible. */
65const editor = new CodeEditor({
66 textarea: elSource,
67 overlay: elHighlight,
68 external: EXTERNAL_OPS,
69 onInput: (value) => {
70 editedSource = value;
71 elRecompile.textContent = 'Recompile *';
72 },
73});
75/** Timesteps submitted per rendered frame. Nothing is read back between them,
76 * so the batch costs one submit and one readback regardless of size. */
77const STEPS_PER_FRAME = 4;
79/**
80 * Steps in a solver-timing burst, and how often to run one.
81 *
82 * Timing the solver needs a `queue.onSubmittedWorkDone()` to know the work
83 * finished, and in a browser that is an IPC round trip into the GPU process — a
84 * fixed cost of a few milliseconds. Spread over one frame's four steps it would
85 * swamp them on a fast GPU and make the solver look far slower than it is. So the
86 * rate is measured in an occasional larger batch, where the single sync is
87 * amortized the way the desktop benchmark amortizes its own. These are ordinary
88 * steps: the simulation advances by them like any others.
89 */
90const MEASURE_BURST = 32;
91const MEASURE_EVERY_MS = 2000;
93// ---------------------------------------------------------------- state
94let device: GPUDevice | null = null;
95let session: ModelSession | null = null;
96let topo: SphereMeshTopology | null = null;
97let scenes: SphereScene[] = [];
98let colorbars: Colorbar[] = [];
99let valueBufs: Float32Array[] = [];
100let colorBufs: Float32Array[] = [];
101let ranges: { lo: number; hi: number }[] = [];
102let resizeObs: ResizeObserver | null = null;
104const initial = resolvePreset(presets[0].key);
105let model: MModel = mModelByKey(initial.model.key)!;
106let params: Params = initial.params;
107/** The .m as edited in the page; `null` while it matches the file. */
108let editedSource: string | null = null;
109let seed = 1;
110let running = false;
111let adapterName = '';
112let pumping = false;
113let solverMs = 0;
114let frameMs = 0;
115let lastMeasure = 0;
116let generation = 0; // bumped on every rebuild to cancel stale pumps
118const source = (): string => editedSource ?? model.source;
120// ---------------------------------------------------------------- UI wiring
121function buildParamInputs(): void {
122 elParams.replaceChildren();
123 for (const spec of model.params) {
124 const label = document.createElement('label');
125 label.textContent = `${spec.label} `;
126 const input = document.createElement('input');
127 input.type = 'number';
128 input.min = String(spec.min);
129 input.max = String(spec.max);
130 input.step = String(spec.step);
131 input.value = String(params[spec.key]);
132 input.addEventListener('change', () => {
133 const v = Number(input.value);
134 if (Number.isFinite(v)) params[spec.key] = v;
135 // Parameters are uniforms, not constants baked into the kernels, so a
136 // change costs an upload rather than a recompile.
137 session?.setParams(params);
138 updateCommand();
139 });
140 label.append(input);
141 elParams.append(label);
142 }
143}
145function applyPreset(presetKey: string): void {
146 const resolved = resolvePreset(presetKey);
147 const next = mModelByKey(resolved.model.key);
148 if (!next) {
149 elErr.textContent = `No .m model for '${resolved.model.key}'`;
150 return;
151 }
152 model = next;
153 params = resolved.params;
154 editedSource = null;
155 editor.value = model.source;
156 elEditorTitle.textContent =
157 `models/${model.key}.m — init() and step(), compiled to WebGPU`;
158 buildParamInputs();
159 elBlurb.textContent = model.blurb;
160 updateCommand();
161}
163/** The run currently on screen, as the benchmark's RunSpec. */
164function currentSpec(): RunSpec {
165 return {
166 preset: elModel.value,
167 lmax: Number(elLmax.value),
168 seed,
169 steps: DEFAULT_STEPS,
170 warmup: DEFAULT_WARMUP,
171 params,
172 };
173}
175function updateCommand(): void {
176 elCmd.textContent = formatCommand(currentSpec());
177}
179elModel.addEventListener('change', () => {
180 applyPreset(elModel.value);
181 void rebuild();
182});
183elLmax.addEventListener('change', () => void rebuild());
184elColormap.addEventListener('change', () => void draw());
186function setRunning(next: boolean): void {
187 running = next;
188 elRunPause.textContent = running ? 'Pause' : 'Run';
189 if (running) void pump();
190}
192elRunPause.addEventListener('click', () => setRunning(!running));
193elBenchmark.addEventListener('click', () => void benchmark());
194elReseed.addEventListener('click', () => {
195 seed = (Math.random() * 2 ** 31) >>> 0;
196 setRunning(false);
197 updateCommand();
198 void reseed();
199});
200elResetView.addEventListener('click', () => {
201 for (const s of scenes) s.resetCamera();
202});
204elRecompile.addEventListener('click', () => {
205 editedSource = editor.value;
206 void rebuild();
207});
208elRevert.addEventListener('click', () => {
209 editedSource = null;
210 editor.value = model.source;
211 void rebuild();
212});
214// The command reproduces this run's parameters on the desktop; keep it
215// selectable even where the clipboard API is unavailable.
216elCopyCmd.addEventListener('click', () => {
217 const text = elCmd.textContent ?? '';
218 const flash = (msg: string): void => {
219 elCopyCmd.textContent = msg;
220 setTimeout(() => (elCopyCmd.textContent = 'Copy'), 1200);
221 };
222 const selectCommand = (): void => {
223 const range = document.createRange();
224 range.selectNodeContents(elCmd);
225 const sel = getSelection();
226 sel?.removeAllRanges();
227 sel?.addRange(range);
228 flash('Selected');
229 };
230 if (!navigator.clipboard) return selectCommand();
231 navigator.clipboard.writeText(text).then(() => flash('Copied'), selectCommand);
232});
234// ---------------------------------------------------------------- setup
235function disposeView(): void {
236 for (const s of scenes) s.dispose();
237 scenes = [];
238 colorbars = [];
239 resizeObs?.disconnect();
240 resizeObs = null;
241 elPanels.replaceChildren();
242}
244/** Report a compile failure, and select the offending text in the editor. */
245function reportCompileError(e: unknown): void {
246 elErr.textContent = formatFailure(e, source());
247 elCompiled.textContent = '';
248 if (e instanceof ModelCompileError && e.start !== undefined) {
249 editor.select(e.start, e.end ?? e.start);
250 }
251}
253async function rebuild(): Promise<void> {
254 generation++;
255 const gen = generation;
256 setRunning(false);
257 disposeView();
258 session?.destroy();
259 session = null;
260 solverMs = 0;
261 frameMs = 0;
262 lastMeasure = 0;
263 elErr.textContent = '';
264 updateCommand();
265 if (!device) return;
267 try {
268 session = await ModelSession.create({
269 device,
270 model,
271 params,
272 lmax: Number(elLmax.value),
273 source: source(),
274 });
275 } catch (e) {
276 reportCompileError(e);
277 return;
278 }
279 if (gen !== generation) return;
281 session.seed(seed);
283 const plan = session.describe();
284 elCompiled.textContent =
285 `one step compiled to ${plan.step.length} GPU operations:\n` +
286 plan.step.map((l) => ` ${l}`).join('\n');
287 elRecompile.textContent = 'Recompile';
289 // mesh + scenes
290 const { nphi } = session.cfg;
291 const phi = new Float64Array(nphi);
292 for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
293 topo = buildTopology(session.sht.cosTheta, phi);
295 const sphereBg = getComputedStyle(document.documentElement)
296 .getPropertyValue('--sphere-bg')
297 .trim();
298 for (let k = 0; k < model.species.length; k++) {
299 const panel = document.createElement('div');
300 panel.className = 'panel';
301 const box = document.createElement('div');
302 box.className = 'sphere-box';
303 const tag = document.createElement('div');
304 tag.className = 'species-tag';
305 tag.textContent = model.species[k];
306 box.append(tag);
307 const side = document.createElement('div');
308 panel.append(box, side);
309 elPanels.append(panel);
311 const scene = new SphereScene(
312 box,
313 topo.numVertices,
314 topo.indices,
315 topo.sphereRef,
316 sphereBg || undefined,
317 );
318 scene.fitCamera();
319 scenes.push(scene);
320 colorbars.push(new Colorbar(side));
321 valueBufs[k] = new Float32Array(topo.numVertices);
322 colorBufs[k] = new Float32Array(topo.numVertices * 3);
323 ranges[k] = { lo: NaN, hi: NaN };
324 }
325 for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
327 resizeObs = new ResizeObserver(() => {
328 const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
329 boxes.forEach((box, i) => {
330 scenes[i]?.resize(box.clientWidth, box.clientHeight);
331 });
332 });
333 elPanels
334 .querySelectorAll<HTMLElement>('.sphere-box')
335 .forEach((box) => resizeObs!.observe(box));
337 await draw();
338 updateStats();
339 void pump();
340}
342async function reseed(): Promise<void> {
343 if (!session) return;
344 const gen = generation;
345 session.seed(seed);
346 if (gen !== generation) return;
347 for (const r of ranges) {
348 r.lo = NaN;
349 r.hi = NaN;
350 }
351 await draw();
352 updateStats();
353}
355// ---------------------------------------------------------------- drawing
356async function draw(): Promise<void> {
357 if (!session || !topo) return;
358 const gen = generation;
359 const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
360 for (let k = 0; k < model.species.length; k++) {
361 // The one readback per frame — the loop is otherwise entirely on the GPU.
362 // A rebuild can land while this is in flight and destroy the buffer being
363 // mapped, which rejects the map; that result is stale anyway, so drop it.
364 let field: Float32Array;
365 try {
366 field = await session.read(model.species[k]);
367 } catch (e) {
368 if (gen !== generation) return;
369 throw e;
370 }
371 if (gen !== generation || !topo) return;
372 fillFieldValues(valueBufs[k], field, topo);
373 let lo = Infinity;
374 let hi = -Infinity;
375 for (const v of valueBufs[k]) {
376 if (v < lo) lo = v;
377 if (v > hi) hi = v;
378 }
379 // smooth the color range in both directions so the shading evolves
380 // gently as the pattern grows (out-of-range values clamp meanwhile)
381 const r = ranges[k];
382 if (!Number.isFinite(r.lo)) {
383 r.lo = lo;
384 r.hi = hi;
385 } else {
386 const a = 0.15;
387 r.lo += a * (lo - r.lo);
388 r.hi += a * (hi - r.hi);
389 }
390 if (r.hi - r.lo < 1e-9) {
391 const mid = (r.hi + r.lo) / 2;
392 r.lo = mid - 5e-10;
393 r.hi = mid + 5e-10;
394 }
395 fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);
396 scenes[k]?.updateColors(colorBufs[k]);
397 colorbars[k]?.update(cmap, r.lo, r.hi);
398 }
399}
401function updateStats(): void {
402 if (!session) return;
403 const { nlat, nphi } = session.cfg;
404 const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
405 const solver =
406 solverMs > 0
407 ? `<b>${solverMs.toFixed(2)} ms/step</b> (${(1000 / solverMs).toFixed(0)} steps/s, ` +
408 `batch of ${MEASURE_BURST}, no readback)`
409 : '—';
410 const frame =
411 frameMs > 0
412 ? `${frameMs.toFixed(1)} ms/frame (${STEPS_PER_FRAME} steps + readback + render)`
413 : '—';
414 elStats.innerHTML =
415 `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
416 `${session.sht.fourierMode.toUpperCase()} · solver ${solver} · ${frame} · ` +
417 `t = <b>${session.t.toFixed(2)}</b> (${session.steps} steps)`;
418}
420// ---------------------------------------------------------------- sim loop
421const nextFrame = () => new Promise<number>(requestAnimationFrame);
423async function pump(): Promise<void> {
424 if (pumping) return;
425 pumping = true;
426 const gen = generation;
427 try {
428 while (running && session && gen === generation) {
429 // Occasionally, a burst purely to measure the solver rate: many steps,
430 // one sync, nothing read back — directly comparable to the desktop
431 // benchmark's throughput number.
432 if (performance.now() - lastMeasure > MEASURE_EVERY_MS) {
433 const m0 = performance.now();
434 session.step(MEASURE_BURST);
435 await session.sync();
436 if (gen !== generation) break;
437 solverMs = (performance.now() - m0) / MEASURE_BURST;
438 lastMeasure = performance.now();
439 }
441 // The frame itself. No explicit sync here — draw()'s readback already
442 // waits for the steps, so asking twice would only add a round trip.
443 const t0 = performance.now();
444 session.step(STEPS_PER_FRAME);
445 await draw();
446 if (gen !== generation) break;
447 frameMs = frameMs === 0
448 ? performance.now() - t0
449 : frameMs + 0.05 * (performance.now() - t0 - frameMs);
450 updateStats();
451 await nextFrame();
452 }
453 if (gen === generation) {
454 await draw();
455 updateStats();
456 }
457 } finally {
458 pumping = false;
459 }
460}
462/**
463 * Sustained solver benchmark, in the page.
464 *
465 * The same measurement `npm run bench` makes: batches of steps submitted
466 * together, waited for, never read back, with no rendering and no animation
467 * pacing in between. That makes it directly comparable to the terminal number,
468 * which is the only way to tell a genuinely slower browser GPU stack apart from
469 * the costs the app adds on top.
470 *
471 * It also reports the ramp — the first third of the run against the last. GPUs
472 * downclock when idle, and an animation-paced loop leaves them idle most of every
473 * frame, so a large ramp means the app's steady-state number is limited by clocks
474 * rather than by the work.
475 *
476 * These are ordinary steps: the simulation advances by them.
477 */
478async function benchmark(): Promise<void> {
479 if (!session) return;
480 setRunning(false);
481 const BATCH = 32;
482 const DURATION_MS = 2000;
483 elBenchResult.textContent = 'benchmarking…';
484 await nextFrame();
486 const gen = generation;
487 const perStep: number[] = [];
488 const t0 = performance.now();
489 while (performance.now() - t0 < DURATION_MS) {
490 const b0 = performance.now();
491 session.step(BATCH);
492 await session.sync();
493 if (gen !== generation) return;
494 perStep.push((performance.now() - b0) / BATCH);
495 }
497 const mean = (xs: number[]): number => xs.reduce((a, b) => a + b, 0) / xs.length;
498 const all = mean(perStep);
499 const best = Math.min(...perStep);
500 const third = Math.max(1, Math.floor(perStep.length / 3));
501 const first = mean(perStep.slice(0, third));
502 const last = mean(perStep.slice(-third));
503 const steps = perStep.length * BATCH;
505 elBenchResult.innerHTML =
506 `sustained solver: <b>${all.toFixed(2)} ms/step</b> ` +
507 `(${(1000 / all).toFixed(0)} steps/s) · best ${best.toFixed(2)} · ` +
508 `ramp ${(first / last).toFixed(2)}× (${first.toFixed(2)} → ${last.toFixed(2)}) · ` +
509 `${steps} steps in batches of ${BATCH} · ` +
510 `compare with <code>npm run bench -- --lmax ${session.cfg.lmax}</code>`;
511 await draw();
512 updateStats();
513}
515// ---------------------------------------------------------------- boot
516async function boot(): Promise<void> {
517 elModel.value = presets[0].key;
518 applyPreset(presets[0].key);
519 try {
520 device = await requestShtDevice();
521 adapterName = await describeAdapter(device);
522 } catch (e) {
523 device = null;
524 elErr.textContent =
525 `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
526 `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +
527 `needs a WebGPU-capable browser (Chrome/Edge 113+).`;
528 return;
529 }
530 device.lost.then((info) => {
531 if (info.reason !== 'destroyed') {
532 elErr.textContent = `WebGPU device lost: ${info.message}`;
533 }
534 });
535 await rebuild();
536}
538void boot();