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';
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 20import { ShtPlan, type ShtBinding, type ShtBatchBinding } 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';
24import {
25 buildKernel,
26 UnsupportedOnGpu,
27 WORKGROUP_SIZE,
28 type KernelInputs,
29} from './wgsl.ts';
31const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
32const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
33const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
35/**
36 * The compile-time value of a scalar expression, if it has one. A literal
37 * carries its own; a variable carries one when it was bound to a `const` (the
38 * host's fixed scalars) or computed from constants, because numbl propagates
39 * `exact` through the type lattice.
40 */
41const exactValue = (e: IRExpr): number | undefined => {
42 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
43 return e.kind === 'NumLit' ? e.value : undefined;
44};
46/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
47 * its own pipelines at compile time and its own dispatches per step — so a
48 * runaway bound should be a clear error rather than a hang. */
49const MAX_UNROLL = 64;
51interface Slot {
52 buffer: GPUBuffer;
53 count: number;
54}
56const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
57 device.createBuffer({
58 label,
59 size: 4 * count,
60 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
61 });
63/**
64 * Buffers for host-bound variables, shared across plans.
65 *
66 * A model is two programs — `init` and `step` — compiled separately but
67 * operating on the same state. `U` in the step must be the very buffer `init`
68 * wrote, so the buffers for host bindings live here rather than inside either
69 * plan.
70 */
71export class HostBuffers {
72 #device: GPUDevice;
73 #slots = new Map<string, Slot>();
75 constructor(device: GPUDevice) {
76 this.#device = device;
77 }
79 ensure(name: string, count: number): Slot {
80 const existing = this.#slots.get(name);
81 if (existing) {
82 if (existing.count !== count) {
83 throw new UnsupportedOnGpu(
84 `'${name}' is ${existing.count} elements in one program and ` +
85 `${count} in another`,
86 );
87 }
88 return existing;
89 }
90 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
91 this.#slots.set(name, slot);
92 return slot;
93 }
95 get(name: string): Slot | undefined {
96 return this.#slots.get(name);
97 }
99 /** Upload initial data for a host binding. */
100 upload(name: string, data: Float32Array): void {
101 const slot = this.#slots.get(name);
102 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
103 if (data.length !== slot.count) {
104 throw new Error(
105 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
106 );
107 }
108 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
109 }
111 destroy(): void {
112 for (const s of this.#slots.values()) s.buffer.destroy();
113 this.#slots.clear();
114 }
115}
117type Op =
118 | {
119 kind: 'kernel';
120 pipeline: GPUComputePipeline;
121 bindGroup: GPUBindGroup;
122 count: number;
123 label: string;
124 /** Set when the kernel had to write to scratch because its output
125 * aliases one of its inputs; copied back after the dispatch. */
126 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
127 }
128 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 129 | { kind: 'synth-batch' | 'analys-batch'; binding: ShtBatchBinding; labels: string[] }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 130 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 131 | { kind: 'dthetac' | 'dphic'; bindGroup: GPUBindGroup; label: string }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 132 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
135 * A transform op as planned, before bindings exist: `in`/`out` are the
136 * caller-side buffers (spectral in / grid out for synth, the reverse for
137 * analys). Kept unbound until every statement is planned so that adjacent
138 * independent transforms of the same kind can be grouped into one batched
139 * dispatch (ShtPlan.createSynthBatchBinding) — the Legendre recurrence is
140 * the expensive shared part, and a batch walks it once for all lanes.
141 */
142interface PendingSht {
143 pending: true;
144 kind: 'synth' | 'analys';
145 in: GPUBuffer;
146 out: GPUBuffer;
147 label: string;
148}
150type Planned = Op | PendingSht;
152const isPending = (op: Planned): op is PendingSht => 'pending' in op;
154/**
155 * Group maximal runs of adjacent same-kind transforms into batches of the
156 * widest compiled lane count, and create all bindings. Only literal
157 * adjacency in the op sequence is batched — no reordering — so the models
158 * are written to keep batchable transforms consecutive (see the solve loops
159 * in models/*.m). Batching changes dispatch shape only: per-lane arithmetic
160 * is identical to the scalar kernels', so results do not depend on batchK.
161 */
162function materializeTransforms(planned: Planned[], sht: ShtPlan): Op[] {
163 /** Lanes must not collide: distinct outputs, and no lane reading another's
164 * output (repeated read-only inputs would be harmless, but WebGPU also
165 * forbids aliasing a writable binding, so outputs are the hard rule). */
166 const disjoint = (members: PendingSht[]): boolean => {
167 const outs = new Set<GPUBuffer>();
168 for (const m of members) {
169 if (outs.has(m.out)) return false;
170 outs.add(m.out);
171 }
172 return members.every((m) => !outs.has(m.in));
173 };
174 const bind = (m: PendingSht): Op =>
175 m.kind === 'synth'
176 ? { kind: 'synth', binding: sht.createSynthBinding(m.in, m.out), label: m.label }
177 : { kind: 'analys', binding: sht.createAnalysBinding(m.in, m.out), label: m.label };
178 const bindBatch = (members: PendingSht[]): Op =>
179 members[0].kind === 'synth'
180 ? {
181 kind: 'synth-batch',
182 binding: sht.createSynthBatchBinding(
183 members.map((m) => ({ qlmIn: m.in, spatOut: m.out })),
184 ),
185 labels: members.map((m) => m.label),
186 }
187 : {
188 kind: 'analys-batch',
189 binding: sht.createAnalysBatchBinding(
190 members.map((m) => ({ spatIn: m.in, qlmOut: m.out })),
191 ),
192 labels: members.map((m) => m.label),
193 };
195 const out: Op[] = [];
196 let i = 0;
197 while (i < planned.length) {
198 const op = planned[i];
199 if (!isPending(op)) {
200 out.push(op);
201 i++;
202 continue;
203 }
204 let j = i;
205 while (j < planned.length) {
206 const p = planned[j];
207 if (!isPending(p) || p.kind !== op.kind) break;
208 j++;
209 }
210 const run = planned.slice(i, j) as PendingSht[];
211 let s = 0;
212 while (s < run.length) {
213 let take = 1;
214 for (const K of [4, 2]) {
215 if (K > sht.batchK || s + K > run.length) continue;
216 if (disjoint(run.slice(s, s + K))) {
217 take = K;
218 break;
219 }
220 }
221 out.push(take === 1 ? bind(run[s]) : bindBatch(run.slice(s, s + take)));
222 s += take;
223 }
224 i = j;
225 }
226 return out;
227}
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 229export interface PlanSpec {
230 /** The specialized function this plan executes. */
231 fn: CompiledFunction;
232 /** Output index -> host binding name to copy the result into after the run,
233 * so the next call reads it (the new spectral state feeds the old). */
234 feedback: (string | null)[];
235}
237/**
238 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
239 * buffers after it, then the params buffer.
240 *
241 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
242 * only contains the bindings the shader actually references — so a kernel that
243 * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
244 * binding and no longer match the bind group. An explicit layout may carry
245 * bindings the shader ignores.
246 */
247function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
248 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
249 binding,
250 visibility: GPUShaderStage.COMPUTE,
251 buffer: { type: 'read-only-storage' },
252 });
253 return device.createBindGroupLayout({
254 entries: [
255 {
256 binding: 0,
257 visibility: GPUShaderStage.COMPUTE,
258 buffer: { type: 'storage' },
259 },
260 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
261 readOnly(inputs + 1),
262 ],
263 });
264}
266async function makePipeline(
267 device: GPUDevice,
268 code: string,
269 label: string,
270 bindGroupLayout: GPUBindGroupLayout,
271): Promise<GPUComputePipeline> {
272 device.pushErrorScope('validation');
273 const module = device.createShaderModule({ code, label });
274 const info = await module.getCompilationInfo();
275 const errors = info.messages.filter((m) => m.type === 'error');
276 if (errors.length) {
277 throw new UnsupportedOnGpu(
278 `generated WGSL failed to compile for '${label}':\n` +
279 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
280 `\n--- shader ---\n${code}`,
281 );
282 }
283 const pipeline = await device.createComputePipelineAsync({
284 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
285 compute: { module, entryPoint: 'main' },
286 label,
287 });
288 const err = await device.popErrorScope();
289 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
290 return pipeline;
291}
293/** A compiled .m step, ready to run on the GPU. */
294export class ModelPlan {
295 /** Scalar parameter names, in the order the params buffer expects them. */
296 readonly paramNames: string[];
298 #device: GPUDevice;
299 #sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 300 #deriv?: DerivPlan;
302 #owned: GPUBuffer[];
303 #paramBuf: GPUBuffer;
304 #paramData: Float32Array;
305 /** Public name -> buffer, for uploading initial state and reading results. */
306 #byName: Map<string, Slot>;
308 private constructor(init: {
309 device: GPUDevice;
310 sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 311 deriv?: DerivPlan;
313 byName: Map<string, Slot>;
314 owned: GPUBuffer[];
315 paramBuf: GPUBuffer;
316 paramData: Float32Array;
317 paramNames: string[];
318 }) {
319 this.#device = init.device;
320 this.#sht = init.sht;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 321 this.#deriv = init.deriv;
323 this.#byName = init.byName;
324 this.#owned = init.owned;
325 this.#paramBuf = init.paramBuf;
326 this.#paramData = init.paramData;
327 this.paramNames = init.paramNames;
328 }
330 static async create(
331 device: GPUDevice,
332 sht: ShtPlan,
333 spec: PlanSpec,
334 host: HostBuffers,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 335 /** Computes dtheta/dphi — only needed if the .m calls them. */
336 deriv?: DerivPlan,
338 const { fn } = spec;
340 const slots = new Map<string, Slot>();
341 const byName = new Map<string, Slot>();
342 const owned: GPUBuffer[] = [];
343 /** Scalars the .m computes from its parameters, by cName. */
344 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
346 const alloc = (label: string, count: number): Slot => {
347 const buffer = makeBuffer(device, label, count);
348 owned.push(buffer);
349 return { buffer, count };
350 };
352 // Arguments, bound by what the function's signature declares. Array
353 // arguments come from the shared pool, so a value one function returns is
354 // the same buffer the next one reads. Scalar parameters share one small
355 // storage buffer, in signature order.
356 const paramNames: string[] = [];
357 const paramSlots = new Map<string, number>();
358 for (const p of fn.params) {
359 if (p.binding.kind === 'tensor') {
360 const count = p.binding.shape.reduce((x, y) => x * y, 1);
361 const slot = host.ensure(p.name, count);
362 slots.set(p.cName, slot);
363 byName.set(p.name, slot);
364 } else if (p.binding.kind === 'param') {
365 paramSlots.set(p.cName, paramNames.length);
366 paramNames.push(p.name);
367 }
368 // `const` arguments are exact in the IR and fold into the kernels.
369 }
370 const paramData = new Float32Array(Math.max(1, paramNames.length));
371 const paramBuf = device.createBuffer({
372 label: 'mgpu-params',
373 size: 4 * paramData.length,
374 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
375 });
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 377 const planned: Planned[] = [];
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 378 for (const stmt of fn.body) {
379 await planStatement(stmt);
380 }
382 // Feed declared outputs back into the argument buffers they replace.
383 fn.outputs.forEach((out, i) => {
384 const to = spec.feedback[i];
385 if (!to) return;
386 const src = slots.get(out.cName);
387 const dst = host.get(to);
388 if (!src) {
389 throw new UnsupportedOnGpu(
390 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
391 );
392 }
393 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
394 if (src.count !== dst.count) {
395 throw new UnsupportedOnGpu(
396 `'${out.name}' (${src.count} elements) cannot feed ` +
397 `'${to}' (${dst.count})`,
398 );
399 }
402 from: src.buffer,
403 to: dst.buffer,
404 bytes: 4 * src.count,
405 label: `${out.name} -> ${to}`,
406 });
407 });
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 409 // Group adjacent independent transforms into batched dispatches and
410 // create every binding.
411 const ops = materializeTransforms(planned, sht);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 414 device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
417 async function planStatement(stmt: IRStmt): Promise<void> {
418 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
419 if (stmt.kind === 'For') return planFor(stmt);
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 420 if (stmt.kind === 'MultiAssignCall') return planMultiTransform(stmt);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 421 if (stmt.kind !== 'Assign') {
422 throw new UnsupportedOnGpu(
423 `a model function body may only contain assignments ` +
424 `(found '${stmt.kind}')`,
425 stmt.span,
426 );
427 }
428 if (!isNumeric(stmt.ty)) {
429 throw new UnsupportedOnGpu(
430 `'${stmt.name}' is not a numeric value`,
431 stmt.span,
432 );
433 }
434 if (!isTensor(stmt.ty)) {
435 // A scalar the model derives from its parameters (`us = a + b`). It
436 // gets no buffer and no dispatch: the kernels that read it bind it as
437 // a `let` in their prologue.
438 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
439 return;
440 }
441 const count = numel(stmt.ty);
443 // Reuse the destination buffer across steps: the same cName always maps
444 // to the same buffer, so a step allocates nothing.
445 let dest = slots.get(stmt.cName);
446 if (!dest) {
447 dest = alloc(`mgpu-${stmt.name}`, count);
448 slots.set(stmt.cName, dest);
449 } else if (dest.count !== count) {
450 throw new UnsupportedOnGpu(
451 `'${stmt.name}' changes size between assignments`,
452 stmt.span,
453 );
454 }
455 byName.set(stmt.name, dest);
457 const ext = externalCall(stmt);
458 if (ext) {
459 const argSlot = slots.get(ext.argCName);
460 if (!argSlot) {
461 throw new UnsupportedOnGpu(
462 `'${ext.name}' reads '${ext.argName}', which has no buffer`,
463 stmt.span,
464 );
465 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 466 const label = `${stmt.name} = ${ext.name}(${ext.argName})`;
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 467 if (ext.name === 'synth' || ext.name === 'analys') {
468 // Left unbound until materializeTransforms has grouped adjacent
469 // independent transforms into batched dispatches.
470 planned.push({
471 pending: true,
472 kind: ext.name,
473 in: argSlot.buffer,
474 out: dest.buffer,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 475 label,
476 });
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 477 } else if (
478 ext.name === 'dtheta' || ext.name === 'dphi' ||
479 ext.name === 'dthetac' || ext.name === 'dphic'
480 ) {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 481 if (!deriv) {
482 throw new UnsupportedOnGpu(
483 `'${ext.name}' needs the surface's derivative transforms, ` +
484 `which this plan was not given`,
485 stmt.span,
486 );
487 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 488 if (ext.name === 'dthetac' || ext.name === 'dphic') {
489 // Coefficient-space shuffles read at l+-1 (dthetac) or in place
490 // (dphic) and cannot alias their output: WebGPU forbids one buffer
491 // being readable and writable storage in the same dispatch, and
492 // there is no scratch-copy fallback here — refuse rather than
493 // silently reroute.
494 if (argSlot.buffer === dest.buffer) {
495 throw new UnsupportedOnGpu(
496 `'${stmt.name} = ${ext.name}(${ext.argName})' reads and ` +
497 `writes the same buffer; assign to a new name instead`,
498 stmt.span,
499 );
500 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 502 kind: ext.name,
503 bindGroup:
504 ext.name === 'dthetac'
505 ? deriv.createDthetacBinding(argSlot.buffer, dest.buffer)
506 : deriv.createDphicBinding(argSlot.buffer, dest.buffer),
507 label,
508 });
509 return;
510 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 512 ext.name === 'dtheta'
513 ? { kind: 'dtheta', binding: deriv.createDthetaBinding(argSlot.buffer, dest.buffer), label }
514 : { kind: 'dphi', binding: deriv.createDphiBinding(argSlot.buffer, dest.buffer), label },
515 );
516 } else {
517 throw new UnsupportedOnGpu(`unknown external op '${ext.name}'`, stmt.span);
518 }
520 }
522 // Element-wise kernel. Collect the distinct tensor operands and give
523 // them dense binding slots.
524 const tensors = new Map<string, number>();
525 collectTensorVars(stmt.expr, (cName) => {
526 if (!tensors.has(cName)) tensors.set(cName, tensors.size);
527 });
529 const label = `${stmt.name} = <${count} elements, element-wise>`;
530 const kernel = buildKernel(
531 stmt,
532 {
533 tensors,
534 params: paramSlots,
535 scalars: derivedScalars,
536 } satisfies KernelInputs,
537 count,
538 label,
539 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 540
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 541 const bindGroupLayout = kernelLayout(device, tensors.size);
542 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
544 // WebGPU forbids aliasing a writable storage binding with another
545 // binding in the same group, so an in-place update (`u = u + 1`) writes
546 // to scratch and copies back. Element-wise kernels only ever touch
547 // their own index, so the copy is the only cost.
548 const aliased = tensors.has(stmt.cName);
549 const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
551 const entries: GPUBindGroupEntry[] = [
552 { binding: 0, resource: { buffer: target.buffer } },
553 ];
554 for (const [cName, i] of tensors) {
555 const s = slots.get(cName);
556 if (!s) {
557 throw new UnsupportedOnGpu(
558 `'${stmt.name}' reads a value with no buffer`,
559 stmt.span,
560 );
561 }
562 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
563 }
564 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
568 pipeline,
569 bindGroup: device.createBindGroup({
570 layout: bindGroupLayout,
571 entries,
572 }),
573 count,
574 label,
575 copyBack: aliased
576 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
577 : undefined,
578 });
579 }
582 * `[a, b] = synth(x, y)` / `[a, b] = analys(x, y)`: an explicitly grouped
583 * transform — output k is the transform of argument k. The group is
584 * planned as consecutive pending transforms, which materializeTransforms
585 * then chunks into whatever batched dispatch widths the device supports
586 * (one x4 batch, two x2, or scalars with SHT_BATCH=0) — the syntax
587 * promises grouping intent, never a lane width, so the same source
588 * compiles everywhere.
589 */
590 function planMultiTransform(stmt: MultiAssignCall): void {
591 if (stmt.name !== 'synth' && stmt.name !== 'analys') {
592 throw new UnsupportedOnGpu(
593 `'${stmt.name}' does not return multiple values here — only the ` +
594 `transforms ('synth', 'analys') support [a, b] = op(x, y) grouping`,
595 stmt.span,
596 );
597 }
598 const kind = stmt.name;
599 for (let i = 0; i < stmt.outputs.length; i++) {
600 const slot = stmt.outputs[i];
601 const arg = stmt.args[i];
602 if (!slot.binding) {
603 throw new UnsupportedOnGpu(
604 `every output of '${kind}' must be bound to a name — output ` +
605 `${i + 1} is dropped, but each input costs a transform`,
606 stmt.span,
607 );
608 }
609 if (!arg || arg.kind !== 'Var') {
610 throw new UnsupportedOnGpu(
611 `'${kind}' must be applied to variables (argument ${i + 1})`,
612 stmt.span,
613 );
614 }
615 const argSlot = slots.get(arg.cName);
616 if (!argSlot) {
617 throw new UnsupportedOnGpu(
618 `'${kind}' reads '${arg.name}', which has no buffer`,
619 stmt.span,
620 );
621 }
622 if (!isNumeric(slot.ty) || !isTensor(slot.ty)) {
623 throw new UnsupportedOnGpu(
624 `'${slot.binding.name}' is not a numeric array`,
625 stmt.span,
626 );
627 }
628 const count = numel(slot.ty);
629 let dest = slots.get(slot.binding.cName);
630 if (!dest) {
631 dest = alloc(`mgpu-${slot.binding.name}`, count);
632 slots.set(slot.binding.cName, dest);
633 } else if (dest.count !== count) {
634 throw new UnsupportedOnGpu(
635 `'${slot.binding.name}' changes size between assignments`,
636 stmt.span,
637 );
638 }
639 byName.set(slot.binding.name, dest);
640 planned.push({
641 pending: true,
642 kind,
643 in: argSlot.buffer,
644 out: dest.buffer,
645 label: `${slot.binding.name} = ${kind}(${arg.name})`,
646 });
647 }
648 }
651 * Unroll a counted loop into the op sequence.
652 *
653 * A plan is a fixed list of GPU operations with no branching, which is what
654 * makes a timestep pure command recording. A `for` with compile-time-known
655 * bounds still fits that: it is the same body planned once per iteration.
656 * Nothing else changes — numbl gives a variable one cName for every
657 * assignment to it, so the buffer an iteration writes is the buffer the
658 * next one reads, which is exactly a loop-carried value.
659 *
660 * The loop variable gets no buffer either: it is bound as a derived scalar
661 * to this iteration's literal value, so a kernel that reads `k` folds the
662 * number in. The binding is overwritten per iteration, before that
663 * iteration's body is planned and its WGSL emitted.
664 */
665 async function planFor(stmt: For): Promise<void> {
666 const from = exactValue(stmt.start);
667 const to = exactValue(stmt.end);
668 if (from === undefined || to === undefined) {
669 throw new UnsupportedOnGpu(
670 `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
671 `be known when the model is compiled — ` +
672 `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
673 `runtime value. Use a whole number, or a count the app supplies ` +
674 `as a fixed argument (changing it recompiles).`,
675 stmt.span,
676 );
677 }
678 const trips = Math.floor((to - from) / stmt.step) + 1;
679 if (!Number.isFinite(trips)) {
680 throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
681 }
682 if (trips > MAX_UNROLL) {
683 throw new UnsupportedOnGpu(
684 `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
685 `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
686 `long loop compiles slowly and runs no faster than writing it out.`,
687 stmt.span,
688 );
689 }
690 for (let i = 0; i < trips; i++) {
691 const value = from + i * stmt.step;
692 derivedScalars.set(stmt.cVar, {
693 name: stmt.varName,
694 expr: {
695 kind: 'NumLit',
696 value,
697 ty: scalarDouble(
698 value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
699 value,
700 ),
701 span: stmt.span,
702 },
703 });
704 for (const s of stmt.body) await planStatement(s);
705 }
706 }
707 }
709 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
710 setParams(values: Record<string, number>): void {
711 this.paramNames.forEach((name, i) => {
712 const v = values[name];
713 this.#paramData[i] = Number.isFinite(v) ? v : 0;
714 });
715 this.#device.queue.writeBuffer(
716 this.#paramBuf,
717 0,
718 this.#paramData as Float32Array<ArrayBuffer>,
719 );
720 }
722 /** Buffer holding the named value, or undefined if the .m never binds it. */
723 buffer(name: string): GPUBuffer | undefined {
724 return this.#byName.get(name)?.buffer;
725 }
727 elementCount(name: string): number | undefined {
728 return this.#byName.get(name)?.count;
729 }
731 /**
732 * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
733 * ops share one compute pass, which WebGPU executes in submission order
734 * with a barrier between dispatches.
735 */
736 encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
737 for (let s = 0; s < steps; s++) {
738 let pass: GPUComputePassEncoder | null = null;
739 const inPass = (): GPUComputePassEncoder => {
740 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
741 return pass;
742 };
743 const endPass = (): void => {
744 if (pass) {
745 pass.end();
746 pass = null;
747 }
748 };
749 for (const op of this.#ops) {
750 switch (op.kind) {
751 case 'kernel': {
752 const p = inPass();
753 p.setPipeline(op.pipeline);
754 p.setBindGroup(0, op.bindGroup);
755 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
756 if (op.copyBack) {
757 endPass();
758 encoder.copyBufferToBuffer(
759 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
760 );
761 }
762 break;
763 }
764 case 'synth':
765 this.#shtInto(inPass(), op);
766 break;
767 case 'analys':
768 this.#shtInto(inPass(), op);
769 break;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 770 case 'dtheta':
771 this.#derivInto(inPass(), op);
772 break;
773 case 'dphi':
774 this.#derivInto(inPass(), op);
775 break;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 776 case 'dthetac':
777 this.#deriv!.encodeDthetacInto(inPass(), op.bindGroup);
778 break;
779 case 'dphic':
780 this.#deriv!.encodeDphicInto(inPass(), op.bindGroup);
781 break;
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 782 case 'synth-batch':
783 this.#sht.encodeSynthBatchInto(inPass(), op.binding);
784 break;
785 case 'analys-batch':
786 this.#sht.encodeAnalysBatchInto(inPass(), op.binding);
787 break;
789 endPass();
790 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
791 break;
792 }
793 }
794 endPass();
795 }
796 }
798 #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
799 if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
800 else this.#sht.encodeAnalysInto(pass, op.binding);
801 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 803 #derivInto(pass: GPUComputePassEncoder, op: Op & { kind: 'dtheta' | 'dphi' }): void {
804 // planStatement already refused to plan a dtheta/dphi op without a
805 // DerivPlan, so #deriv is guaranteed set whenever an op of this kind exists.
806 if (op.kind === 'dtheta') this.#deriv!.encodeDthetaInto(pass, op.binding);
807 else this.#deriv!.encodeDphiInto(pass, op.binding);
808 }
811 * Human-readable op sequence — what the .m actually compiled to. Batched
812 * transforms list one line per lane, annotated: the line count equals the
813 * logical op count regardless of the device's batch width, so op-count
814 * assertions in the tests are batch-invariant.
815 */
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 817 return this.#ops.flatMap((op) => {
818 if ('labels' in op) {
819 const kind = op.kind === 'synth-batch' ? 'synth' : 'analys';
820 return op.labels.map(
821 (label, i) =>
822 `${kind.padEnd(7)} ${label} [batch lane ${i + 1}/${op.binding.size}]`,
823 );
824 }
825 return [`${op.kind.padEnd(7)} ${op.label}`];
826 });
829 destroy(): void {
830 for (const b of this.#owned) b.destroy();
831 this.#paramBuf.destroy();
832 this.#owned.length = 0;
833 }
834}
836/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
837function externalCall(
838 stmt: Assign,
839): { name: string; argCName: string; argName: string } | null {
840 const e = stmt.expr;
841 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
842 if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
843 throw new UnsupportedOnGpu(
844 `'${e.name}' must be applied to a single variable`,
845 stmt.span,
846 );
847 }
848 const arg = e.args[0];
849 return { name: e.name, argCName: arg.cName, argName: arg.name };
850}
852function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
853 const walk = (x: IRExpr): void => {
854 switch (x.kind) {
855 case 'Var':
856 if (isTensor(x.ty)) visit(x.cName);
857 return;
858 case 'Binary':
859 walk(x.left);
860 walk(x.right);
861 return;
862 case 'Unary':
863 walk(x.operand);
864 return;
865 case 'Call':
866 x.args.forEach(walk);
867 return;
868 default:
869 return;
870 }
871 };
872 walk(e);
873}