2 * The generic runner: everything the host does with an instrument file's
3 * WGSL, with no knowledge of what the shaders compute.
4 *
5 * A phase (precompute or render) runs as: create its f32 storage buffers at
6 * the sizes the derived expressions give, upload any named inputs, build
7 * one compute pipeline per shader with the scope values the shader lists as
8 * override constants, run the setup passes, run the loop, and read back the
9 * requested buffers. Two conveniences the schedules lean on:
10 *
11 * - `__tick`: a u32 the host zeroes before the loop and bumps after every
12 * iteration, so a shader knows which timestep or tile it is in without
13 * any per-iteration upload.
14 * - `rotate`: buffer groups cycled one place per iteration, which is how
15 * a leapfrog's state histories advance without copies. With a group
16 * [p, pm, pn], the buffer a pass sees as `p` at iteration k is the one
17 * it saw as `pn` at iteration k-1.
18 *
19 * Iterations are encoded in batches with an await between them, which
20 * bounds how far the queue runs ahead and paces progress and cancellation.
21 */
22import { evalExpr } from '../expr.ts';
23import type { PassSpec, PhaseSpec, ShaderSpec } from '../format.ts';
25export const WORKGROUP_SIZE = 64;
27/** Workgroup grids over 65535 spill into y; shaders linearize back with
28 * `gid.x + gid.y * 4194240u` (65535 * 64, one x-row of threads). */
29const MAX_WG_X = 65535;
31export class PhaseCancelled extends Error {
32 constructor() {
33 super('phase cancelled');
34 }
35}
37export interface PhaseRunOptions {
38 device: GPUDevice;
39 shaders: Record<string, ShaderSpec>;
40 phase: PhaseSpec;
41 /** Fully evaluated scope: params, derived, and any runtime values. */
42 scope: Record<string, number>;
43 /** Uploaded into same-named declared buffers before setup. */
44 inputs?: Record<string, Float32Array>;
45 /** Buffer names to read back when the phase completes. */
46 outputs: string[];
47 label: string;
48 onProgress?: (done: number, total: number) => void;
49 cancelled?: () => boolean;
50}
52const bumpShader = `
53@group(0) @binding(0) var<storage, read_write> tick: array<u32>;
54@compute @workgroup_size(1)
55fn main() { tick[0] = tick[0] + 1u; }
56`;
58export async function runPhase(opts: PhaseRunOptions): Promise<Record<string, Float32Array>> {
59 const { device, shaders, phase, scope, label } = opts;
60 const owned: GPUBuffer[] = [];
61 const mk = (name: string, bytes: number, extra: GPUBufferUsageFlags = 0): GPUBuffer => {
62 const b = device.createBuffer({
63 label: `${label}-${name}`,
64 size: Math.max(4, bytes),
65 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST | extra,
66 });
67 owned.push(b);
68 return b;
69 };
71 try {
72 // Buffers. Sizes are expressions; contents start zeroed (WebGPU
73 // guarantees zero-initialized buffers), inputs overwrite by name.
74 const buffers = new Map<string, GPUBuffer>();
75 for (const b of phase.buffers) {
76 const count = evalExpr(b.count, scope);
77 if (!Number.isInteger(count) || count < 0) {
78 throw new Error(`buffer '${b.name}': count '${b.count}' = ${count} is not a whole number`);
79 }
80 buffers.set(b.name, mk(b.name, 4 * count));
81 }
82 buffers.set('__tick', mk('__tick', 4));
83 for (const [name, data] of Object.entries(opts.inputs ?? {})) {
84 const buf = buffers.get(name);
85 if (!buf) throw new Error(`input '${name}' has no declared buffer`);
86 if (data.byteLength > buf.size) {
87 throw new Error(`input '${name}' (${data.byteLength} B) exceeds its buffer (${buf.size} B)`);
88 }
89 device.queue.writeBuffer(buf, 0, data as Float32Array<ArrayBuffer>);
90 }
92 // Pipelines: one per distinct shader used, override constants from the
93 // scope. An unknown constant name is a real authoring error; surface it.
94 const passes: PassSpec[] = [...phase.setup, ...(phase.loop?.body ?? [])];
95 const pipelines = new Map<string, GPUComputePipeline>();
96 device.pushErrorScope('validation');
97 for (const p of passes) {
98 if (pipelines.has(p.shader)) continue;
99 const s = shaders[p.shader];
100 const constants: Record<string, number> = {};
101 for (const c of s.constants) {
102 if (!(c in scope)) throw new Error(`shader '${p.shader}': constant '${c}' is not in scope`);
103 constants[c] = scope[c];
104 }
105 pipelines.set(
106 p.shader,
107 await device.createComputePipelineAsync({
108 label: `${label}-${p.shader}`,
109 layout: 'auto',
110 compute: {
111 module: device.createShaderModule({ code: s.code, label: `${label}-${p.shader}` }),
112 entryPoint: 'main',
113 constants,
114 },
115 }),
116 );
117 }
118 const bumpPipe = await device.createComputePipelineAsync({
119 label: `${label}-bump`,
120 layout: 'auto',
121 compute: { module: device.createShaderModule({ code: bumpShader }), entryPoint: 'main' },
122 });
123 const pipeErr = await device.popErrorScope();
124 if (pipeErr) throw new Error(`${label}: shader compilation failed: ${pipeErr.message}`);
126 // Rotation: per phase index, each logical name in a rotate group maps to
127 // the group's physical buffer (index - k) mod n. Bind groups are built
128 // per pass per rotation phase; passes outside the loop use phase 0.
129 const rotate = phase.loop?.rotate ?? [];
130 const nPhases = rotate.reduce((a, g) => lcm(a, g.length), 1);
131 const physical = (name: string, k: number): GPUBuffer => {
132 for (const g of rotate) {
133 const l = g.indexOf(name);
134 if (l >= 0) {
135 const idx = ((l - k) % g.length + g.length) % g.length;
136 const buf = buffers.get(g[idx]);
137 if (!buf) throw new Error(`rotate group buffer '${g[idx]}' is not declared`);
138 return buf;
139 }
140 }
141 const buf = buffers.get(name);
142 if (!buf) throw new Error(`pass binds undeclared buffer '${name}'`);
143 return buf;
144 };
145 const bindOf = (p: PassSpec, k: number): GPUBindGroup =>
146 device.createBindGroup({
147 layout: pipelines.get(p.shader)!.getBindGroupLayout(0),
148 entries: p.buffers.map((name, i) => ({
149 binding: i,
150 resource: { buffer: physical(name, k) },
151 })),
152 });
153 const setupBinds = phase.setup.map((p) => bindOf(p, 0));
154 const loopBinds = (phase.loop?.body ?? []).map((p) =>
155 Array.from({ length: nPhases }, (_, k) => bindOf(p, k)),
156 );
157 const bumpBind = device.createBindGroup({
158 layout: bumpPipe.getBindGroupLayout(0),
159 entries: [{ binding: 0, resource: { buffer: buffers.get('__tick')! } }],
160 });
162 const dispatchCount = new Map<PassSpec, number>();
163 for (const p of passes) dispatchCount.set(p, Math.ceil(evalExpr(p.dispatch, scope) / WORKGROUP_SIZE));
165 const encodePass = (
166 cp: GPUComputePassEncoder,
167 p: PassSpec,
168 bind: GPUBindGroup,
169 ): void => {
170 cp.setPipeline(pipelines.get(p.shader)!);
171 cp.setBindGroup(0, bind);
172 const wg = dispatchCount.get(p)!;
173 cp.dispatchWorkgroups(Math.min(wg, MAX_WG_X), Math.ceil(wg / MAX_WG_X));
174 };
176 // Setup.
177 {
178 const enc = device.createCommandEncoder({ label: `${label}-setup` });
179 const cp = enc.beginComputePass();
180 phase.setup.forEach((p, i) => encodePass(cp, p, setupBinds[i]));
181 cp.end();
182 device.queue.submit([enc.finish()]);
183 }
185 // The loop, batched.
186 if (phase.loop) {
187 const total = evalExpr(phase.loop.count, scope);
188 if (!Number.isInteger(total) || total < 0) {
189 throw new Error(`loop count '${phase.loop.count}' = ${total} is not a whole number`);
190 }
191 const batch = phase.loop.batch ?? 256;
192 const body = phase.loop.body;
193 let done = 0;
194 while (done < total) {
195 if (opts.cancelled?.()) throw new PhaseCancelled();
196 const n = Math.min(batch, total - done);
197 const enc = device.createCommandEncoder({ label: `${label}-loop` });
198 const cp = enc.beginComputePass();
199 for (let j = 0; j < n; j++) {
200 const k = (done + j) % nPhases;
201 body.forEach((p, i) => encodePass(cp, p, loopBinds[i][k]));
202 cp.setPipeline(bumpPipe);
203 cp.setBindGroup(0, bumpBind);
204 cp.dispatchWorkgroups(1);
205 }
206 cp.end();
207 device.queue.submit([enc.finish()]);
208 await device.queue.onSubmittedWorkDone();
209 done += n;
210 opts.onProgress?.(done, total);
211 }
212 }
214 // Readback.
215 const out: Record<string, Float32Array> = {};
216 for (const name of opts.outputs) {
217 const src = buffers.get(name);
218 if (!src) throw new Error(`output '${name}' has no declared buffer`);
219 const rb = device.createBuffer({
220 label: `${label}-read-${name}`,
221 size: src.size,
222 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
223 });
224 owned.push(rb);
225 const enc = device.createCommandEncoder();
226 enc.copyBufferToBuffer(src, 0, rb, 0, src.size);
227 device.queue.submit([enc.finish()]);
228 await rb.mapAsync(GPUMapMode.READ);
229 out[name] = new Float32Array(rb.getMappedRange().slice(0));
230 rb.unmap();
231 }
232 return out;
233 } finally {
234 for (const b of owned) b.destroy();
235 }
236}
238const gcd = (a: number, b: number): number => (b === 0 ? a : gcd(b, a % b));
239const lcm = (a: number, b: number): number => (a / gcd(a, b)) * b;