concept-collection / math-webgpu-sandbox
Show the generated WGSL kernels in the compiled-plan pane
Jeremy Magland <jeremy.magland@gmail.com> committed commit 4ded6ff8608e parent 9c655b8 Browse files
4 changed files+44−3
index.htmlmodified+9−1View file
@@ -121,6 +121,13 @@
121121 background: var(--pane-bg); border: 1px solid var(--line); border-radius: 6px;
122122 font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
123123 }
124+ #kernels details { margin: 6px 0 0 14px; }
125+ #kernels summary {
126+ cursor: pointer; color: var(--ink-2);
127+ font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
128+ }
129+ #kernels summary b { color: var(--accent); font-weight: 600; }
130+ #kernels pre { max-height: 24em; }
124131 #blurb { margin-top: 14px; font-size: 13px; color: var(--ink-2); max-width: 72em; }
125132 #blurb code { font-size: 12px; }
126133 .tok-com { color: var(--tok-com); }
@@ -165,8 +172,9 @@
165172 </div>
166173 </div>
167174 <details id="plandetails" hidden>
168- <summary>compiled GPU plan</summary>
175+ <summary>compiled GPU plan &amp; kernels</summary>
169176 <pre id="plan"></pre>
177+ <div id="kernels"></div>
170178 </details>
171179 <p id="blurb">
172180 The GPU computes in f32 (WebGPU has no f64) โ€” for the closest MATLAB comparison,
src/main.tsmodified+13−0View file
@@ -32,6 +32,7 @@ const copyBtn = $<HTMLButtonElement>('copy');
3232 const exampleSel = $<HTMLSelectElement>('example');
3333 const planDetails = $<HTMLElement>('plandetails');
3434 const planEl = $<HTMLPreElement>('plan');
35+const kernelsEl = $<HTMLElement>('kernels');
3536 const runState = $<HTMLElement>('runstate');
3637 const deviceEl = $<HTMLElement>('device');
3738 const gpuWarn = $<HTMLElement>('gpuwarn');
@@ -112,6 +113,18 @@ async function runGpu(): Promise<void> {
112113 'meta',
113114 );
114115 planEl.textContent = run.planDescription.join('\n') || '(no GPU ops)';
116+ kernelsEl.textContent = '';
117+ for (const k of run.kernels) {
118+ const d = document.createElement('details');
119+ const summary = document.createElement('summary');
120+ const b = document.createElement('b');
121+ b.textContent = k.label;
122+ summary.append('wgsl ยท ', b);
123+ const pre = document.createElement('pre');
124+ pre.textContent = k.code.trim();
125+ d.append(summary, pre);
126+ kernelsEl.appendChild(d);
127+ }
115128 planDetails.hidden = false;
116129 runState.textContent = '';
117130 } catch (e) {
src/mgpu/plan.tsmodified+13−0View file
@@ -125,10 +125,20 @@ export type Op =
125125 shape: number[];
126126 };
127127
128+/** One generated compute shader, for the "show me the kernels" pane. */
129+export interface KernelSource {
130+ /** The label of the statement that first created it. */
131+ label: string;
132+ code: string;
133+}
134+
128135 export interface ScriptPlan {
129136 ops: Op[];
130137 /** Human-readable op sequence โ€” what the script actually compiled to. */
131138 describe(): string[];
139+ /** Every distinct WGSL kernel, in creation order (reused kernels appear
140+ * once, under their first label). */
141+ kernels: KernelSource[];
132142 destroy(): void;
133143 }
134144
@@ -153,6 +163,7 @@ export async function planScript(
153163 /** Loop-variable uniform buffers, by cVar. */
154164 const loopUniforms = new Map<string, GPUBuffer>();
155165 const pipelines = new Map<string, GPUComputePipeline>();
166+ const kernelSources: KernelSource[] = [];
156167 const describeLines: string[] = [];
157168
158169 let seedCounter = 1;
@@ -253,6 +264,7 @@ export async function planScript(
253264 const err = await device.popErrorScope();
254265 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
255266 pipelines.set(code, p);
267+ kernelSources.push({ label, code });
256268 return p;
257269 }
258270
@@ -1213,6 +1225,7 @@ export async function planScript(
12131225 return {
12141226 ops: opStack[0],
12151227 describe: () => describeLines,
1228+ kernels: kernelSources,
12161229 destroy: () => {
12171230 for (const b of owned) b.destroy();
12181231 owned.length = 0;
src/mgpu/session.tsmodified+9−2View file
@@ -5,7 +5,7 @@
55 * from execution โ€” the execution numbers are the ones to compare with MATLAB.
66 */
77 import { compileScript } from './compile.ts';
8-import { planScript, type ScriptPlan } from './plan.ts';
8+import { planScript, type KernelSource, type ScriptPlan } from './plan.ts';
99 import { executePlan, type RunResult } from './run.ts';
1010 import { asCompileError } from './errors.ts';
1111
@@ -13,6 +13,8 @@ export interface ScriptRun {
1313 compileSeconds: number;
1414 /** Human-readable op sequence, for the "what did this compile to" pane. */
1515 planDescription: string[];
16+ /** The WGSL of every distinct kernel the script compiled to. */
17+ kernels: KernelSource[];
1618 result: RunResult;
1719 }
1820
@@ -62,7 +64,12 @@ export async function runScript(
6264 const compileSeconds = (performance.now() - t0) / 1000;
6365 try {
6466 const result = await executePlan(device, plan, onOutput);
65- return { compileSeconds, planDescription: plan.describe(), result };
67+ return {
68+ compileSeconds,
69+ planDescription: plan.describe(),
70+ kernels: plan.kernels,
71+ result,
72+ };
6673 } finally {
6774 plan.destroy();
6875 }