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 { 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';
15import type { CompiledFunction } from './compile.ts';
16import { EXTERNAL_OPS } from './externals.ts';
17import {
18 buildKernel,
19 UnsupportedOnGpu,
20 WORKGROUP_SIZE,
21 type KernelInputs,
22} from './wgsl.ts';
24const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
25const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
26const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
28/**
29 * The compile-time value of a scalar expression, if it has one. A literal
30 * carries its own; a variable carries one when it was bound to a `const` (the
31 * host's fixed scalars) or computed from constants, because numbl propagates
32 * `exact` through the type lattice.
33 */
34const exactValue = (e: IRExpr): number | undefined => {
35 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
36 return e.kind === 'NumLit' ? e.value : undefined;
37};
39/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
40 * its own pipelines at compile time and its own dispatches per step — so a
41 * runaway bound should be a clear error rather than a hang. */
42const MAX_UNROLL = 64;
44interface Slot {
45 buffer: GPUBuffer;
46 count: number;
47}
49const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
50 device.createBuffer({
51 label,
52 size: 4 * count,
53 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
54 });
56/**
57 * Buffers for host-bound variables, shared across plans.
58 *
59 * A model is two programs — `init` and `step` — compiled separately but
60 * operating on the same state. `U` in the step must be the very buffer `init`
61 * wrote, so the buffers for host bindings live here rather than inside either
62 * plan.
63 */
64export class HostBuffers {
65 #device: GPUDevice;
66 #slots = new Map<string, Slot>();
68 constructor(device: GPUDevice) {
69 this.#device = device;
70 }
72 ensure(name: string, count: number): Slot {
73 const existing = this.#slots.get(name);
74 if (existing) {
75 if (existing.count !== count) {
76 throw new UnsupportedOnGpu(
77 `'${name}' is ${existing.count} elements in one program and ` +
78 `${count} in another`,
79 );
80 }
81 return existing;
82 }
83 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
84 this.#slots.set(name, slot);
85 return slot;
86 }
88 get(name: string): Slot | undefined {
89 return this.#slots.get(name);
90 }
92 /** Upload initial data for a host binding. */
93 upload(name: string, data: Float32Array): void {
94 const slot = this.#slots.get(name);
95 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
96 if (data.length !== slot.count) {
97 throw new Error(
98 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
99 );
100 }
101 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
102 }
104 destroy(): void {
105 for (const s of this.#slots.values()) s.buffer.destroy();
106 this.#slots.clear();
107 }
108}
110type Op =
111 | {
112 kind: 'kernel';
113 pipeline: GPUComputePipeline;
114 bindGroup: GPUBindGroup;
115 count: number;
116 label: string;
117 /** Set when the kernel had to write to scratch because its output
118 * aliases one of its inputs; copied back after the dispatch. */
119 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
120 }
121 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
122 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
124export interface PlanSpec {
125 /** The specialized function this plan executes. */
126 fn: CompiledFunction;
127 /** Output index -> host binding name to copy the result into after the run,
128 * so the next call reads it (the new spectral state feeds the old). */
129 feedback: (string | null)[];
130}
132/**
133 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
134 * buffers after it, then the params buffer.
135 *
136 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
137 * only contains the bindings the shader actually references — so a kernel that
138 * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
139 * binding and no longer match the bind group. An explicit layout may carry
140 * bindings the shader ignores.
141 */
142function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
143 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
144 binding,
145 visibility: GPUShaderStage.COMPUTE,
146 buffer: { type: 'read-only-storage' },
147 });
148 return device.createBindGroupLayout({
149 entries: [
150 {
151 binding: 0,
152 visibility: GPUShaderStage.COMPUTE,
153 buffer: { type: 'storage' },
154 },
155 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
156 readOnly(inputs + 1),
157 ],
158 });
159}
161async function makePipeline(
162 device: GPUDevice,
163 code: string,
164 label: string,
165 bindGroupLayout: GPUBindGroupLayout,
166): Promise<GPUComputePipeline> {
167 device.pushErrorScope('validation');
168 const module = device.createShaderModule({ code, label });
169 const info = await module.getCompilationInfo();
170 const errors = info.messages.filter((m) => m.type === 'error');
171 if (errors.length) {
172 throw new UnsupportedOnGpu(
173 `generated WGSL failed to compile for '${label}':\n` +
174 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
175 `\n--- shader ---\n${code}`,
176 );
177 }
178 const pipeline = await device.createComputePipelineAsync({
179 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
180 compute: { module, entryPoint: 'main' },
181 label,
182 });
183 const err = await device.popErrorScope();
184 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
185 return pipeline;
186}
188/** A compiled .m step, ready to run on the GPU. */
189export class ModelPlan {
190 /** Scalar parameter names, in the order the params buffer expects them. */
191 readonly paramNames: string[];
193 #device: GPUDevice;
194 #sht: ShtPlan;
195 #ops: Op[];
196 #owned: GPUBuffer[];
197 #paramBuf: GPUBuffer;
198 #paramData: Float32Array;
199 /** Public name -> buffer, for uploading initial state and reading results. */
200 #byName: Map<string, Slot>;
202 private constructor(init: {
203 device: GPUDevice;
204 sht: ShtPlan;
205 ops: Op[];
206 byName: Map<string, Slot>;
207 owned: GPUBuffer[];
208 paramBuf: GPUBuffer;
209 paramData: Float32Array;
210 paramNames: string[];
211 }) {
212 this.#device = init.device;
213 this.#sht = init.sht;
214 this.#ops = init.ops;
215 this.#byName = init.byName;
216 this.#owned = init.owned;
217 this.#paramBuf = init.paramBuf;
218 this.#paramData = init.paramData;
219 this.paramNames = init.paramNames;
220 }
222 static async create(
223 device: GPUDevice,
224 sht: ShtPlan,
225 spec: PlanSpec,
226 host: HostBuffers,
227 ): Promise<ModelPlan> {
228 const { fn } = spec;
230 const slots = new Map<string, Slot>();
231 const byName = new Map<string, Slot>();
232 const owned: GPUBuffer[] = [];
233 /** Scalars the .m computes from its parameters, by cName. */
234 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
236 const alloc = (label: string, count: number): Slot => {
237 const buffer = makeBuffer(device, label, count);
238 owned.push(buffer);
239 return { buffer, count };
240 };
242 // Arguments, bound by what the function's signature declares. Array
243 // arguments come from the shared pool, so a value one function returns is
244 // the same buffer the next one reads. Scalar parameters share one small
245 // storage buffer, in signature order.
246 const paramNames: string[] = [];
247 const paramSlots = new Map<string, number>();
248 for (const p of fn.params) {
249 if (p.binding.kind === 'tensor') {
250 const count = p.binding.shape.reduce((x, y) => x * y, 1);
251 const slot = host.ensure(p.name, count);
252 slots.set(p.cName, slot);
253 byName.set(p.name, slot);
254 } else if (p.binding.kind === 'param') {
255 paramSlots.set(p.cName, paramNames.length);
256 paramNames.push(p.name);
257 }
258 // `const` arguments are exact in the IR and fold into the kernels.
259 }
260 const paramData = new Float32Array(Math.max(1, paramNames.length));
261 const paramBuf = device.createBuffer({
262 label: 'mgpu-params',
263 size: 4 * paramData.length,
264 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
265 });
267 const ops: Op[] = [];
268 for (const stmt of fn.body) {
269 await planStatement(stmt);
270 }
272 // Feed declared outputs back into the argument buffers they replace.
273 fn.outputs.forEach((out, i) => {
274 const to = spec.feedback[i];
275 if (!to) return;
276 const src = slots.get(out.cName);
277 const dst = host.get(to);
278 if (!src) {
279 throw new UnsupportedOnGpu(
280 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
281 );
282 }
283 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
284 if (src.count !== dst.count) {
285 throw new UnsupportedOnGpu(
286 `'${out.name}' (${src.count} elements) cannot feed ` +
287 `'${to}' (${dst.count})`,
288 );
289 }
290 ops.push({
291 kind: 'copy',
292 from: src.buffer,
293 to: dst.buffer,
294 bytes: 4 * src.count,
295 label: `${out.name} -> ${to}`,
296 });
297 });
299 return new ModelPlan({
300 device, sht, ops, byName, owned, paramBuf, paramData, paramNames,
301 });
303 async function planStatement(stmt: IRStmt): Promise<void> {
304 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
305 if (stmt.kind === 'For') return planFor(stmt);
306 if (stmt.kind !== 'Assign') {
307 throw new UnsupportedOnGpu(
308 `a model function body may only contain assignments ` +
309 `(found '${stmt.kind}')`,
310 stmt.span,
311 );
312 }
313 if (!isNumeric(stmt.ty)) {
314 throw new UnsupportedOnGpu(
315 `'${stmt.name}' is not a numeric value`,
316 stmt.span,
317 );
318 }
319 if (!isTensor(stmt.ty)) {
320 // A scalar the model derives from its parameters (`us = a + b`). It
321 // gets no buffer and no dispatch: the kernels that read it bind it as
322 // a `let` in their prologue.
323 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
324 return;
325 }
326 const count = numel(stmt.ty);
328 // Reuse the destination buffer across steps: the same cName always maps
329 // to the same buffer, so a step allocates nothing.
330 let dest = slots.get(stmt.cName);
331 if (!dest) {
332 dest = alloc(`mgpu-${stmt.name}`, count);
333 slots.set(stmt.cName, dest);
334 } else if (dest.count !== count) {
335 throw new UnsupportedOnGpu(
336 `'${stmt.name}' changes size between assignments`,
337 stmt.span,
338 );
339 }
340 byName.set(stmt.name, dest);
342 const ext = externalCall(stmt);
343 if (ext) {
344 const argSlot = slots.get(ext.argCName);
345 if (!argSlot) {
346 throw new UnsupportedOnGpu(
347 `'${ext.name}' reads '${ext.argName}', which has no buffer`,
348 stmt.span,
349 );
350 }
351 ops.push(
352 ext.name === 'synth'
353 ? {
354 kind: 'synth',
355 binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
356 label: `${stmt.name} = synth(${ext.argName})`,
357 }
358 : {
359 kind: 'analys',
360 binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
361 label: `${stmt.name} = analys(${ext.argName})`,
362 },
363 );
364 return;
365 }
367 // Element-wise kernel. Collect the distinct tensor operands and give
368 // them dense binding slots.
369 const tensors = new Map<string, number>();
370 collectTensorVars(stmt.expr, (cName) => {
371 if (!tensors.has(cName)) tensors.set(cName, tensors.size);
372 });
374 const label = `${stmt.name} = <${count} elements, element-wise>`;
375 const kernel = buildKernel(
376 stmt,
377 {
378 tensors,
379 params: paramSlots,
380 scalars: derivedScalars,
381 } satisfies KernelInputs,
382 count,
383 label,
384 );
385 const bindGroupLayout = kernelLayout(device, tensors.size);
386 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
388 // WebGPU forbids aliasing a writable storage binding with another
389 // binding in the same group, so an in-place update (`u = u + 1`) writes
390 // to scratch and copies back. Element-wise kernels only ever touch
391 // their own index, so the copy is the only cost.
392 const aliased = tensors.has(stmt.cName);
393 const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
395 const entries: GPUBindGroupEntry[] = [
396 { binding: 0, resource: { buffer: target.buffer } },
397 ];
398 for (const [cName, i] of tensors) {
399 const s = slots.get(cName);
400 if (!s) {
401 throw new UnsupportedOnGpu(
402 `'${stmt.name}' reads a value with no buffer`,
403 stmt.span,
404 );
405 }
406 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
407 }
408 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
410 ops.push({
411 kind: 'kernel',
412 pipeline,
413 bindGroup: device.createBindGroup({
414 layout: bindGroupLayout,
415 entries,
416 }),
417 count,
418 label,
419 copyBack: aliased
420 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
421 : undefined,
422 });
423 }
425 /**
426 * Unroll a counted loop into the op sequence.
427 *
428 * A plan is a fixed list of GPU operations with no branching, which is what
429 * makes a timestep pure command recording. A `for` with compile-time-known
430 * bounds still fits that: it is the same body planned once per iteration.
431 * Nothing else changes — numbl gives a variable one cName for every
432 * assignment to it, so the buffer an iteration writes is the buffer the
433 * next one reads, which is exactly a loop-carried value.
434 *
435 * The loop variable gets no buffer either: it is bound as a derived scalar
436 * to this iteration's literal value, so a kernel that reads `k` folds the
437 * number in. The binding is overwritten per iteration, before that
438 * iteration's body is planned and its WGSL emitted.
439 */
440 async function planFor(stmt: For): Promise<void> {
441 const from = exactValue(stmt.start);
442 const to = exactValue(stmt.end);
443 if (from === undefined || to === undefined) {
444 throw new UnsupportedOnGpu(
445 `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
446 `be known when the model is compiled — ` +
447 `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
448 `runtime value. Use a whole number, or a count the app supplies ` +
449 `as a fixed argument (changing it recompiles).`,
450 stmt.span,
451 );
452 }
453 const trips = Math.floor((to - from) / stmt.step) + 1;
454 if (!Number.isFinite(trips)) {
455 throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
456 }
457 if (trips > MAX_UNROLL) {
458 throw new UnsupportedOnGpu(
459 `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
460 `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
461 `long loop compiles slowly and runs no faster than writing it out.`,
462 stmt.span,
463 );
464 }
465 for (let i = 0; i < trips; i++) {
466 const value = from + i * stmt.step;
467 derivedScalars.set(stmt.cVar, {
468 name: stmt.varName,
469 expr: {
470 kind: 'NumLit',
471 value,
472 ty: scalarDouble(
473 value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
474 value,
475 ),
476 span: stmt.span,
477 },
478 });
479 for (const s of stmt.body) await planStatement(s);
480 }
481 }
482 }
484 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
485 setParams(values: Record<string, number>): void {
486 this.paramNames.forEach((name, i) => {
487 const v = values[name];
488 this.#paramData[i] = Number.isFinite(v) ? v : 0;
489 });
490 this.#device.queue.writeBuffer(
491 this.#paramBuf,
492 0,
493 this.#paramData as Float32Array<ArrayBuffer>,
494 );
495 }
497 /** Buffer holding the named value, or undefined if the .m never binds it. */
498 buffer(name: string): GPUBuffer | undefined {
499 return this.#byName.get(name)?.buffer;
500 }
502 elementCount(name: string): number | undefined {
503 return this.#byName.get(name)?.count;
504 }
506 /**
507 * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
508 * ops share one compute pass, which WebGPU executes in submission order
509 * with a barrier between dispatches.
510 */
511 encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
512 for (let s = 0; s < steps; s++) {
513 let pass: GPUComputePassEncoder | null = null;
514 const inPass = (): GPUComputePassEncoder => {
515 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
516 return pass;
517 };
518 const endPass = (): void => {
519 if (pass) {
520 pass.end();
521 pass = null;
522 }
523 };
524 for (const op of this.#ops) {
525 switch (op.kind) {
526 case 'kernel': {
527 const p = inPass();
528 p.setPipeline(op.pipeline);
529 p.setBindGroup(0, op.bindGroup);
530 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
531 if (op.copyBack) {
532 endPass();
533 encoder.copyBufferToBuffer(
534 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
535 );
536 }
537 break;
538 }
539 case 'synth':
540 this.#shtInto(inPass(), op);
541 break;
542 case 'analys':
543 this.#shtInto(inPass(), op);
544 break;
545 case 'copy':
546 endPass();
547 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
548 break;
549 }
550 }
551 endPass();
552 }
553 }
555 #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
556 if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
557 else this.#sht.encodeAnalysInto(pass, op.binding);
558 }
560 /** Human-readable op sequence — what the .m actually compiled to. */
561 describe(): string[] {
562 return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
563 }
565 destroy(): void {
566 for (const b of this.#owned) b.destroy();
567 this.#paramBuf.destroy();
568 this.#owned.length = 0;
569 }
570}
572/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
573function externalCall(
574 stmt: Assign,
575): { name: string; argCName: string; argName: string } | null {
576 const e = stmt.expr;
577 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
578 if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
579 throw new UnsupportedOnGpu(
580 `'${e.name}' must be applied to a single variable`,
581 stmt.span,
582 );
583 }
584 const arg = e.args[0];
585 return { name: e.name, argCName: arg.cName, argName: arg.name };
586}
588function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
589 const walk = (x: IRExpr): void => {
590 switch (x.kind) {
591 case 'Var':
592 if (isTensor(x.ty)) visit(x.cName);
593 return;
594 case 'Binary':
595 walk(x.left);
596 walk(x.right);
597 return;
598 case 'Unary':
599 walk(x.operand);
600 return;
601 case 'Call':
602 x.args.forEach(walk);
603 return;
604 default:
605 return;
606 }
607 };
608 walk(e);
609}