4f822e1turing-surface-cache: reaction-diffusion solutions at a chosen end time, shared through a cloud cacheJeremy Magland 1/**
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';
12import type {
13 Assign,
14 For,
15 IRExpr,
16 IRStmt,
17 MultiAssignCall,
18} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
19import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
20import { ShtPlan, type ShtBinding, type ShtBatchBinding, type ShtDphigBinding } from '../sht/sht.ts';
21import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
22import type { CompiledFunction } from './compile.ts';
23import { EXTERNAL_OPS } from './externals.ts';
24import {
25 MODE_BUFFER,
26 INITIAL_MODES,
27 modeTableLength,
28 randnfun3Chunks,
29 randnfun3WGSL,
30} from './randnfun3.ts';
31import {
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 }
106 /**
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 }
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 };
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;
166 }
167 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
168 | { kind: 'synth-batch' | 'analys-batch'; binding: ShtBatchBinding; labels: string[] }
169 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
170 | { kind: 'dthetac' | 'dphic'; bindGroup: GPUBindGroup; label: string }
171 | { kind: 'dphig'; binding: ShtDphigBinding; label: string }
172 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
174/**
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}
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[];
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;
341 #device: GPUDevice;
342 #sht: ShtPlan;
343 #deriv?: DerivPlan;
344 #ops: Op[];
345 #owned: GPUBuffer[];
346 #paramBuf: GPUBuffer;
347 #paramData: Float32Array;
348 #rebindRandnfun3: ((table: GPUBuffer) => void) | null;
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;
355 deriv?: DerivPlan;
356 ops: Op[];
357 byName: Map<string, Slot>;
358 owned: GPUBuffer[];
359 paramBuf: GPUBuffer;
360 paramData: Float32Array;
361 paramNames: string[];
362 randnfun3Lambda: Randnfun3Lambda | null;
363 rebindRandnfun3: ((table: GPUBuffer) => void) | null;
364 }) {
365 this.#device = init.device;
366 this.#sht = init.sht;
367 this.#deriv = init.deriv;
368 this.#ops = init.ops;
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;
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);
398 }
400 static async create(
401 device: GPUDevice,
402 sht: ShtPlan,
403 spec: PlanSpec,
404 host: HostBuffers,
405 /** Computes dtheta/dphi — only needed if the .m calls them. */
406 deriv?: DerivPlan,
407 ): Promise<ModelPlan> {
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 });
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;
455 const planned: Planned[] = [];
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 }
478 planned.push({
479 kind: 'copy',
480 from: src.buffer,
481 to: dst.buffer,
482 bytes: 4 * src.count,
483 label: `${out.name} -> ${to}`,
484 });
485 });
487 // Group adjacent independent transforms into batched dispatches and
488 // create every binding.
489 const ops = materializeTransforms(planned, sht);
491 return new ModelPlan({
492 device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
493 randnfun3Lambda, rebindRandnfun3,
494 });
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);
499 if (stmt.kind === 'MultiAssignCall') return planMultiTransform(stmt);
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) {
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);
544 if (!argSlot) {
545 throw new UnsupportedOnGpu(
546 `'${ext.name}' reads '${arg.name}', which has no buffer`,
547 stmt.span,
548 );
549 }
550 const label = `${stmt.name} = ${ext.name}(${arg.name})`;
551 if (ext.name === 'dphig') {
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 }
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,
569 label,
570 });
571 } else if (
572 ext.name === 'dtheta' || ext.name === 'dphi' ||
573 ext.name === 'dthetac' || ext.name === 'dphic'
574 ) {
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 }
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(
590 `'${stmt.name} = ${ext.name}(${arg.name})' reads and ` +
591 `writes the same buffer; assign to a new name instead`,
592 stmt.span,
593 );
594 }
595 planned.push({
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 }
605 planned.push(
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 }
613 return;
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 );
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 } });
660 planned.push({
661 kind: 'kernel',
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 }
675 /**
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 }
744 /**
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 }
853 /**
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 }
934 /**
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 }
980 /**
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 {
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 {
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 };
1003 for (const op of ops) {
1004 switch (op.kind) {
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;
1024 case 'dtheta':
1025 this.#derivInto(inPass(), op);
1026 break;
1027 case 'dphi':
1028 this.#derivInto(inPass(), op);
1029 break;
1030 case 'dthetac':
1031 this.#deriv!.encodeDthetacInto(inPass(), op.bindGroup);
1032 break;
1033 case 'dphic':
1034 this.#deriv!.encodeDphicInto(inPass(), op.bindGroup);
1035 break;
1036 case 'dphig':
1037 this.#sht.encodeDphigInto(inPass(), op.binding);
1038 break;
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;
1045 case 'copy':
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 }
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 }
1067 /**
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 */
1073 describe(): string[] {
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 });
1084 }
1086 destroy(): void {
1087 for (const b of this.#owned) b.destroy();
1088 this.#paramBuf.destroy();
1089 this.#owned.length = 0;
1090 }
1091}
1093/**
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 */
1102function externalCall(
1103 stmt: Assign,
1104): { name: string; args: IRExpr[] } | null {
1105 const e = stmt.expr;
1106 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
1107 const arity = e.name === 'randnfun3' ? 4 : 1;
1108 if (e.args.length !== arity) {
1109 throw new UnsupportedOnGpu(
1110 arity === 1
1111 ? `'${e.name}' must be applied to a single variable`
1112 : `'${e.name}' takes ${arity} arguments, got ${e.args.length}`,
1113 stmt.span,
1114 );
1115 }
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 };
1126}
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]);
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}