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// ---------------------------------------------------------------- state
78let device: GPUDevice | null = null;
79let session: ModelSession | null = null;
80let topo: SphereMeshTopology | null = null;
81let scenes: SphereScene[] = [];
82let colorbars: Colorbar[] = [];
83let valueBufs: Float32Array[] = [];
84let colorBufs: Float32Array[] = [];
85let ranges: { lo: number; hi: number }[] = [];
86let resizeObs: ResizeObserver | null = null;
88const initial = resolvePreset(presets[0].key);
89let model: MModel = mModelByKey(initial.model.key)!;
90let params: Params = initial.params;
91/** The .m as edited in the page; `null` while it matches the file. */
92let editedSource: string | null = null;
93let seed = 1;
94let running = false;
95let adapterName = '';
96let pumping = false;
97let solverMs = 0;
98let frameMs = 0;
99let generation = 0; // bumped on every rebuild to cancel stale pumps
101const source = (): string => editedSource ?? model.source;
103// ---------------------------------------------------------------- UI wiring
104function buildParamInputs(): void {
105 elParams.replaceChildren();
106 for (const spec of model.params) {
107 const label = document.createElement('label');
108 label.textContent = `${spec.label} `;
109 const input = document.createElement('input');
110 input.type = 'number';
111 input.min = String(spec.min);
112 input.max = String(spec.max);
113 input.step = String(spec.step);
114 input.value = String(params[spec.key]);
115 input.addEventListener('change', () => {
116 const v = Number(input.value);
117 if (Number.isFinite(v)) params[spec.key] = v;
118 // Parameters are uniforms, not constants baked into the kernels, so a
119 // change costs an upload rather than a recompile.
120 session?.setParams(params);
121 updateCommand();
122 });
123 label.append(input);
124 elParams.append(label);
125 }
126}
128function applyPreset(presetKey: string): void {
129 const resolved = resolvePreset(presetKey);
130 const next = mModelByKey(resolved.model.key);
131 if (!next) {
132 elErr.textContent = `No .m model for '${resolved.model.key}'`;
133 return;
134 }
135 model = next;
136 params = resolved.params;
137 editedSource = null;
138 editor.value = model.source;
139 elEditorTitle.textContent =
140 `models/${model.key}.m — init() and step(), compiled to WebGPU`;
141 buildParamInputs();
142 elBlurb.textContent = model.blurb;
143 updateCommand();
144}
146/** The run currently on screen, as the benchmark's RunSpec. */
147function currentSpec(): RunSpec {
148 return {
149 preset: elModel.value,
150 lmax: Number(elLmax.value),
151 seed,
152 steps: DEFAULT_STEPS,
153 warmup: DEFAULT_WARMUP,
154 params,
155 };
156}
158function updateCommand(): void {
159 elCmd.textContent = formatCommand(currentSpec());
160}
162elModel.addEventListener('change', () => {
163 applyPreset(elModel.value);
164 void rebuild();
165});
166elLmax.addEventListener('change', () => void rebuild());
167elColormap.addEventListener('change', () => void draw());
169function setRunning(next: boolean): void {
170 running = next;
171 elRunPause.textContent = running ? 'Pause' : 'Run';
172 if (running) void pump();
173}
175elRunPause.addEventListener('click', () => setRunning(!running));
176elReseed.addEventListener('click', () => {
177 seed = (Math.random() * 2 ** 31) >>> 0;
178 setRunning(false);
179 updateCommand();
180 void reseed();
181});
182elResetView.addEventListener('click', () => {
183 for (const s of scenes) s.resetCamera();
184});
186elRecompile.addEventListener('click', () => {
187 editedSource = editor.value;
188 void rebuild();
189});
190elRevert.addEventListener('click', () => {
191 editedSource = null;
192 editor.value = model.source;
193 void rebuild();
194});
196// The command reproduces this run's parameters on the desktop; keep it
197// selectable even where the clipboard API is unavailable.
198elCopyCmd.addEventListener('click', () => {
199 const text = elCmd.textContent ?? '';
200 const flash = (msg: string): void => {
201 elCopyCmd.textContent = msg;
202 setTimeout(() => (elCopyCmd.textContent = 'Copy'), 1200);
203 };
204 const selectCommand = (): void => {
205 const range = document.createRange();
206 range.selectNodeContents(elCmd);
207 const sel = getSelection();
208 sel?.removeAllRanges();
209 sel?.addRange(range);
210 flash('Selected');
211 };
212 if (!navigator.clipboard) return selectCommand();
213 navigator.clipboard.writeText(text).then(() => flash('Copied'), selectCommand);
214});
216// ---------------------------------------------------------------- setup
217function disposeView(): void {
218 for (const s of scenes) s.dispose();
219 scenes = [];
220 colorbars = [];
221 resizeObs?.disconnect();
222 resizeObs = null;
223 elPanels.replaceChildren();
224}
226/** Report a compile failure, and select the offending text in the editor. */
227function reportCompileError(e: unknown): void {
228 elErr.textContent = formatFailure(e, source());
229 elCompiled.textContent = '';
230 if (e instanceof ModelCompileError && e.start !== undefined) {
231 editor.select(e.start, e.end ?? e.start);
232 }
233}
235async function rebuild(): Promise<void> {
236 generation++;
237 const gen = generation;
238 setRunning(false);
239 disposeView();
240 session?.destroy();
241 session = null;
242 solverMs = 0;
243 frameMs = 0;
244 elErr.textContent = '';
245 updateCommand();
246 if (!device) return;
248 try {
249 session = await ModelSession.create({
250 device,
251 model,
252 params,
253 lmax: Number(elLmax.value),
254 source: source(),
255 });
256 } catch (e) {
257 reportCompileError(e);
258 return;
259 }
260 if (gen !== generation) return;
262 session.seed(seed);
264 const plan = session.describe();
265 elCompiled.textContent =
266 `one step compiled to ${plan.step.length} GPU operations:\n` +
267 plan.step.map((l) => ` ${l}`).join('\n');
268 elRecompile.textContent = 'Recompile';
270 // mesh + scenes
271 const { nphi } = session.cfg;
272 const phi = new Float64Array(nphi);
273 for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
274 topo = buildTopology(session.sht.cosTheta, phi);
276 const sphereBg = getComputedStyle(document.documentElement)
277 .getPropertyValue('--sphere-bg')
278 .trim();
279 for (let k = 0; k < model.species.length; k++) {
280 const panel = document.createElement('div');
281 panel.className = 'panel';
282 const box = document.createElement('div');
283 box.className = 'sphere-box';
284 const tag = document.createElement('div');
285 tag.className = 'species-tag';
286 tag.textContent = model.species[k];
287 box.append(tag);
288 const side = document.createElement('div');
289 panel.append(box, side);
290 elPanels.append(panel);
292 const scene = new SphereScene(
293 box,
294 topo.numVertices,
295 topo.indices,
296 topo.sphereRef,
297 sphereBg || undefined,
298 );
299 scene.fitCamera();
300 scenes.push(scene);
301 colorbars.push(new Colorbar(side));
302 valueBufs[k] = new Float32Array(topo.numVertices);
303 colorBufs[k] = new Float32Array(topo.numVertices * 3);
304 ranges[k] = { lo: NaN, hi: NaN };
305 }
306 for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
308 resizeObs = new ResizeObserver(() => {
309 const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
310 boxes.forEach((box, i) => {
311 scenes[i]?.resize(box.clientWidth, box.clientHeight);
312 });
313 });
314 elPanels
315 .querySelectorAll<HTMLElement>('.sphere-box')
316 .forEach((box) => resizeObs!.observe(box));
318 await draw();
319 updateStats();
320 void pump();
321}
323async function reseed(): Promise<void> {
324 if (!session) return;
325 const gen = generation;
326 session.seed(seed);
327 if (gen !== generation) return;
328 for (const r of ranges) {
329 r.lo = NaN;
330 r.hi = NaN;
331 }
332 await draw();
333 updateStats();
334}
336// ---------------------------------------------------------------- drawing
337async function draw(): Promise<void> {
338 if (!session || !topo) return;
339 const gen = generation;
340 const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
341 for (let k = 0; k < model.species.length; k++) {
342 // The one readback per frame — the loop is otherwise entirely on the GPU.
343 // A rebuild can land while this is in flight and destroy the buffer being
344 // mapped, which rejects the map; that result is stale anyway, so drop it.
345 let field: Float32Array;
346 try {
347 field = await session.read(model.species[k]);
348 } catch (e) {
349 if (gen !== generation) return;
350 throw e;
351 }
352 if (gen !== generation || !topo) return;
353 fillFieldValues(valueBufs[k], field, topo);
354 let lo = Infinity;
355 let hi = -Infinity;
356 for (const v of valueBufs[k]) {
357 if (v < lo) lo = v;
358 if (v > hi) hi = v;
359 }
360 // smooth the color range in both directions so the shading evolves
361 // gently as the pattern grows (out-of-range values clamp meanwhile)
362 const r = ranges[k];
363 if (!Number.isFinite(r.lo)) {
364 r.lo = lo;
365 r.hi = hi;
366 } else {
367 const a = 0.15;
368 r.lo += a * (lo - r.lo);
369 r.hi += a * (hi - r.hi);
370 }
371 if (r.hi - r.lo < 1e-9) {
372 const mid = (r.hi + r.lo) / 2;
373 r.lo = mid - 5e-10;
374 r.hi = mid + 5e-10;
375 }
376 fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);
377 scenes[k]?.updateColors(colorBufs[k]);
378 colorbars[k]?.update(cmap, r.lo, r.hi);
379 }
380}
382function updateStats(): void {
383 if (!session) return;
384 const { nlat, nphi } = session.cfg;
385 const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
386 const solver =
387 solverMs > 0
388 ? `<b>${solverMs.toFixed(2)} ms/step</b> (${(1000 / solverMs).toFixed(0)} steps/s)`
389 : '—';
390 const frame =
391 frameMs > 0
392 ? `${frameMs.toFixed(1)} ms/frame incl. readback + render`
393 : '—';
394 elStats.innerHTML =
395 `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
396 `${session.sht.fourierMode.toUpperCase()} · solver ${solver} · ${frame} · ` +
397 `t = <b>${session.t.toFixed(2)}</b> (${session.steps} steps)`;
398}
400// ---------------------------------------------------------------- sim loop
401const nextFrame = () => new Promise<number>(requestAnimationFrame);
403async function pump(): Promise<void> {
404 if (pumping) return;
405 pumping = true;
406 const gen = generation;
407 try {
408 while (running && session && gen === generation) {
409 // Two separate costs, kept separate. The solver is the batch of steps
410 // alone, waited for but not read back — the number the desktop benchmark
411 // reports, and the only one comparable to it. The frame additionally
412 // carries a GPU->CPU readback per species (which in a browser crosses a
413 // process boundary), the colormapping, and the three.js upload, and those
414 // can easily cost more than the steps do.
415 const t0 = performance.now();
416 session.step(STEPS_PER_FRAME);
417 await session.sync();
418 if (gen !== generation) break;
419 const tSolver = performance.now();
420 await draw();
421 if (gen !== generation) break;
422 const ema = (prev: number, next: number): number =>
423 prev === 0 ? next : prev + 0.05 * (next - prev);
424 solverMs = ema(solverMs, (tSolver - t0) / STEPS_PER_FRAME);
425 frameMs = ema(frameMs, performance.now() - t0);
426 updateStats();
427 await nextFrame();
428 }
429 if (gen === generation) {
430 await draw();
431 updateStats();
432 }
433 } finally {
434 pumping = false;
435 }
436}
438// ---------------------------------------------------------------- boot
439async function boot(): Promise<void> {
440 elModel.value = presets[0].key;
441 applyPreset(presets[0].key);
442 try {
443 device = await requestShtDevice();
444 adapterName = await describeAdapter(device);
445 } catch (e) {
446 device = null;
447 elErr.textContent =
448 `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
449 `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +
450 `needs a WebGPU-capable browser (Chrome/Edge 113+).`;
451 return;
452 }
453 device.lost.then((info) => {
454 if (info.reason !== 'destroyed') {
455 elErr.textContent = `WebGPU device lost: ${info.message}`;
456 }
457 });
458 await rebuild();
459}
461void boot();