2 * Statement list -> a replayable sequence of GPU operations.
3 *
4 * Everything expensive happens once, here: pipeline compilation, buffer
5 * allocation, bind-group construction. Because numbl fixes every type and
6 * shape at lowering time, the resulting op sequence is fully static — so
7 * `encodeStep` is pure synchronous command recording, with no allocation, no
8 * pipeline lookup and no readback. That is what lets the whole timestep be
9 * encoded into one submit and keeps the CPU out of the loop.
10 */
11import { isMultiElement, scalarDouble } from 'numbl-src/numbl-core/jit/lowering/types.ts';
13 Assign,
14 For,
15 IRExpr,
16 IRStmt,
17 MultiAssignCall,
18} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 19import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 20import { ShtPlan, type ShtBinding, type ShtBatchBinding, type ShtDphigBinding } from '../sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 21import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 22import type { CompiledFunction } from './compile.ts';
23import { EXTERNAL_OPS } from './externals.ts';
25 MODE_BUFFER,
26 INITIAL_MODES,
27 modeTableLength,
28 randnfun3Chunks,
29 randnfun3WGSL,
30} from './randnfun3.ts';
32 buildKernel,
33 UnsupportedOnGpu,
34 WORKGROUP_SIZE,
35 type KernelInputs,
36} from './wgsl.ts';
38const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
39const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
40const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
42/**
43 * The compile-time value of a scalar expression, if it has one. A literal
44 * carries its own; a variable carries one when it was bound to a `const` (the
45 * host's fixed scalars) or computed from constants, because numbl propagates
46 * `exact` through the type lattice.
47 */
48const exactValue = (e: IRExpr): number | undefined => {
49 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
50 return e.kind === 'NumLit' ? e.value : undefined;
51};
53/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
54 * its own pipelines at compile time and its own dispatches per step — so a
55 * runaway bound should be a clear error rather than a hang. */
56const MAX_UNROLL = 64;
58interface Slot {
59 buffer: GPUBuffer;
60 count: number;
61}
63const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
64 device.createBuffer({
65 label,
66 size: 4 * count,
67 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
68 });
70/**
71 * Buffers for host-bound variables, shared across plans.
72 *
73 * A model is two programs — `init` and `step` — compiled separately but
74 * operating on the same state. `U` in the step must be the very buffer `init`
75 * wrote, so the buffers for host bindings live here rather than inside either
76 * plan.
77 */
78export class HostBuffers {
79 #device: GPUDevice;
80 #slots = new Map<string, Slot>();
82 constructor(device: GPUDevice) {
83 this.#device = device;
84 }
86 ensure(name: string, count: number): Slot {
87 const existing = this.#slots.get(name);
88 if (existing) {
89 if (existing.count !== count) {
90 throw new UnsupportedOnGpu(
91 `'${name}' is ${existing.count} elements in one program and ` +
92 `${count} in another`,
93 );
94 }
95 return existing;
96 }
97 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
98 this.#slots.set(name, slot);
99 return slot;
100 }
102 get(name: string): Slot | undefined {
103 return this.#slots.get(name);
104 }
107 * Replace a slot's buffer with a larger one. Only for buffers whose size is
108 * not fixed by the grid — the randnfun3 mode table, which grows with the
109 * wavelength asked for. The caller must rebuild any bind group holding the
110 * old buffer; it is destroyed here.
111 */
112 resize(name: string, count: number): Slot {
113 const existing = this.#slots.get(name);
114 if (!existing) throw new Error(`resize: no buffer named '${name}'`);
115 if (count <= existing.count) return existing;
116 existing.buffer.destroy();
117 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
118 this.#slots.set(name, slot);
119 return slot;
120 }
122 /** Upload into the front of a slot, leaving any tail as it was. For a
123 * variable-length payload in a buffer sized to its high-water mark. */
124 uploadInto(name: string, data: Float32Array): void {
125 const slot = this.#slots.get(name);
126 if (!slot) throw new Error(`uploadInto: no buffer named '${name}'`);
127 if (data.length > slot.count) {
128 throw new Error(
129 `uploadInto '${name}': ${data.length} elements into a ${slot.count}-element buffer`,
130 );
131 }
132 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
133 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 135 /** Upload initial data for a host binding. */
136 upload(name: string, data: Float32Array): void {
137 const slot = this.#slots.get(name);
138 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
139 if (data.length !== slot.count) {
140 throw new Error(
141 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
142 );
143 }
144 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
145 }
147 destroy(): void {
148 for (const s of this.#slots.values()) s.buffer.destroy();
149 this.#slots.clear();
150 }
151}
153type Op =
154 | {
155 kind: 'kernel';
156 pipeline: GPUComputePipeline;
157 bindGroup: GPUBindGroup;
158 count: number;
159 label: string;
160 /** Set when the kernel had to write to scratch because its output
161 * aliases one of its inputs; copied back after the dispatch. */
162 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 163 /** End the submission here when run through `submitYielding`, so the
164 * GPU is handed back between chunks of a long seed. */
165 yieldAfter?: boolean;
167 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 168 | { kind: 'synth-batch' | 'analys-batch'; binding: ShtBatchBinding; labels: string[] }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 169 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 170 | { kind: 'dthetac' | 'dphic'; bindGroup: GPUBindGroup; label: string }
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 171 | { kind: 'dphig'; binding: ShtDphigBinding; label: string }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 172 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
175 * A transform op as planned, before bindings exist: `in`/`out` are the
176 * caller-side buffers (spectral in / grid out for synth, the reverse for
177 * analys). Kept unbound until every statement is planned so that adjacent
178 * independent transforms of the same kind can be grouped into one batched
179 * dispatch (ShtPlan.createSynthBatchBinding) — the Legendre recurrence is
180 * the expensive shared part, and a batch walks it once for all lanes.
181 */
182interface PendingSht {
183 pending: true;
184 kind: 'synth' | 'analys';
185 in: GPUBuffer;
186 out: GPUBuffer;
187 label: string;
188}
190type Planned = Op | PendingSht;
192const isPending = (op: Planned): op is PendingSht => 'pending' in op;
194/**
195 * Group maximal runs of adjacent same-kind transforms into batches of the
196 * widest compiled lane count, and create all bindings. Only literal
197 * adjacency in the op sequence is batched — no reordering — so the models
198 * are written to keep batchable transforms consecutive (see the solve loops
199 * in models/*.m). Batching changes dispatch shape only: per-lane arithmetic
200 * is identical to the scalar kernels', so results do not depend on batchK.
201 */
202function materializeTransforms(planned: Planned[], sht: ShtPlan): Op[] {
203 /** Lanes must not collide: distinct outputs, and no lane reading another's
204 * output (repeated read-only inputs would be harmless, but WebGPU also
205 * forbids aliasing a writable binding, so outputs are the hard rule). */
206 const disjoint = (members: PendingSht[]): boolean => {
207 const outs = new Set<GPUBuffer>();
208 for (const m of members) {
209 if (outs.has(m.out)) return false;
210 outs.add(m.out);
211 }
212 return members.every((m) => !outs.has(m.in));
213 };
214 const bind = (m: PendingSht): Op =>
215 m.kind === 'synth'
216 ? { kind: 'synth', binding: sht.createSynthBinding(m.in, m.out), label: m.label }
217 : { kind: 'analys', binding: sht.createAnalysBinding(m.in, m.out), label: m.label };
218 const bindBatch = (members: PendingSht[]): Op =>
219 members[0].kind === 'synth'
220 ? {
221 kind: 'synth-batch',
222 binding: sht.createSynthBatchBinding(
223 members.map((m) => ({ qlmIn: m.in, spatOut: m.out })),
224 ),
225 labels: members.map((m) => m.label),
226 }
227 : {
228 kind: 'analys-batch',
229 binding: sht.createAnalysBatchBinding(
230 members.map((m) => ({ spatIn: m.in, qlmOut: m.out })),
231 ),
232 labels: members.map((m) => m.label),
233 };
235 const out: Op[] = [];
236 let i = 0;
237 while (i < planned.length) {
238 const op = planned[i];
239 if (!isPending(op)) {
240 out.push(op);
241 i++;
242 continue;
243 }
244 let j = i;
245 while (j < planned.length) {
246 const p = planned[j];
247 if (!isPending(p) || p.kind !== op.kind) break;
248 j++;
249 }
250 const run = planned.slice(i, j) as PendingSht[];
251 let s = 0;
252 while (s < run.length) {
253 let take = 1;
254 for (const K of [4, 2]) {
255 if (K > sht.batchK || s + K > run.length) continue;
256 if (disjoint(run.slice(s, s + K))) {
257 take = K;
258 break;
259 }
260 }
261 out.push(take === 1 ? bind(run[s]) : bindBatch(run.slice(s, s + take)));
262 s += take;
263 }
264 i = j;
265 }
266 return out;
267}
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 269export interface PlanSpec {
270 /** The specialized function this plan executes. */
271 fn: CompiledFunction;
272 /** Output index -> host binding name to copy the result into after the run,
273 * so the next call reads it (the new spectral state feeds the old). */
274 feedback: (string | null)[];
275}
277/**
278 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
279 * buffers after it, then the params buffer.
280 *
281 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
282 * only contains the bindings the shader actually references — so a kernel that
283 * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
284 * binding and no longer match the bind group. An explicit layout may carry
285 * bindings the shader ignores.
286 */
287function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
288 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
289 binding,
290 visibility: GPUShaderStage.COMPUTE,
291 buffer: { type: 'read-only-storage' },
292 });
293 return device.createBindGroupLayout({
294 entries: [
295 {
296 binding: 0,
297 visibility: GPUShaderStage.COMPUTE,
298 buffer: { type: 'storage' },
299 },
300 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
301 readOnly(inputs + 1),
302 ],
303 });
304}
306async function makePipeline(
307 device: GPUDevice,
308 code: string,
309 label: string,
310 bindGroupLayout: GPUBindGroupLayout,
311): Promise<GPUComputePipeline> {
312 device.pushErrorScope('validation');
313 const module = device.createShaderModule({ code, label });
314 const info = await module.getCompilationInfo();
315 const errors = info.messages.filter((m) => m.type === 'error');
316 if (errors.length) {
317 throw new UnsupportedOnGpu(
318 `generated WGSL failed to compile for '${label}':\n` +
319 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
320 `\n--- shader ---\n${code}`,
321 );
322 }
323 const pipeline = await device.createComputePipelineAsync({
324 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
325 compute: { module, entryPoint: 'main' },
326 label,
327 });
328 const err = await device.popErrorScope();
329 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
330 return pipeline;
331}
333/** A compiled .m step, ready to run on the GPU. */
334export class ModelPlan {
335 /** Scalar parameter names, in the order the params buffer expects them. */
336 readonly paramNames: string[];
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 337 /** The wavelength this plan's `randnfun3` call asked for, or null if it
338 * makes none. The host draws the coefficient table from it. */
339 readonly randnfun3Lambda: Randnfun3Lambda | null;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 340
341 #device: GPUDevice;
342 #sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 343 #deriv?: DerivPlan;
345 #owned: GPUBuffer[];
346 #paramBuf: GPUBuffer;
347 #paramData: Float32Array;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 348 #rebindRandnfun3: ((table: GPUBuffer) => void) | null;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 349 /** Public name -> buffer, for uploading initial state and reading results. */
350 #byName: Map<string, Slot>;
352 private constructor(init: {
353 device: GPUDevice;
354 sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 355 deriv?: DerivPlan;
357 byName: Map<string, Slot>;
358 owned: GPUBuffer[];
359 paramBuf: GPUBuffer;
360 paramData: Float32Array;
361 paramNames: string[];
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 362 randnfun3Lambda: Randnfun3Lambda | null;
363 rebindRandnfun3: ((table: GPUBuffer) => void) | null;
365 this.#device = init.device;
366 this.#sht = init.sht;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 367 this.#deriv = init.deriv;
369 this.#byName = init.byName;
370 this.#owned = init.owned;
371 this.#paramBuf = init.paramBuf;
372 this.#paramData = init.paramData;
373 this.paramNames = init.paramNames;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 374 this.randnfun3Lambda = init.randnfun3Lambda;
375 this.#rebindRandnfun3 = init.rebindRandnfun3;
376 }
378 /**
379 * Point the randnfun3 dispatch at a mode table big enough for `data`,
380 * growing the buffer if this wavelength needs more modes than the last one,
381 * and upload it.
382 */
383 uploadRandnfun3Table(host: HostBuffers, data: Float32Array): void {
384 const slot = host.get(MODE_BUFFER);
385 if (!slot || !this.#rebindRandnfun3) return;
386 if (data.length > slot.count) {
387 const max = this.#device.limits.maxStorageBufferBindingSize;
388 if (4 * data.length > max) {
389 throw new Error(
390 `randnfun3: this wavelength needs a ${(4 * data.length / 1e6).toFixed(0)} MB ` +
391 `mode table, past this device's ${(max / 1e6).toFixed(0)} MB limit ` +
392 `on a single buffer. Use a larger lambda.`,
393 );
394 }
395 this.#rebindRandnfun3(host.resize(MODE_BUFFER, data.length).buffer);
396 }
397 host.uploadInto(MODE_BUFFER, data);
400 static async create(
401 device: GPUDevice,
402 sht: ShtPlan,
403 spec: PlanSpec,
404 host: HostBuffers,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 405 /** Computes dtheta/dphi — only needed if the .m calls them. */
406 deriv?: DerivPlan,
408 const { fn } = spec;
410 const slots = new Map<string, Slot>();
411 const byName = new Map<string, Slot>();
412 const owned: GPUBuffer[] = [];
413 /** Scalars the .m computes from its parameters, by cName. */
414 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
416 const alloc = (label: string, count: number): Slot => {
417 const buffer = makeBuffer(device, label, count);
418 owned.push(buffer);
419 return { buffer, count };
420 };
422 // Arguments, bound by what the function's signature declares. Array
423 // arguments come from the shared pool, so a value one function returns is
424 // the same buffer the next one reads. Scalar parameters share one small
425 // storage buffer, in signature order.
426 const paramNames: string[] = [];
427 const paramSlots = new Map<string, number>();
428 for (const p of fn.params) {
429 if (p.binding.kind === 'tensor') {
430 const count = p.binding.shape.reduce((x, y) => x * y, 1);
431 const slot = host.ensure(p.name, count);
432 slots.set(p.cName, slot);
433 byName.set(p.name, slot);
434 } else if (p.binding.kind === 'param') {
435 paramSlots.set(p.cName, paramNames.length);
436 paramNames.push(p.name);
437 }
438 // `const` arguments are exact in the IR and fold into the kernels.
439 }
440 const paramData = new Float32Array(Math.max(1, paramNames.length));
441 const paramBuf = device.createBuffer({
442 label: 'mgpu-params',
443 size: 4 * paramData.length,
444 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
445 });
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 447 /** Set when the .m calls `randnfun3`: which wavelength it asked for, so
448 * the host draws the coefficient table the kernel reads from exactly
449 * that value (src/mgpu/randnfun3.ts). */
450 let randnfun3Lambda: Randnfun3Lambda | null = null;
451 /** Rebuilds the randnfun3 dispatch's bind group after the mode table is
452 * reallocated for a finer wavelength. */
453 let rebindRandnfun3: ((table: GPUBuffer) => void) | null = null;
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 455 const planned: Planned[] = [];
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 456 for (const stmt of fn.body) {
457 await planStatement(stmt);
458 }
460 // Feed declared outputs back into the argument buffers they replace.
461 fn.outputs.forEach((out, i) => {
462 const to = spec.feedback[i];
463 if (!to) return;
464 const src = slots.get(out.cName);
465 const dst = host.get(to);
466 if (!src) {
467 throw new UnsupportedOnGpu(
468 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
469 );
470 }
471 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
472 if (src.count !== dst.count) {
473 throw new UnsupportedOnGpu(
474 `'${out.name}' (${src.count} elements) cannot feed ` +
475 `'${to}' (${dst.count})`,
476 );
477 }
480 from: src.buffer,
481 to: dst.buffer,
482 bytes: 4 * src.count,
483 label: `${out.name} -> ${to}`,
484 });
485 });
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 487 // Group adjacent independent transforms into batched dispatches and
488 // create every binding.
489 const ops = materializeTransforms(planned, sht);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 492 device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 493 randnfun3Lambda, rebindRandnfun3,
496 async function planStatement(stmt: IRStmt): Promise<void> {
497 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
498 if (stmt.kind === 'For') return planFor(stmt);
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 499 if (stmt.kind === 'MultiAssignCall') return planMultiTransform(stmt);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 500 if (stmt.kind !== 'Assign') {
501 throw new UnsupportedOnGpu(
502 `a model function body may only contain assignments ` +
503 `(found '${stmt.kind}')`,
504 stmt.span,
505 );
506 }
507 if (!isNumeric(stmt.ty)) {
508 throw new UnsupportedOnGpu(
509 `'${stmt.name}' is not a numeric value`,
510 stmt.span,
511 );
512 }
513 if (!isTensor(stmt.ty)) {
514 // A scalar the model derives from its parameters (`us = a + b`). It
515 // gets no buffer and no dispatch: the kernels that read it bind it as
516 // a `let` in their prologue.
517 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
518 return;
519 }
520 const count = numel(stmt.ty);
522 // Reuse the destination buffer across steps: the same cName always maps
523 // to the same buffer, so a step allocates nothing.
524 let dest = slots.get(stmt.cName);
525 if (!dest) {
526 dest = alloc(`mgpu-${stmt.name}`, count);
527 slots.set(stmt.cName, dest);
528 } else if (dest.count !== count) {
529 throw new UnsupportedOnGpu(
530 `'${stmt.name}' changes size between assignments`,
531 stmt.span,
532 );
533 }
534 byName.set(stmt.name, dest);
536 const ext = externalCall(stmt);
537 if (ext) {
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 538 if (ext.name === 'randnfun3') {
539 await planRandnfun3(stmt, ext.args, dest);
540 return;
541 }
542 const arg = ext.args[0] as IRExpr & { kind: 'Var' };
543 const argSlot = slots.get(arg.cName);
545 throw new UnsupportedOnGpu(
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 546 `'${ext.name}' reads '${arg.name}', which has no buffer`,
548 );
549 }
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 550 const label = `${stmt.name} = ${ext.name}(${arg.name})`;
552 // Grid -> grid, staged through the plan's fm scratch; safe even
553 // in place, so no aliasing guard is needed.
554 planned.push({
555 kind: 'dphig',
556 binding: sht.createDphigBinding(argSlot.buffer, dest.buffer),
557 label,
558 });
559 return;
560 }
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 561 if (ext.name === 'synth' || ext.name === 'analys') {
562 // Left unbound until materializeTransforms has grouped adjacent
563 // independent transforms into batched dispatches.
564 planned.push({
565 pending: true,
566 kind: ext.name,
567 in: argSlot.buffer,
568 out: dest.buffer,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 569 label,
570 });
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 571 } else if (
572 ext.name === 'dtheta' || ext.name === 'dphi' ||
573 ext.name === 'dthetac' || ext.name === 'dphic'
574 ) {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 575 if (!deriv) {
576 throw new UnsupportedOnGpu(
577 `'${ext.name}' needs the surface's derivative transforms, ` +
578 `which this plan was not given`,
579 stmt.span,
580 );
581 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 582 if (ext.name === 'dthetac' || ext.name === 'dphic') {
583 // Coefficient-space shuffles read at l+-1 (dthetac) or in place
584 // (dphic) and cannot alias their output: WebGPU forbids one buffer
585 // being readable and writable storage in the same dispatch, and
586 // there is no scratch-copy fallback here — refuse rather than
587 // silently reroute.
588 if (argSlot.buffer === dest.buffer) {
589 throw new UnsupportedOnGpu(
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 590 `'${stmt.name} = ${ext.name}(${arg.name})' reads and ` +
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 591 `writes the same buffer; assign to a new name instead`,
592 stmt.span,
593 );
594 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 596 kind: ext.name,
597 bindGroup:
598 ext.name === 'dthetac'
599 ? deriv.createDthetacBinding(argSlot.buffer, dest.buffer)
600 : deriv.createDphicBinding(argSlot.buffer, dest.buffer),
601 label,
602 });
603 return;
604 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 606 ext.name === 'dtheta'
607 ? { kind: 'dtheta', binding: deriv.createDthetaBinding(argSlot.buffer, dest.buffer), label }
608 : { kind: 'dphi', binding: deriv.createDphiBinding(argSlot.buffer, dest.buffer), label },
609 );
610 } else {
611 throw new UnsupportedOnGpu(`unknown external op '${ext.name}'`, stmt.span);
612 }
614 }
616 // Element-wise kernel. Collect the distinct tensor operands and give
617 // them dense binding slots.
618 const tensors = new Map<string, number>();
619 collectTensorVars(stmt.expr, (cName) => {
620 if (!tensors.has(cName)) tensors.set(cName, tensors.size);
621 });
623 const label = `${stmt.name} = <${count} elements, element-wise>`;
624 const kernel = buildKernel(
625 stmt,
626 {
627 tensors,
628 params: paramSlots,
629 scalars: derivedScalars,
630 } satisfies KernelInputs,
631 count,
632 label,
633 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 634
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 635 const bindGroupLayout = kernelLayout(device, tensors.size);
636 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
638 // WebGPU forbids aliasing a writable storage binding with another
639 // binding in the same group, so an in-place update (`u = u + 1`) writes
640 // to scratch and copies back. Element-wise kernels only ever touch
641 // their own index, so the copy is the only cost.
642 const aliased = tensors.has(stmt.cName);
643 const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
645 const entries: GPUBindGroupEntry[] = [
646 { binding: 0, resource: { buffer: target.buffer } },
647 ];
648 for (const [cName, i] of tensors) {
649 const s = slots.get(cName);
650 if (!s) {
651 throw new UnsupportedOnGpu(
652 `'${stmt.name}' reads a value with no buffer`,
653 stmt.span,
654 );
655 }
656 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
657 }
658 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
662 pipeline,
663 bindGroup: device.createBindGroup({
664 layout: bindGroupLayout,
665 entries,
666 }),
667 count,
668 label,
669 copyBack: aliased
670 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
671 : undefined,
672 });
673 }
676 * `[a, b] = synth(x, y)` / `[a, b] = analys(x, y)`: an explicitly grouped
677 * transform — output k is the transform of argument k. The group is
678 * planned as consecutive pending transforms, which materializeTransforms
679 * then chunks into whatever batched dispatch widths the device supports
680 * (one x4 batch, two x2, or scalars with SHT_BATCH=0) — the syntax
681 * promises grouping intent, never a lane width, so the same source
682 * compiles everywhere.
683 */
684 function planMultiTransform(stmt: MultiAssignCall): void {
685 if (stmt.name !== 'synth' && stmt.name !== 'analys') {
686 throw new UnsupportedOnGpu(
687 `'${stmt.name}' does not return multiple values here — only the ` +
688 `transforms ('synth', 'analys') support [a, b] = op(x, y) grouping`,
689 stmt.span,
690 );
691 }
692 const kind = stmt.name;
693 for (let i = 0; i < stmt.outputs.length; i++) {
694 const slot = stmt.outputs[i];
695 const arg = stmt.args[i];
696 if (!slot.binding) {
697 throw new UnsupportedOnGpu(
698 `every output of '${kind}' must be bound to a name — output ` +
699 `${i + 1} is dropped, but each input costs a transform`,
700 stmt.span,
701 );
702 }
703 if (!arg || arg.kind !== 'Var') {
704 throw new UnsupportedOnGpu(
705 `'${kind}' must be applied to variables (argument ${i + 1})`,
706 stmt.span,
707 );
708 }
709 const argSlot = slots.get(arg.cName);
710 if (!argSlot) {
711 throw new UnsupportedOnGpu(
712 `'${kind}' reads '${arg.name}', which has no buffer`,
713 stmt.span,
714 );
715 }
716 if (!isNumeric(slot.ty) || !isTensor(slot.ty)) {
717 throw new UnsupportedOnGpu(
718 `'${slot.binding.name}' is not a numeric array`,
719 stmt.span,
720 );
721 }
722 const count = numel(slot.ty);
723 let dest = slots.get(slot.binding.cName);
724 if (!dest) {
725 dest = alloc(`mgpu-${slot.binding.name}`, count);
726 slots.set(slot.binding.cName, dest);
727 } else if (dest.count !== count) {
728 throw new UnsupportedOnGpu(
729 `'${slot.binding.name}' changes size between assignments`,
730 stmt.span,
731 );
732 }
733 byName.set(slot.binding.name, dest);
734 planned.push({
735 pending: true,
736 kind,
737 in: argSlot.buffer,
738 out: dest.buffer,
739 label: `${slot.binding.name} = ${kind}(${arg.name})`,
740 });
741 }
742 }
745 * `f = randnfun3(lambda, gx, gy, gz)`: the seeded random field, summed
746 * over its Fourier modes at every surface point.
747 *
748 * One dispatch, one thread per point. The coefficient table is not an
749 * argument — it is a host buffer this plan binds and the host refills per
750 * seed, the way `synth` reads Legendre matrices the .m never names. What
751 * the .m *does* choose is the wavelength, which is recorded here so the
752 * host draws the table for exactly that value.
753 */
754 async function planRandnfun3(
755 stmt: Assign,
756 args: IRExpr[],
757 dest: Slot,
758 ): Promise<void> {
759 const lam = args[0];
760 const lambda: Randnfun3Lambda | null =
761 lam.kind === 'NumLit'
762 ? { kind: 'const', value: lam.value }
763 : lam.kind === 'Var' && paramSlots.has(lam.cName)
764 ? { kind: 'param', name: lam.name }
765 : null;
766 if (!lambda) {
767 throw new UnsupportedOnGpu(
768 `randnfun3's wavelength is drawn on the host before the step runs, ` +
769 `so it must be a number or a model parameter — not a value ` +
770 `computed on the GPU`,
771 stmt.span,
772 );
773 }
774 if (randnfun3Lambda && !sameLambda(randnfun3Lambda, lambda)) {
775 throw new UnsupportedOnGpu(
776 `this function calls randnfun3 with two different wavelengths; ` +
777 `one coefficient table is drawn per plan, so only one is supported`,
778 stmt.span,
779 );
780 }
781 randnfun3Lambda = lambda;
783 const points = args.slice(1).map((a) => {
784 const v = a as IRExpr & { kind: 'Var' };
785 const slot = slots.get(v.cName);
786 if (!slot) {
787 throw new UnsupportedOnGpu(
788 `randnfun3 reads '${v.name}', which has no buffer`,
789 stmt.span,
790 );
791 }
792 return { slot, name: v.name };
793 });
795 const modes = host.ensure(MODE_BUFFER, modeTableLength(INITIAL_MODES));
796 const label =
797 `${stmt.name} = randnfun3(${
798 lambda.kind === 'const' ? lambda.value : lambda.name
799 }, ${points.map((p) => p.name).join(', ')})`;
801 const bindGroupLayout = device.createBindGroupLayout({
802 label: 'mgpu-randnfun3',
803 entries: [0, 1, 2, 3, 4].map((binding) => ({
804 binding,
805 visibility: GPUShaderStage.COMPUTE,
806 buffer: { type: binding === 0 ? ('storage' as const) : ('read-only-storage' as const) },
807 })),
808 });
809 // The table is sized to whatever wavelength is actually asked for, so a
810 // finer one reallocates it — and with it these bind groups, which are
811 // the only things holding the old buffer.
812 const bind = (table: GPUBuffer): GPUBindGroup =>
813 device.createBindGroup({
814 layout: bindGroupLayout,
815 entries: [
816 { binding: 0, resource: { buffer: dest.buffer } },
817 ...points.map((p, i) => ({
818 binding: i + 1,
819 resource: { buffer: p.slot.buffer },
820 })),
821 { binding: 4, resource: { buffer: table } },
822 ],
823 });
825 // One dispatch per slice of the mode table — see randnfun3Chunks. Each
826 // reads the same table and accumulates into the same output, so they
827 // share a bind group and differ only in their compiled slice index.
828 const ops: (Op & { kind: 'kernel' })[] = [];
829 for (let chunk = 0; chunk < randnfun3Chunks; chunk++) {
830 const chunkLabel = `${label} [${chunk + 1}/${randnfun3Chunks}]`;
831 const op = {
832 kind: 'kernel' as const,
833 pipeline: await makePipeline(
834 device,
835 randnfun3WGSL(dest.count, chunk),
836 chunkLabel,
837 bindGroupLayout,
838 ),
839 bindGroup: bind(modes.buffer),
840 count: dest.count,
841 label: chunkLabel,
842 yieldAfter: true,
843 };
844 ops.push(op);
845 planned.push(op);
846 }
847 rebindRandnfun3 = (table: GPUBuffer): void => {
848 const group = bind(table);
849 for (const op of ops) op.bindGroup = group;
850 };
851 }
854 * Unroll a counted loop into the op sequence.
855 *
856 * A plan is a fixed list of GPU operations with no branching, which is what
857 * makes a timestep pure command recording. A `for` with compile-time-known
858 * bounds still fits that: it is the same body planned once per iteration.
859 * Nothing else changes — numbl gives a variable one cName for every
860 * assignment to it, so the buffer an iteration writes is the buffer the
861 * next one reads, which is exactly a loop-carried value.
862 *
863 * The loop variable gets no buffer either: it is bound as a derived scalar
864 * to this iteration's literal value, so a kernel that reads `k` folds the
865 * number in. The binding is overwritten per iteration, before that
866 * iteration's body is planned and its WGSL emitted.
867 */
868 async function planFor(stmt: For): Promise<void> {
869 const from = exactValue(stmt.start);
870 const to = exactValue(stmt.end);
871 if (from === undefined || to === undefined) {
872 throw new UnsupportedOnGpu(
873 `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
874 `be known when the model is compiled — ` +
875 `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
876 `runtime value. Use a whole number, or a count the app supplies ` +
877 `as a fixed argument (changing it recompiles).`,
878 stmt.span,
879 );
880 }
881 const trips = Math.floor((to - from) / stmt.step) + 1;
882 if (!Number.isFinite(trips)) {
883 throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
884 }
885 if (trips > MAX_UNROLL) {
886 throw new UnsupportedOnGpu(
887 `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
888 `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
889 `long loop compiles slowly and runs no faster than writing it out.`,
890 stmt.span,
891 );
892 }
893 for (let i = 0; i < trips; i++) {
894 const value = from + i * stmt.step;
895 derivedScalars.set(stmt.cVar, {
896 name: stmt.varName,
897 expr: {
898 kind: 'NumLit',
899 value,
900 ty: scalarDouble(
901 value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
902 value,
903 ),
904 span: stmt.span,
905 },
906 });
907 for (const s of stmt.body) await planStatement(s);
908 }
909 }
910 }
912 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
913 setParams(values: Record<string, number>): void {
914 this.paramNames.forEach((name, i) => {
915 const v = values[name];
916 this.#paramData[i] = Number.isFinite(v) ? v : 0;
917 });
918 this.#device.queue.writeBuffer(
919 this.#paramBuf,
920 0,
921 this.#paramData as Float32Array<ArrayBuffer>,
922 );
923 }
925 /** Buffer holding the named value, or undefined if the .m never binds it. */
926 buffer(name: string): GPUBuffer | undefined {
927 return this.#byName.get(name)?.buffer;
928 }
930 elementCount(name: string): number | undefined {
931 return this.#byName.get(name)?.count;
932 }
935 * Run one pass of this plan, submitting in pieces so the GPU is not held for
936 * the whole of it.
937 *
938 * For `init` only, and only because the seed field's mode sum can be huge:
939 * at a fine wavelength the dispatches add up to tens of seconds, and a
940 * browser's GPU process is shared with compositing, so one submission that
941 * long stops the whole browser painting — the user's tabs included. Ops
942 * marked `yieldAfter` (the randnfun3 chunks) end their submission and give
943 * the queue back before the next one is recorded, which turns a freeze into
944 * a wait. Everything else is recorded exactly as `encodeSteps` would.
945 */
946 async submitYielding(label: string): Promise<void> {
947 let encoder = this.#device.createCommandEncoder({ label });
948 let any = false;
949 for (const group of this.#yieldGroups()) {
950 if (any) {
951 // Let the queue drain, then hand the event loop back, so compositing
952 // and input get a turn between chunks.
953 await this.#device.queue.onSubmittedWorkDone();
954 await new Promise((r) => setTimeout(r, 0));
955 encoder = this.#device.createCommandEncoder({ label });
956 }
957 this.#encodeOps(encoder, group);
958 this.#device.queue.submit([encoder.finish()]);
959 any = true;
960 }
961 if (!any) {
962 this.#encodeOps(encoder, []);
963 this.#device.queue.submit([encoder.finish()]);
964 }
965 }
967 /** The op list split at every `yieldAfter` boundary. */
968 *#yieldGroups(): Generator<Op[]> {
969 let group: Op[] = [];
970 for (const op of this.#ops) {
971 group.push(op);
972 if (op.kind === 'kernel' && op.yieldAfter) {
973 yield group;
974 group = [];
975 }
976 }
977 if (group.length) yield group;
978 }
981 * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
982 * ops share one compute pass, which WebGPU executes in submission order
983 * with a barrier between dispatches.
984 */
985 encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 986 for (let s = 0; s < steps; s++) this.#encodeOps(encoder, this.#ops);
987 }
989 /** Record one pass over `ops` into `encoder`. */
990 #encodeOps(encoder: GPUCommandEncoder, ops: Op[]): void {
991 {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 992 let pass: GPUComputePassEncoder | null = null;
993 const inPass = (): GPUComputePassEncoder => {
994 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
995 return pass;
996 };
997 const endPass = (): void => {
998 if (pass) {
999 pass.end();
1000 pass = null;
1001 }
1002 };
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 1003 for (const op of ops) {
1005 case 'kernel': {
1006 const p = inPass();
1007 p.setPipeline(op.pipeline);
1008 p.setBindGroup(0, op.bindGroup);
1009 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
1010 if (op.copyBack) {
1011 endPass();
1012 encoder.copyBufferToBuffer(
1013 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
1014 );
1015 }
1016 break;
1017 }
1018 case 'synth':
1019 this.#shtInto(inPass(), op);
1020 break;
1021 case 'analys':
1022 this.#shtInto(inPass(), op);
1023 break;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 1024 case 'dtheta':
1025 this.#derivInto(inPass(), op);
1026 break;
1027 case 'dphi':
1028 this.#derivInto(inPass(), op);
1029 break;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 1030 case 'dthetac':
1031 this.#deriv!.encodeDthetacInto(inPass(), op.bindGroup);
1032 break;
1033 case 'dphic':
1034 this.#deriv!.encodeDphicInto(inPass(), op.bindGroup);
1035 break;
1037 this.#sht.encodeDphigInto(inPass(), op.binding);
1038 break;
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 1039 case 'synth-batch':
1040 this.#sht.encodeSynthBatchInto(inPass(), op.binding);
1041 break;
1042 case 'analys-batch':
1043 this.#sht.encodeAnalysBatchInto(inPass(), op.binding);
1044 break;
1046 endPass();
1047 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
1048 break;
1049 }
1050 }
1051 endPass();
1052 }
1053 }
1055 #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
1056 if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
1057 else this.#sht.encodeAnalysInto(pass, op.binding);
1058 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 1060 #derivInto(pass: GPUComputePassEncoder, op: Op & { kind: 'dtheta' | 'dphi' }): void {
1061 // planStatement already refused to plan a dtheta/dphi op without a
1062 // DerivPlan, so #deriv is guaranteed set whenever an op of this kind exists.
1063 if (op.kind === 'dtheta') this.#deriv!.encodeDthetaInto(pass, op.binding);
1064 else this.#deriv!.encodeDphiInto(pass, op.binding);
1065 }
1068 * Human-readable op sequence — what the .m actually compiled to. Batched
1069 * transforms list one line per lane, annotated: the line count equals the
1070 * logical op count regardless of the device's batch width, so op-count
1071 * assertions in the tests are batch-invariant.
1072 */
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 1074 return this.#ops.flatMap((op) => {
1075 if ('labels' in op) {
1076 const kind = op.kind === 'synth-batch' ? 'synth' : 'analys';
1077 return op.labels.map(
1078 (label, i) =>
1079 `${kind.padEnd(7)} ${label} [batch lane ${i + 1}/${op.binding.size}]`,
1080 );
1081 }
1082 return [`${op.kind.padEnd(7)} ${op.label}`];
1083 });
1086 destroy(): void {
1087 for (const b of this.#owned) b.destroy();
1088 this.#paramBuf.destroy();
1089 this.#owned.length = 0;
1090 }
1091}
1094 * `x = synth(y)` / `x = randnfun3(lam, gx, gy, gz)` -> the call's name and
1095 * arguments.
1096 *
1097 * Every external op but `randnfun3` takes exactly one array; `randnfun3`
1098 * takes a wavelength and the three surface coordinates. Its wavelength may
1099 * be a literal, so arguments are returned as expressions and the caller
1100 * decides which it needs as a buffer.
1101 */
1103 stmt: Assign,
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 1104): { name: string; args: IRExpr[] } | null {
1106 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 1107 const arity = e.name === 'randnfun3' ? 4 : 1;
1108 if (e.args.length !== arity) {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 1109 throw new UnsupportedOnGpu(
1111 ? `'${e.name}' must be applied to a single variable`
1112 : `'${e.name}' takes ${arity} arguments, got ${e.args.length}`,
1114 );
1115 }
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 1116 // Only the wavelength may be something other than a plain variable.
1117 for (let i = e.name === 'randnfun3' ? 1 : 0; i < e.args.length; i++) {
1118 if (e.args[i].kind !== 'Var') {
1119 throw new UnsupportedOnGpu(
1120 `'${e.name}' must be applied to variables, not expressions`,
1121 stmt.span,
1122 );
1123 }
1124 }
1125 return { name: e.name, args: e.args };
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 1128/** A `randnfun3` wavelength argument: a literal, or the parameter to read it
1129 * from when the host fills the coefficient table. */
1130export type Randnfun3Lambda =
1131 | { kind: 'const'; value: number }
1132 | { kind: 'param'; name: string };
1134const sameLambda = (a: Randnfun3Lambda, b: Randnfun3Lambda): boolean =>
1135 a.kind === 'const' && b.kind === 'const'
1136 ? a.value === b.value
1137 : a.kind === 'param' && b.kind === 'param' && a.name === b.name;
1139/** The wavelength value a plan's `randnfun3` call resolves to. */
1140export const resolveLambda = (
1141 lambda: Randnfun3Lambda,
1142 params: Record<string, number>,
1143): number => (lambda.kind === 'const' ? lambda.value : params[lambda.name]);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 1145function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
1146 const walk = (x: IRExpr): void => {
1147 switch (x.kind) {
1148 case 'Var':
1149 if (isTensor(x.ty)) visit(x.cName);
1150 return;
1151 case 'Binary':
1152 walk(x.left);
1153 walk(x.right);
1154 return;
1155 case 'Unary':
1156 walk(x.operand);
1157 return;
1158 case 'Call':
1159 x.args.forEach(walk);
1160 return;
1161 default:
1162 return;
1163 }
1164 };
1165 walk(e);
1166}