src/main.ts at 35d91fae9dbbc96029de1a0a74b776de1bf16e96

concept-collection / turing-sphere

443 lines · 13.4 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// ---------------------------------------------------------------- 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 stepMs = 0;

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

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

101

102// ---------------------------------------------------------------- UI wiring

103function buildParamInputs(): void {

104 elParams.replaceChildren();

105 for (const spec of model.params) {

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

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

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

109 input.type = 'number';

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

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

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

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

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

115 const v = Number(input.value);

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

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

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

119 session?.setParams(params);

120 updateCommand();

121 });

122 label.append(input);

123 elParams.append(label);

124 }

125}

126

127function applyPreset(presetKey: string): void {

128 const resolved = resolvePreset(presetKey);

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

130 if (!next) {

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

132 return;

133 }

134 model = next;

135 params = resolved.params;

136 editedSource = null;

137 editor.value = model.source;

138 elEditorTitle.textContent =

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

140 buildParamInputs();

141 elBlurb.textContent = model.blurb;

142 updateCommand();

143}

144

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

146function currentSpec(): RunSpec {

147 return {

148 preset: elModel.value,

149 lmax: Number(elLmax.value),

150 seed,

151 steps: DEFAULT_STEPS,

152 warmup: DEFAULT_WARMUP,

153 params,

154 };

155}

156

157function updateCommand(): void {

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

159}

160

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

162 applyPreset(elModel.value);

163 void rebuild();

164});

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

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

167

168function setRunning(next: boolean): void {

169 running = next;

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

171 if (running) void pump();

172}

173

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

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

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

177 setRunning(false);

178 updateCommand();

179 void reseed();

180});

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

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

183});

184

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

186 editedSource = editor.value;

187 void rebuild();

188});

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

190 editedSource = null;

191 editor.value = model.source;

192 void rebuild();

193});

194

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

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

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

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

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

200 elCopyCmd.textContent = msg;

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

202 };

203 const selectCommand = (): void => {

204 const range = document.createRange();

205 range.selectNodeContents(elCmd);

206 const sel = getSelection();

207 sel?.removeAllRanges();

208 sel?.addRange(range);

209 flash('Selected');

210 };

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

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

213});

214

215// ---------------------------------------------------------------- setup

216function disposeView(): void {

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

218 scenes = [];

219 colorbars = [];

220 resizeObs?.disconnect();

221 resizeObs = null;

222 elPanels.replaceChildren();

223}

224

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

226function reportCompileError(e: unknown): void {

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

228 elCompiled.textContent = '';

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

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

231 }

232}

233

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

235 generation++;

236 const gen = generation;

237 setRunning(false);

238 disposeView();

239 session?.destroy();

240 session = null;

241 stepMs = 0;

242 elErr.textContent = '';

243 updateCommand();

244 if (!device) return;

245

246 try {

247 session = await ModelSession.create({

248 device,

249 model,

250 params,

251 lmax: Number(elLmax.value),

252 source: source(),

253 });

254 } catch (e) {

255 reportCompileError(e);

256 return;

257 }

258 if (gen !== generation) return;

259

260 session.seed(seed);

261

262 const plan = session.describe();

263 elCompiled.textContent =

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

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

266 elRecompile.textContent = 'Recompile';

267

268 // mesh + scenes

269 const { nphi } = session.cfg;

270 const phi = new Float64Array(nphi);

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

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

273

274 const sphereBg = getComputedStyle(document.documentElement)

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

276 .trim();

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

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

279 panel.className = 'panel';

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

281 box.className = 'sphere-box';

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

283 tag.className = 'species-tag';

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

285 box.append(tag);

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

287 panel.append(box, side);

288 elPanels.append(panel);

289

290 const scene = new SphereScene(

291 box,

292 topo.numVertices,

293 topo.indices,

294 topo.sphereRef,

295 sphereBg || undefined,

296 );

297 scene.fitCamera();

298 scenes.push(scene);

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

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

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

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

303 }

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

305

306 resizeObs = new ResizeObserver(() => {

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

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

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

310 });

311 });

312 elPanels

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

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

315

316 await draw();

317 updateStats();

318 void pump();

319}

320

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

322 if (!session) return;

323 const gen = generation;

324 session.seed(seed);

325 if (gen !== generation) return;

326 for (const r of ranges) {

327 r.lo = NaN;

328 r.hi = NaN;

329 }

330 await draw();

331 updateStats();

332}

333

334// ---------------------------------------------------------------- drawing

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

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

337 const gen = generation;

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

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

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

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

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

343 let field: Float32Array;

344 try {

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

346 } catch (e) {

347 if (gen !== generation) return;

348 throw e;

349 }

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

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

352 let lo = Infinity;

353 let hi = -Infinity;

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

355 if (v < lo) lo = v;

356 if (v > hi) hi = v;

357 }

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

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

360 const r = ranges[k];

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

362 r.lo = lo;

363 r.hi = hi;

364 } else {

365 const a = 0.15;

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

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

368 }

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

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

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

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

373 }

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

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

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

377 }

378}

379

380function updateStats(): void {

381 if (!session) return;

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

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

384 const rate = stepMs > 0 ? `${(1000 / stepMs).toFixed(1)} steps/s` : '—';

385 elStats.innerHTML =

386 `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +

387 `${stepMs > 0 ? stepMs.toFixed(1) : '—'} ms/step · ${rate} · ` +

388 `t = <b>${session.t.toFixed(2)}</b> (${session.steps} steps)`;

389}

390

391// ---------------------------------------------------------------- sim loop

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

393

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

395 if (pumping) return;

396 pumping = true;

397 const gen = generation;

398 try {

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

400 const t0 = performance.now();

401 session.step(STEPS_PER_FRAME);

402 // draw() awaits the readback, which also waits for the batch to finish,

403 // so this measures the real end-to-end cost per step.

404 await draw();

405 if (gen !== generation) break;

406 const dtMs = (performance.now() - t0) / STEPS_PER_FRAME;

407 stepMs = stepMs === 0 ? dtMs : stepMs + 0.05 * (dtMs - stepMs);

408 updateStats();

409 await nextFrame();

410 }

411 if (gen === generation) {

412 await draw();

413 updateStats();

414 }

415 } finally {

416 pumping = false;

417 }

418}

419

420// ---------------------------------------------------------------- boot

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

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

423 applyPreset(presets[0].key);

424 try {

425 device = await requestShtDevice();

426 adapterName = await describeAdapter(device);

427 } catch (e) {

428 device = null;

429 elErr.textContent =

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

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

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

433 return;

434 }

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

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

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

438 }

439 });

440 await rebuild();

441}

442

443void boot();