/ concept-collection / math-webgpu-sandbox
Sign in
concept-collection / math-webgpu-sandbox
math-webgpu-sandbox / src / main.ts
220 lines · 7.6 KBBlameHistoryRaw
1/**
2 * The sandbox page: editor on the left, output on the right, a timing table
3 * underneath that lines up GPU tic/toc segments with the CPU run's.
4 */
5import { CodeEditor } from './editor/codeEditor.ts';
6import { requestSandboxDevice, runScript, type SandboxGpu } from './mgpu/session.ts';
7import { formatFailure } from './mgpu/errors.ts';
8import type { CpuRunHandle } from './cpu/cpuRunner.ts';
9import fusedLoop from '../examples/fused_loop.m?raw';
10import matmul from '../examples/matmul.m?raw';
11import monteCarloPi from '../examples/monte_carlo_pi.m?raw';
12import logisticEnsemble from '../examples/logistic_ensemble.m?raw';
13import reductions from '../examples/reductions.m?raw';
15const EXAMPLES: { key: string; title: string; source: string }[] = [
16 { key: 'fused_loop', title: 'fused elementwise loop', source: fusedLoop },
17 { key: 'matmul', title: 'matrix multiply (GEMM)', source: matmul },
18 { key: 'monte_carlo_pi', title: 'Monte Carlo pi (masks)', source: monteCarloPi },
19 { key: 'logistic_ensemble', title: 'logistic map ensemble', source: logisticEnsemble },
20 { key: 'reductions', title: 'fused reductions', source: reductions },
21];
23const $ = <T extends HTMLElement>(id: string): T => {
24 const el = document.getElementById(id);
25 if (!el) throw new Error(`missing #${id}`);
26 return el as T;
27};
29const consoleEl = $<HTMLPreElement>('console');
30const cpuConsoleEl = $<HTMLPreElement>('cpuconsole');
31const runBtn = $<HTMLButtonElement>('run');
32const runCpuBtn = $<HTMLButtonElement>('runcpu');
33const copyBtn = $<HTMLButtonElement>('copy');
34const exampleSel = $<HTMLSelectElement>('example');
35const timingsEl = $<HTMLTableElement>('timings');
36const planDetails = $<HTMLElement>('plandetails');
37const planEl = $<HTMLPreElement>('plan');
38const runState = $<HTMLElement>('runstate');
39const deviceEl = $<HTMLElement>('device');
40const gpuWarn = $<HTMLElement>('gpuwarn');
42const editor = new CodeEditor({
43 textarea: $<HTMLTextAreaElement>('source'),
44 overlay: $('highlight'),
45 onInput: () => {
46 $('editstate').textContent = '';
47 },
48});
50for (const ex of EXAMPLES) {
51 const opt = document.createElement('option');
52 opt.value = ex.key;
53 opt.textContent = ex.title;
54 exampleSel.appendChild(opt);
56exampleSel.addEventListener('change', () => {
57 const ex = EXAMPLES.find((e) => e.key === exampleSel.value);
58 if (ex) editor.value = ex.source;
59});
60editor.value = EXAMPLES[0].source;
62const append = (el: HTMLElement, text: string, cls?: string): void => {
63 if (cls) {
64 const span = document.createElement('span');
65 span.className = cls;
66 span.textContent = text;
67 el.appendChild(span);
68 } else {
69 el.appendChild(document.createTextNode(text));
70 }
71 el.scrollTop = el.scrollHeight;
72};
74/** Latest timing results, merged into the table by tic..toc pair. */
75const lastSeconds: { gpu: number[]; cpu: number[] } = { gpu: [], cpu: [] };
76let gpuTotal: number | null = null;
77let cpuTotal: number | null = null;
79function renderTimings(): void {
80 const n = Math.max(lastSeconds.gpu.length, lastSeconds.cpu.length);
81 if (n === 0 && gpuTotal === null && cpuTotal === null) {
82 timingsEl.style.display = 'none';
83 return;
84 }
85 const fmt = (s: number | undefined | null): string =>
86 s === undefined || s === null ? '—' : s >= 0.1 ? `${s.toFixed(3)} s` : `${(s * 1000).toFixed(2)} ms`;
87 const rows: string[] = [
88 '<tr><th></th><th>GPU</th><th>CPU (numbl)</th><th>CPU / GPU</th></tr>',
89 ];
90 for (let i = 0; i < n; i++) {
91 const g = lastSeconds.gpu[i];
92 const c = lastSeconds.cpu[i];
93 const ratio = g !== undefined && c !== undefined ? `${(c / g).toFixed(1)}×` : '—';
94 rows.push(
95 `<tr><td>tic…toc #${i + 1}</td><td>${fmt(g)}</td><td>${fmt(c)}</td><td>${ratio}</td></tr>`,
96 );
97 }
98 const totalRatio =
99 gpuTotal !== null && cpuTotal !== null ? `${(cpuTotal / gpuTotal).toFixed(1)}×` : '—';
100 rows.push(
101 `<tr><td>whole script</td><td>${fmt(gpuTotal)}</td><td>${fmt(cpuTotal)}</td><td>${totalRatio}</td></tr>`,
102 );
103 timingsEl.innerHTML = rows.join('');
104 timingsEl.style.display = 'table';
107let gpu: SandboxGpu | null = null;
108(async () => {
109 try {
110 gpu = await requestSandboxDevice();
111 deviceEl.textContent = `GPU: ${gpu.description}`;
112 gpu.device.lost.then((info) => {
113 gpu = null;
114 gpuWarn.hidden = false;
115 gpuWarn.textContent = `The GPU device was lost (${info.message}); reload the page.`;
116 runBtn.disabled = true;
117 });
118 } catch (e) {
119 gpuWarn.hidden = false;
120 gpuWarn.innerHTML =
121 `<b>No WebGPU here.</b> ${e instanceof Error ? e.message : String(e)} — ` +
122 `Chrome/Edge have it on by default; Firefox and Safari are rolling it out. ` +
123 `The CPU run still works.`;
124 runBtn.disabled = true;
125 }
126})();
128runBtn.addEventListener('click', () => {
129 void runGpu();
130});
131async function runGpu(): Promise<void> {
132 if (!gpu) return;
133 runBtn.disabled = true;
134 runState.textContent = 'compiling…';
135 consoleEl.textContent = '';
136 lastSeconds.gpu = [];
137 gpuTotal = null;
138 renderTimings();
139 planDetails.hidden = true;
140 const source = editor.value;
141 try {
142 const t0 = performance.now();
143 const run = await runScript(gpu.device, source, (text) => {
144 runState.textContent = 'running…';
145 append(consoleEl, text);
146 });
147 void t0;
148 if (run.result.error) {
149 append(consoleEl, `\n${run.result.error}\n`, 'err');
150 }
151 append(
152 consoleEl,
153 `\n[compile ${(run.compileSeconds * 1000).toFixed(0)} ms · ` +
154 `run ${run.result.totalSeconds.toFixed(3)} s · f32 · ${gpu.description}]\n`,
155 'meta',
156 );
157 lastSeconds.gpu = run.result.segments.map((s) => s.seconds);
158 gpuTotal = run.result.totalSeconds;
159 planEl.textContent = run.planDescription.join('\n') || '(no GPU ops)';
160 planDetails.hidden = false;
161 renderTimings();
162 runState.textContent = '';
163 } catch (e) {
164 append(consoleEl, formatFailure(e, source) + '\n', 'err');
165 runState.textContent = '';
166 } finally {
167 runBtn.disabled = !gpu;
168 }
171let cpuHandle: CpuRunHandle | null = null;
172runCpuBtn.addEventListener('click', () => {
173 if (cpuHandle) {
174 cpuHandle.cancel();
175 cpuHandle = null;
176 runCpuBtn.textContent = 'Run on CPU (numbl)';
177 return;
178 }
179 void runCpu();
180});
181async function runCpu(): Promise<void> {
182 cpuConsoleEl.style.display = 'block';
183 cpuConsoleEl.textContent = '';
184 append(cpuConsoleEl, '[CPU · numbl engine in a worker · f64]\n', 'meta');
185 lastSeconds.cpu = [];
186 cpuTotal = null;
187 renderTimings();
188 runCpuBtn.textContent = 'Stop CPU run';
189 // numbl's browser engine is a couple of megabytes; load it on first use.
190 const { runOnCpu } = await import('./cpu/cpuRunner.ts');
191 const handle = runOnCpu(editor.value, (text) => append(cpuConsoleEl, text));
192 cpuHandle = handle;
193 try {
194 const res = await handle.result;
195 if (res.error) append(cpuConsoleEl, `\n${res.error}\n`, 'err');
196 if (res.aborted) append(cpuConsoleEl, `\n[stopped]\n`, 'meta');
197 else {
198 append(cpuConsoleEl, `[total ${res.totalSeconds.toFixed(3)} s]\n`, 'meta');
199 // Line the CPU's tic..toc pairs up with the GPU's by order.
200 lastSeconds.cpu = [...res.output.matchAll(/Elapsed time is ([0-9.eE+-]+) seconds/g)]
201 .map((m) => Number(m[1]));
202 cpuTotal = res.totalSeconds;
203 renderTimings();
204 }
205 } catch (e) {
206 append(cpuConsoleEl, `${e instanceof Error ? e.message : String(e)}\n`, 'err');
207 } finally {
208 if (cpuHandle === handle) cpuHandle = null;
209 runCpuBtn.textContent = 'Run on CPU (numbl)';
210 }
213copyBtn.addEventListener('click', () => {
214 void navigator.clipboard.writeText(editor.value).then(() => {
215 copyBtn.textContent = 'Copied ✓';
216 setTimeout(() => {
217 copyBtn.textContent = 'Copy script';
218 }, 1200);
219 });
220});
moveopenescclose