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 { Assign, For, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
13import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
14import { ShtPlan, type ShtBinding } from '../sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 15import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 16import { ReducePlan, type DotBinding } from './reduce.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 17import type { CompiledFunction } from './compile.ts';
18import { EXTERNAL_OPS } from './externals.ts';
19import {
20 buildKernel,
21 UnsupportedOnGpu,
22 WORKGROUP_SIZE,
23 type KernelInputs,
24} from './wgsl.ts';
26const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
27const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
28const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 30/** Scalar arithmetic a plan-time evaluator can fold. */
31const PLAN_BINOPS: Record<string, (l: number, r: number) => number> = {
32 plus: (l, r) => l + r,
33 minus: (l, r) => l - r,
34 times: (l, r) => l * r,
35 mtimes: (l, r) => l * r,
36 rdivide: (l, r) => l / r,
37 mrdivide: (l, r) => l / r,
38 power: (l, r) => Math.pow(l, r),
39 mpower: (l, r) => Math.pow(l, r),
42/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
43 * its own pipelines at compile time and its own dispatches per step — so a
44 * runaway bound should be a clear error rather than a hang. */
45const MAX_UNROLL = 64;
47interface Slot {
48 buffer: GPUBuffer;
49 count: number;
50}
52const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
53 device.createBuffer({
54 label,
55 size: 4 * count,
56 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
57 });
59/**
60 * Buffers for host-bound variables, shared across plans.
61 *
62 * A model is two programs — `init` and `step` — compiled separately but
63 * operating on the same state. `U` in the step must be the very buffer `init`
64 * wrote, so the buffers for host bindings live here rather than inside either
65 * plan.
66 */
67export class HostBuffers {
68 #device: GPUDevice;
69 #slots = new Map<string, Slot>();
71 constructor(device: GPUDevice) {
72 this.#device = device;
73 }
75 ensure(name: string, count: number): Slot {
76 const existing = this.#slots.get(name);
77 if (existing) {
78 if (existing.count !== count) {
79 throw new UnsupportedOnGpu(
80 `'${name}' is ${existing.count} elements in one program and ` +
81 `${count} in another`,
82 );
83 }
84 return existing;
85 }
86 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
87 this.#slots.set(name, slot);
88 return slot;
89 }
91 get(name: string): Slot | undefined {
92 return this.#slots.get(name);
93 }
95 /** Upload initial data for a host binding. */
96 upload(name: string, data: Float32Array): void {
97 const slot = this.#slots.get(name);
98 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
99 if (data.length !== slot.count) {
100 throw new Error(
101 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
102 );
103 }
104 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
105 }
107 destroy(): void {
108 for (const s of this.#slots.values()) s.buffer.destroy();
109 this.#slots.clear();
110 }
111}
113type Op =
114 | {
115 kind: 'kernel';
116 pipeline: GPUComputePipeline;
117 bindGroup: GPUBindGroup;
118 count: number;
119 label: string;
120 /** Set when the kernel had to write to scratch because its output
121 * aliases one of its inputs; copied back after the dispatch. */
122 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
123 }
124 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 125 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 126 | { kind: 'dot'; binding: DotBinding; label: string }
127 | {
128 kind: 'copy';
129 from: GPUBuffer;
130 to: GPUBuffer;
131 bytes: number;
132 label: string;
133 /** Byte offsets, for the indexed-access ops. Absent means 0. */
134 fromOffset?: number;
135 toOffset?: number;
136 };
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 137
138export interface PlanSpec {
139 /** The specialized function this plan executes. */
140 fn: CompiledFunction;
141 /** Output index -> host binding name to copy the result into after the run,
142 * so the next call reads it (the new spectral state feeds the old). */
143 feedback: (string | null)[];
144}
146/**
147 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
148 * buffers after it, then the params buffer.
149 *
150 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
151 * only contains the bindings the shader actually references — so a kernel that
152 * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
153 * binding and no longer match the bind group. An explicit layout may carry
154 * bindings the shader ignores.
155 */
156function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
157 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
158 binding,
159 visibility: GPUShaderStage.COMPUTE,
160 buffer: { type: 'read-only-storage' },
161 });
162 return device.createBindGroupLayout({
163 entries: [
164 {
165 binding: 0,
166 visibility: GPUShaderStage.COMPUTE,
167 buffer: { type: 'storage' },
168 },
169 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
170 readOnly(inputs + 1),
171 ],
172 });
173}
175async function makePipeline(
176 device: GPUDevice,
177 code: string,
178 label: string,
179 bindGroupLayout: GPUBindGroupLayout,
180): Promise<GPUComputePipeline> {
181 device.pushErrorScope('validation');
182 const module = device.createShaderModule({ code, label });
183 const info = await module.getCompilationInfo();
184 const errors = info.messages.filter((m) => m.type === 'error');
185 if (errors.length) {
186 throw new UnsupportedOnGpu(
187 `generated WGSL failed to compile for '${label}':\n` +
188 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
189 `\n--- shader ---\n${code}`,
190 );
191 }
192 const pipeline = await device.createComputePipelineAsync({
193 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
194 compute: { module, entryPoint: 'main' },
195 label,
196 });
197 const err = await device.popErrorScope();
198 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
199 return pipeline;
200}
202/** A compiled .m step, ready to run on the GPU. */
203export class ModelPlan {
204 /** Scalar parameter names, in the order the params buffer expects them. */
205 readonly paramNames: string[];
207 #device: GPUDevice;
208 #sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 209 #deriv?: DerivPlan;
211 #owned: GPUBuffer[];
212 #paramBuf: GPUBuffer;
213 #paramData: Float32Array;
214 /** Public name -> buffer, for uploading initial state and reading results. */
215 #byName: Map<string, Slot>;
217 private constructor(init: {
218 device: GPUDevice;
219 sht: ShtPlan;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 220 deriv?: DerivPlan;
222 byName: Map<string, Slot>;
223 owned: GPUBuffer[];
224 paramBuf: GPUBuffer;
225 paramData: Float32Array;
226 paramNames: string[];
227 }) {
228 this.#device = init.device;
229 this.#sht = init.sht;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 230 this.#deriv = init.deriv;
232 this.#byName = init.byName;
233 this.#owned = init.owned;
234 this.#paramBuf = init.paramBuf;
235 this.#paramData = init.paramData;
236 this.paramNames = init.paramNames;
237 }
239 static async create(
240 device: GPUDevice,
241 sht: ShtPlan,
242 spec: PlanSpec,
243 host: HostBuffers,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 244 /** Computes dtheta/dphi — only needed if the .m calls them. */
245 deriv?: DerivPlan,
247 const { fn } = spec;
249 const slots = new Map<string, Slot>();
250 const byName = new Map<string, Slot>();
251 const owned: GPUBuffer[] = [];
252 /** Scalars the .m computes from its parameters, by cName. */
253 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
255 const alloc = (label: string, count: number): Slot => {
256 const buffer = makeBuffer(device, label, count);
257 owned.push(buffer);
258 return { buffer, count };
259 };
261 // Arguments, bound by what the function's signature declares. Array
262 // arguments come from the shared pool, so a value one function returns is
263 // the same buffer the next one reads. Scalar parameters share one small
264 // storage buffer, in signature order.
265 const paramNames: string[] = [];
266 const paramSlots = new Map<string, number>();
267 for (const p of fn.params) {
268 if (p.binding.kind === 'tensor') {
269 const count = p.binding.shape.reduce((x, y) => x * y, 1);
270 const slot = host.ensure(p.name, count);
271 slots.set(p.cName, slot);
272 byName.set(p.name, slot);
273 } else if (p.binding.kind === 'param') {
274 paramSlots.set(p.cName, paramNames.length);
275 paramNames.push(p.name);
276 }
277 // `const` arguments are exact in the IR and fold into the kernels.
278 }
279 const paramData = new Float32Array(Math.max(1, paramNames.length));
280 const paramBuf = device.createBuffer({
281 label: 'mgpu-params',
282 size: 4 * paramData.length,
283 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
284 });
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 286 /** Built on first use — only a model that calls `dot` pays for it. */
287 let reduce: ReducePlan | null = null;
289 /** Does this expression read any GPU-resident value? Decides whether a
290 * scalar assignment can stay a compile-time derived scalar or needs a
291 * 1-element kernel. Plan-order matters and is correct: a name is
292 * buffer-backed from the statement that first computes it into one. */
293 const readsBufferValue = (e: IRExpr): boolean => {
294 let found = false;
295 collectVars(e, (v) => {
296 if (slots.has(v.cName)) found = true;
297 });
298 return found;
299 };
301 /**
302 * The value a scalar expression has *at this point in the plan*, if it
303 * is decidable. A literal carries its own; a variable carries one via
304 * numbl's `exact` lattice or — the case the lattice cannot see — via its
305 * derived-scalar binding, which is how an unrolled loop's variable (and
306 * anything computed from it, like an index or an inner loop bound)
307 * resolves to that iteration's literal. A buffer-backed name is a
308 * runtime value and never resolves.
309 */
310 const planTimeValue = (e: IRExpr): number | undefined => {
311 if (e.kind === 'NumLit') return e.value;
312 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
313 switch (e.kind) {
314 case 'Var': {
315 if (slots.has(e.cName)) return undefined;
316 const d = derivedScalars.get(e.cName);
317 return d ? planTimeValue(d.expr) : undefined;
318 }
319 case 'Binary': {
320 const op = PLAN_BINOPS[e.builtin];
321 if (!op) return undefined;
322 const l = planTimeValue(e.left);
323 const r = planTimeValue(e.right);
324 return l === undefined || r === undefined ? undefined : op(l, r);
325 }
326 case 'Unary': {
327 const v = planTimeValue(e.operand);
328 if (v === undefined) return undefined;
329 if (e.builtin === 'uminus') return -v;
330 if (e.builtin === 'uplus') return v;
331 return undefined;
332 }
333 default:
334 return undefined;
335 }
336 };
338 /** A plan-time index: integral and 1-based. */
339 const planTimeIndex = (e: IRExpr, what: string, span: unknown): number => {
340 const v = planTimeValue(e);
341 if (v === undefined) {
342 throw new UnsupportedOnGpu(
343 `${what} must be known when the model compiles — a literal, a fixed ` +
344 `argument, or a value of the unrolled loop's variable`,
345 span,
346 );
347 }
348 if (!Number.isInteger(v) || v < 1) {
349 throw new UnsupportedOnGpu(`${what} must be a positive integer (got ${v})`, span);
350 }
351 return v;
352 };
355 for (const stmt of fn.body) {
356 await planStatement(stmt);
357 }
359 // Feed declared outputs back into the argument buffers they replace.
360 fn.outputs.forEach((out, i) => {
361 const to = spec.feedback[i];
362 if (!to) return;
363 const src = slots.get(out.cName);
364 const dst = host.get(to);
365 if (!src) {
366 throw new UnsupportedOnGpu(
367 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
368 );
369 }
370 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
371 if (src.count !== dst.count) {
372 throw new UnsupportedOnGpu(
373 `'${out.name}' (${src.count} elements) cannot feed ` +
374 `'${to}' (${dst.count})`,
375 );
376 }
377 ops.push({
378 kind: 'copy',
379 from: src.buffer,
380 to: dst.buffer,
381 bytes: 4 * src.count,
382 label: `${out.name} -> ${to}`,
383 });
384 });
386 return new ModelPlan({
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 387 device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
390 async function planStatement(stmt: IRStmt): Promise<void> {
391 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
392 if (stmt.kind === 'For') return planFor(stmt);
393 if (stmt.kind !== 'Assign') {
394 throw new UnsupportedOnGpu(
395 `a model function body may only contain assignments ` +
396 `(found '${stmt.kind}')`,
397 stmt.span,
398 );
399 }
400 if (!isNumeric(stmt.ty)) {
401 throw new UnsupportedOnGpu(
402 `'${stmt.name}' is not a numeric value`,
403 stmt.span,
404 );
405 }
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 406 const ext = externalCall(stmt);
407 if (!isTensor(stmt.ty) && !ext && !readsBufferValue(stmt.expr)) {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 408 // A scalar the model derives from its parameters (`us = a + b`). It
409 // gets no buffer and no dispatch: the kernels that read it bind it as
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 410 // a `let` in their prologue. A scalar computed from GPU-resident
411 // values (a `dot` result, or anything downstream of one) instead
412 // falls through to a 1-element kernel, because its inputs live in
413 // buffers the CPU never sees.
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 414 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
415 return;
416 }
417 const count = numel(stmt.ty);
419 // Reuse the destination buffer across steps: the same cName always maps
420 // to the same buffer, so a step allocates nothing.
421 let dest = slots.get(stmt.cName);
422 if (!dest) {
423 dest = alloc(`mgpu-${stmt.name}`, count);
424 slots.set(stmt.cName, dest);
425 } else if (dest.count !== count) {
426 throw new UnsupportedOnGpu(
427 `'${stmt.name}' changes size between assignments`,
428 stmt.span,
429 );
430 }
431 byName.set(stmt.name, dest);
433 if (ext) {
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 434 // Lazy per-argument resolution: buffer arguments must have slots,
435 // while index arguments are plan-time scalars with no buffer at all.
436 const argSlot = (i: number): Slot => {
437 const a = ext.args[i];
438 if (a.kind !== 'Var') {
439 throw new UnsupportedOnGpu(
440 `'${ext.name}' needs a plain variable here — assign the ` +
441 `expression to a variable first`,
442 stmt.span,
443 );
444 }
445 const s = slots.get(a.cName);
446 if (!s) {
447 throw new UnsupportedOnGpu(
448 `'${ext.name}' reads '${a.name}', which has no buffer`,
449 stmt.span,
450 );
451 }
452 return s;
453 };
454 const label = `${stmt.name} = ${ext.name}(${ext.args.map(extArgName).join(', ')})`;
455 if (ext.name === 'dot') {
456 const a = argSlot(0);
457 const b = argSlot(1);
458 if (a.count !== b.count) {
459 throw new UnsupportedOnGpu(
460 `'dot' needs equal-length arguments (${a.count} vs ${b.count})`,
461 stmt.span,
462 );
463 }
464 if (a.buffer === dest.buffer || b.buffer === dest.buffer) {
465 throw new UnsupportedOnGpu(
466 `'dot' cannot write over one of its own arguments`,
467 stmt.span,
468 );
469 }
470 reduce ??= new ReducePlan(device);
471 ops.push({
472 kind: 'dot',
473 binding: await reduce.createDotBinding(a.buffer, b.buffer, dest.buffer, a.count),
474 label,
475 });
476 return;
477 }
478 if (ext.name === 'getslab' || ext.name === 'setslab') {
479 const slabElems = 2 * sht.nlm;
480 const bank = argSlot(0);
481 const nslabs = Math.floor(bank.count / slabElems);
482 const kArg = ext.args[ext.name === 'getslab' ? 1 : 2];
483 const k = planTimeIndex(kArg, `'${ext.name}'s index '${extArgName(kArg)}'`, stmt.span);
484 if (bank.count % slabElems !== 0 || k > nslabs) {
485 throw new UnsupportedOnGpu(
486 `'${ext.name}': slab ${k} is out of range for a bank of ` +
487 `${nslabs} spectral fields`,
488 stmt.span,
489 );
490 }
491 const slabBytes = 4 * slabElems;
492 if (ext.name === 'getslab') {
493 if (dest.count !== slabElems || bank.buffer === dest.buffer) {
494 throw new UnsupportedOnGpu(`'getslab' cannot read into its own bank`, stmt.span);
495 }
496 ops.push({
497 kind: 'copy', from: bank.buffer, fromOffset: (k - 1) * slabBytes,
498 to: dest.buffer, bytes: slabBytes, label,
499 });
500 } else {
501 const field = argSlot(1);
502 if (field.count !== slabElems || field.buffer === dest.buffer) {
503 throw new UnsupportedOnGpu(
504 `'setslab' needs a distinct 2 x nlm field to write`,
505 stmt.span,
506 );
507 }
508 // Functional update: writing back over the base is the in-place
509 // fast path; a fresh destination first takes a copy of the bank.
510 if (dest.buffer !== bank.buffer) {
511 ops.push({
512 kind: 'copy', from: bank.buffer, to: dest.buffer,
513 bytes: 4 * bank.count, label: `${label} (bank copy)`,
514 });
515 }
516 ops.push({
517 kind: 'copy', from: field.buffer,
518 to: dest.buffer, toOffset: (k - 1) * slabBytes,
519 bytes: slabBytes, label,
520 });
521 }
522 return;
523 }
524 if (ext.name === 'getat' || ext.name === 'setat') {
525 const base = argSlot(0);
526 const baseTy = ext.args[0].ty;
527 if (ext.args[0].kind !== 'Var') {
528 throw new UnsupportedOnGpu(`'${ext.name}' needs a variable base`, stmt.span);
529 }
530 const shape = isNumeric(baseTy) ? baseTy.shape : undefined;
531 if (!shape) {
532 throw new UnsupportedOnGpu(`'${ext.name}' needs a base of known shape`, stmt.span);
533 }
534 const idxArgs = ext.args.slice(ext.name === 'getat' ? 1 : 2);
535 const idx = idxArgs.map(
536 (a) => planTimeIndex(a, `'${ext.name}'s index '${extArgName(a)}'`, stmt.span) - 1,
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 538 // Column-major, like everything else in the 2 x nlm layout: a
539 // 2-index access is (i-1) + (j-1)*rows, a 1-index access is linear.
540 let offset: number;
541 if (idx.length === 2) {
542 const [i, j] = idx;
543 if (i >= shape[0] || j >= (shape[1] ?? 1)) {
544 throw new UnsupportedOnGpu(
545 `'${ext.name}': (${i + 1}, ${j + 1}) is outside ` +
546 `${shape.join('x')} '${extArgName(ext.args[0])}'`,
547 stmt.span,
548 );
549 }
550 offset = i + j * shape[0];
551 } else {
552 offset = idx[0];
553 if (offset >= base.count) {
554 throw new UnsupportedOnGpu(
555 `'${ext.name}': index ${offset + 1} is outside ` +
556 `${base.count}-element '${extArgName(ext.args[0])}'`,
557 stmt.span,
558 );
559 }
560 }
561 if (ext.name === 'getat') {
562 if (dest.count !== 1 || base.buffer === dest.buffer) {
563 throw new UnsupportedOnGpu(`'getat' cannot read into its own base`, stmt.span);
564 }
565 ops.push({
566 kind: 'copy', from: base.buffer, fromOffset: 4 * offset,
567 to: dest.buffer, bytes: 4, label,
568 });
569 } else {
570 const value = argSlot(1);
571 if (value.count !== 1 || value.buffer === dest.buffer) {
572 throw new UnsupportedOnGpu(
573 `'setat' needs a distinct 1-element value to write — compute ` +
574 `it into a variable first`,
575 stmt.span,
576 );
577 }
578 if (dest.buffer !== base.buffer) {
579 ops.push({
580 kind: 'copy', from: base.buffer, to: dest.buffer,
581 bytes: 4 * base.count, label: `${label} (base copy)`,
582 });
583 }
584 ops.push({
585 kind: 'copy', from: value.buffer,
586 to: dest.buffer, toOffset: 4 * offset, bytes: 4, label,
587 });
588 }
589 return;
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 591 const src = argSlot(0);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 592 if (ext.name === 'synth') {
593 ops.push({
594 kind: 'synth',
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 595 binding: sht.createSynthBinding(src.buffer, dest.buffer),
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 596 label,
597 });
598 } else if (ext.name === 'analys') {
599 ops.push({
600 kind: 'analys',
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 601 binding: sht.createAnalysBinding(src.buffer, dest.buffer),
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 602 label,
603 });
604 } else if (ext.name === 'dtheta' || ext.name === 'dphi') {
605 if (!deriv) {
606 throw new UnsupportedOnGpu(
607 `'${ext.name}' needs the surface's derivative transforms, ` +
608 `which this plan was not given`,
609 stmt.span,
610 );
611 }
612 ops.push(
613 ext.name === 'dtheta'
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 614 ? { kind: 'dtheta', binding: deriv.createDthetaBinding(src.buffer, dest.buffer), label }
615 : { kind: 'dphi', binding: deriv.createDphiBinding(src.buffer, dest.buffer), label },
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 616 );
617 } else {
618 throw new UnsupportedOnGpu(`unknown external op '${ext.name}'`, stmt.span);
619 }
621 }
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 623 // Element-wise kernel. Collect the distinct buffer-backed operands —
624 // multi-element tensors, plus any single-element value living in a
625 // buffer (a dot result or a scalar computed from one) — and give them
626 // dense binding slots. The kernel reads a single-element operand as
627 // `in<slot>[0]`, which is what broadcasts it across the output.
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 628 const tensors = new Map<string, number>();
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 629 collectVars(stmt.expr, (v) => {
630 if (!isTensor(v.ty) && !slots.has(v.cName)) return;
631 if (!tensors.has(v.cName)) tensors.set(v.cName, tensors.size);
634 const label = `${stmt.name} = <${count} elements, element-wise>`;
635 const kernel = buildKernel(
636 stmt,
637 {
638 tensors,
639 params: paramSlots,
640 scalars: derivedScalars,
641 } satisfies KernelInputs,
642 count,
643 label,
644 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 645
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 646 const bindGroupLayout = kernelLayout(device, tensors.size);
647 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
649 // WebGPU forbids aliasing a writable storage binding with another
650 // binding in the same group, so an in-place update (`u = u + 1`) writes
651 // to scratch and copies back. Element-wise kernels only ever touch
652 // their own index, so the copy is the only cost.
653 const aliased = tensors.has(stmt.cName);
654 const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
656 const entries: GPUBindGroupEntry[] = [
657 { binding: 0, resource: { buffer: target.buffer } },
658 ];
659 for (const [cName, i] of tensors) {
660 const s = slots.get(cName);
661 if (!s) {
662 throw new UnsupportedOnGpu(
663 `'${stmt.name}' reads a value with no buffer`,
664 stmt.span,
665 );
666 }
667 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
668 }
669 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
671 ops.push({
672 kind: 'kernel',
673 pipeline,
674 bindGroup: device.createBindGroup({
675 layout: bindGroupLayout,
676 entries,
677 }),
678 count,
679 label,
680 copyBack: aliased
681 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
682 : undefined,
683 });
684 }
686 /**
687 * Unroll a counted loop into the op sequence.
688 *
689 * A plan is a fixed list of GPU operations with no branching, which is what
690 * makes a timestep pure command recording. A `for` with compile-time-known
691 * bounds still fits that: it is the same body planned once per iteration.
692 * Nothing else changes — numbl gives a variable one cName for every
693 * assignment to it, so the buffer an iteration writes is the buffer the
694 * next one reads, which is exactly a loop-carried value.
695 *
696 * The loop variable gets no buffer either: it is bound as a derived scalar
697 * to this iteration's literal value, so a kernel that reads `k` folds the
698 * number in. The binding is overwritten per iteration, before that
699 * iteration's body is planned and its WGSL emitted.
700 */
701 async function planFor(stmt: For): Promise<void> {
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 702 const from = planTimeValue(stmt.start);
703 const to = planTimeValue(stmt.end);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 704 if (from === undefined || to === undefined) {
705 throw new UnsupportedOnGpu(
706 `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
707 `be known when the model is compiled — ` +
708 `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 709 `runtime value. Use a whole number, a count the app supplies ` +
710 `as a fixed argument (changing it recompiles), or an enclosing ` +
711 `unrolled loop's variable.`,
713 );
714 }
715 const trips = Math.floor((to - from) / stmt.step) + 1;
716 if (!Number.isFinite(trips)) {
717 throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
718 }
719 if (trips > MAX_UNROLL) {
720 throw new UnsupportedOnGpu(
721 `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
722 `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
723 `long loop compiles slowly and runs no faster than writing it out.`,
724 stmt.span,
725 );
726 }
727 for (let i = 0; i < trips; i++) {
728 const value = from + i * stmt.step;
729 derivedScalars.set(stmt.cVar, {
730 name: stmt.varName,
731 expr: {
732 kind: 'NumLit',
733 value,
734 ty: scalarDouble(
735 value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
736 value,
737 ),
738 span: stmt.span,
739 },
740 });
741 for (const s of stmt.body) await planStatement(s);
742 }
743 }
744 }
746 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
747 setParams(values: Record<string, number>): void {
748 this.paramNames.forEach((name, i) => {
749 const v = values[name];
750 this.#paramData[i] = Number.isFinite(v) ? v : 0;
751 });
752 this.#device.queue.writeBuffer(
753 this.#paramBuf,
754 0,
755 this.#paramData as Float32Array<ArrayBuffer>,
756 );
757 }
759 /** Buffer holding the named value, or undefined if the .m never binds it. */
760 buffer(name: string): GPUBuffer | undefined {
761 return this.#byName.get(name)?.buffer;
762 }
764 elementCount(name: string): number | undefined {
765 return this.#byName.get(name)?.count;
766 }
768 /**
769 * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
770 * ops share one compute pass, which WebGPU executes in submission order
771 * with a barrier between dispatches.
772 */
773 encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
774 for (let s = 0; s < steps; s++) {
775 let pass: GPUComputePassEncoder | null = null;
776 const inPass = (): GPUComputePassEncoder => {
777 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
778 return pass;
779 };
780 const endPass = (): void => {
781 if (pass) {
782 pass.end();
783 pass = null;
784 }
785 };
786 for (const op of this.#ops) {
787 switch (op.kind) {
788 case 'kernel': {
789 const p = inPass();
790 p.setPipeline(op.pipeline);
791 p.setBindGroup(0, op.bindGroup);
792 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
793 if (op.copyBack) {
794 endPass();
795 encoder.copyBufferToBuffer(
796 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
797 );
798 }
799 break;
800 }
801 case 'synth':
802 this.#shtInto(inPass(), op);
803 break;
804 case 'analys':
805 this.#shtInto(inPass(), op);
806 break;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 807 case 'dtheta':
808 this.#derivInto(inPass(), op);
809 break;
810 case 'dphi':
811 this.#derivInto(inPass(), op);
812 break;
814 const p = inPass();
815 p.setPipeline(op.binding.pipeline);
816 p.setBindGroup(0, op.binding.bindGroup);
817 p.dispatchWorkgroups(1);
818 break;
819 }
821 endPass();
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 822 encoder.copyBufferToBuffer(
823 op.from, op.fromOffset ?? 0, op.to, op.toOffset ?? 0, op.bytes,
824 );
826 }
827 }
828 endPass();
829 }
830 }
832 #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
833 if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
834 else this.#sht.encodeAnalysInto(pass, op.binding);
835 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 837 #derivInto(pass: GPUComputePassEncoder, op: Op & { kind: 'dtheta' | 'dphi' }): void {
838 // planStatement already refused to plan a dtheta/dphi op without a
839 // DerivPlan, so #deriv is guaranteed set whenever an op of this kind exists.
840 if (op.kind === 'dtheta') this.#deriv!.encodeDthetaInto(pass, op.binding);
841 else this.#deriv!.encodeDphiInto(pass, op.binding);
842 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 844 /** Human-readable op sequence — what the .m actually compiled to. */
845 describe(): string[] {
846 return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
847 }
849 destroy(): void {
850 for (const b of this.#owned) b.destroy();
851 this.#paramBuf.destroy();
852 this.#owned.length = 0;
853 }
854}
857 * `x = synth(y)` / `x = dot(y, z)` -> the call's name and arguments. A
858 * buffer argument must be a plain variable (an expression would need its own
859 * buffer, which is exactly what writing it on its own line provides — the
860 * per-argument check is in the planner); an index argument may be any
861 * expression the plan can evaluate (`j + 1`).
862 */
863function externalCall(stmt: Assign): { name: string; args: IRExpr[] } | null {
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 865 if (e.kind !== 'Call') return null;
866 const arity = EXTERNAL_OPS.get(e.name);
867 if (!arity) return null;
868 if (e.args.length < arity.minArgs || e.args.length > arity.maxArgs) {
869 const want =
870 arity.minArgs === arity.maxArgs
871 ? `${arity.minArgs}`
872 : `${arity.minArgs} to ${arity.maxArgs}`;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 873 throw new UnsupportedOnGpu(
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 874 `'${e.name}' takes ${want} argument${arity.maxArgs === 1 ? '' : 's'}`,
876 );
877 }
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 878 return { name: e.name, args: e.args };
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 881const extArgName = (a: IRExpr): string =>
882 a.kind === 'Var' ? a.name : a.kind === 'NumLit' ? String(a.value) : '<expression>';
884function collectVars(
885 e: IRExpr,
886 visit: (v: Extract<IRExpr, { kind: 'Var' }>) => void,
887): void {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 888 const walk = (x: IRExpr): void => {
889 switch (x.kind) {
890 case 'Var':
893 case 'Binary':
894 walk(x.left);
895 walk(x.right);
896 return;
897 case 'Unary':
898 walk(x.operand);
899 return;
900 case 'Call':
901 x.args.forEach(walk);
902 return;
903 default:
904 return;
905 }
906 };
907 walk(e);
908}