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 } from 'numbl-src/numbl-core/jit/lowering/types.ts';
12import type { Assign, 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';
4166b48Tidy up after the SHTNS comparisonJeremy Magland 16import { EXTERNAL_OPS } from './externals.ts';
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);
28interface Slot {
29 buffer: GPUBuffer;
30 count: number;
31}
33const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
34 device.createBuffer({
35 label,
36 size: 4 * count,
37 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
38 });
40/**
41 * Buffers for host-bound variables, shared across plans.
42 *
43 * A model is two programs — `init` and `step` — compiled separately but
44 * operating on the same state. `U` in the step must be the very buffer `init`
45 * wrote, so the buffers for host bindings live here rather than inside either
46 * plan.
47 */
48export class HostBuffers {
49 #device: GPUDevice;
50 #slots = new Map<string, Slot>();
52 constructor(device: GPUDevice) {
53 this.#device = device;
54 }
56 ensure(name: string, count: number): Slot {
57 const existing = this.#slots.get(name);
58 if (existing) {
59 if (existing.count !== count) {
60 throw new UnsupportedOnGpu(
61 `'${name}' is ${existing.count} elements in one program and ` +
62 `${count} in another`,
63 );
64 }
65 return existing;
66 }
67 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
68 this.#slots.set(name, slot);
69 return slot;
70 }
72 get(name: string): Slot | undefined {
73 return this.#slots.get(name);
74 }
76 /** Upload initial data for a host binding. */
77 upload(name: string, data: Float32Array): void {
78 const slot = this.#slots.get(name);
79 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
80 if (data.length !== slot.count) {
81 throw new Error(
82 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
83 );
84 }
85 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
86 }
88 destroy(): void {
89 for (const s of this.#slots.values()) s.buffer.destroy();
90 this.#slots.clear();
91 }
92}
94type Op =
95 | {
96 kind: 'kernel';
97 pipeline: GPUComputePipeline;
98 bindGroup: GPUBindGroup;
99 count: number;
100 label: string;
101 /** Set when the kernel had to write to scratch because its output
102 * aliases one of its inputs; copied back after the dispatch. */
103 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
104 }
105 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
106 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
108export interface PlanSpec {
109 /** The specialized function this plan executes. */
110 fn: CompiledFunction;
111 /** Output index -> host binding name to copy the result into after the run,
112 * so the next call reads it (the new spectral state feeds the old). */
113 feedback: (string | null)[];
114}
116/**
117 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
118 * buffers after it, then the params buffer.
119 *
120 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
121 * only contains the bindings the shader actually references — so a kernel that
122 * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
123 * binding and no longer match the bind group. An explicit layout may carry
124 * bindings the shader ignores.
125 */
126function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
127 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
128 binding,
129 visibility: GPUShaderStage.COMPUTE,
130 buffer: { type: 'read-only-storage' },
131 });
132 return device.createBindGroupLayout({
133 entries: [
134 {
135 binding: 0,
136 visibility: GPUShaderStage.COMPUTE,
137 buffer: { type: 'storage' },
138 },
139 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
140 readOnly(inputs + 1),
141 ],
142 });
143}
145async function makePipeline(
146 device: GPUDevice,
147 code: string,
148 label: string,
149 bindGroupLayout: GPUBindGroupLayout,
150): Promise<GPUComputePipeline> {
151 device.pushErrorScope('validation');
152 const module = device.createShaderModule({ code, label });
153 const info = await module.getCompilationInfo();
154 const errors = info.messages.filter((m) => m.type === 'error');
155 if (errors.length) {
156 throw new UnsupportedOnGpu(
157 `generated WGSL failed to compile for '${label}':\n` +
158 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
159 `\n--- shader ---\n${code}`,
160 );
161 }
162 const pipeline = await device.createComputePipelineAsync({
163 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
164 compute: { module, entryPoint: 'main' },
165 label,
166 });
167 const err = await device.popErrorScope();
168 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
169 return pipeline;
170}
172/** A compiled .m step, ready to run on the GPU. */
173export class ModelPlan {
174 /** Scalar parameter names, in the order the params buffer expects them. */
175 readonly paramNames: string[];
177 #device: GPUDevice;
178 #sht: ShtPlan;
179 #ops: Op[];
180 #owned: GPUBuffer[];
181 #paramBuf: GPUBuffer;
182 #paramData: Float32Array;
183 /** Public name -> buffer, for uploading initial state and reading results. */
184 #byName: Map<string, Slot>;
186 private constructor(init: {
187 device: GPUDevice;
188 sht: ShtPlan;
189 ops: Op[];
190 byName: Map<string, Slot>;
191 owned: GPUBuffer[];
192 paramBuf: GPUBuffer;
193 paramData: Float32Array;
194 paramNames: string[];
195 }) {
196 this.#device = init.device;
197 this.#sht = init.sht;
198 this.#ops = init.ops;
199 this.#byName = init.byName;
200 this.#owned = init.owned;
201 this.#paramBuf = init.paramBuf;
202 this.#paramData = init.paramData;
203 this.paramNames = init.paramNames;
204 }
206 static async create(
207 device: GPUDevice,
208 sht: ShtPlan,
209 spec: PlanSpec,
210 host: HostBuffers,
211 ): Promise<ModelPlan> {
212 const { fn } = spec;
214 const slots = new Map<string, Slot>();
215 const byName = new Map<string, Slot>();
216 const owned: GPUBuffer[] = [];
217 /** Scalars the .m computes from its parameters, by cName. */
218 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
220 const alloc = (label: string, count: number): Slot => {
221 const buffer = makeBuffer(device, label, count);
222 owned.push(buffer);
223 return { buffer, count };
224 };
226 // Arguments, bound by what the function's signature declares. Array
227 // arguments come from the shared pool, so a value one function returns is
228 // the same buffer the next one reads. Scalar parameters share one small
229 // storage buffer, in signature order.
230 const paramNames: string[] = [];
231 const paramSlots = new Map<string, number>();
232 for (const p of fn.params) {
233 if (p.binding.kind === 'tensor') {
234 const count = p.binding.shape.reduce((x, y) => x * y, 1);
235 const slot = host.ensure(p.name, count);
236 slots.set(p.cName, slot);
237 byName.set(p.name, slot);
238 } else if (p.binding.kind === 'param') {
239 paramSlots.set(p.cName, paramNames.length);
240 paramNames.push(p.name);
241 }
242 // `const` arguments are exact in the IR and fold into the kernels.
243 }
244 const paramData = new Float32Array(Math.max(1, paramNames.length));
245 const paramBuf = device.createBuffer({
246 label: 'mgpu-params',
247 size: 4 * paramData.length,
248 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
249 });
251 const ops: Op[] = [];
252 for (const stmt of fn.body) {
253 await planStatement(stmt);
254 }
256 // Feed declared outputs back into the argument buffers they replace.
257 fn.outputs.forEach((out, i) => {
258 const to = spec.feedback[i];
259 if (!to) return;
260 const src = slots.get(out.cName);
261 const dst = host.get(to);
262 if (!src) {
263 throw new UnsupportedOnGpu(
264 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
265 );
266 }
267 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
268 if (src.count !== dst.count) {
269 throw new UnsupportedOnGpu(
270 `'${out.name}' (${src.count} elements) cannot feed ` +
271 `'${to}' (${dst.count})`,
272 );
273 }
274 ops.push({
275 kind: 'copy',
276 from: src.buffer,
277 to: dst.buffer,
278 bytes: 4 * src.count,
279 label: `${out.name} -> ${to}`,
280 });
281 });
283 return new ModelPlan({
284 device, sht, ops, byName, owned, paramBuf, paramData, paramNames,
285 });
287 async function planStatement(stmt: IRStmt): Promise<void> {
288 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
289 if (stmt.kind !== 'Assign') {
290 throw new UnsupportedOnGpu(
291 `a model function body may only contain assignments ` +
292 `(found '${stmt.kind}')`,
293 stmt.span,
294 );
295 }
296 if (!isNumeric(stmt.ty)) {
297 throw new UnsupportedOnGpu(
298 `'${stmt.name}' is not a numeric value`,
299 stmt.span,
300 );
301 }
302 if (!isTensor(stmt.ty)) {
303 // A scalar the model derives from its parameters (`us = a + b`). It
304 // gets no buffer and no dispatch: the kernels that read it bind it as
305 // a `let` in their prologue.
306 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
307 return;
308 }
309 const count = numel(stmt.ty);
311 // Reuse the destination buffer across steps: the same cName always maps
312 // to the same buffer, so a step allocates nothing.
313 let dest = slots.get(stmt.cName);
314 if (!dest) {
315 dest = alloc(`mgpu-${stmt.name}`, count);
316 slots.set(stmt.cName, dest);
317 } else if (dest.count !== count) {
318 throw new UnsupportedOnGpu(
319 `'${stmt.name}' changes size between assignments`,
320 stmt.span,
321 );
322 }
323 byName.set(stmt.name, dest);
325 const ext = externalCall(stmt);
326 if (ext) {
327 const argSlot = slots.get(ext.argCName);
328 if (!argSlot) {
329 throw new UnsupportedOnGpu(
330 `'${ext.name}' reads '${ext.argName}', which has no buffer`,
331 stmt.span,
332 );
333 }
334 ops.push(
335 ext.name === 'synth'
336 ? {
337 kind: 'synth',
338 binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
339 label: `${stmt.name} = synth(${ext.argName})`,
340 }
341 : {
342 kind: 'analys',
343 binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
344 label: `${stmt.name} = analys(${ext.argName})`,
345 },
346 );
347 return;
348 }
350 // Element-wise kernel. Collect the distinct tensor operands and give
351 // them dense binding slots.
352 const tensors = new Map<string, number>();
353 collectTensorVars(stmt.expr, (cName) => {
354 if (!tensors.has(cName)) tensors.set(cName, tensors.size);
355 });
357 const label = `${stmt.name} = <${count} elements, element-wise>`;
358 const kernel = buildKernel(
359 stmt,
360 {
361 tensors,
362 params: paramSlots,
363 scalars: derivedScalars,
364 } satisfies KernelInputs,
365 count,
366 label,
367 );
368 const bindGroupLayout = kernelLayout(device, tensors.size);
369 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
371 // WebGPU forbids aliasing a writable storage binding with another
372 // binding in the same group, so an in-place update (`u = u + 1`) writes
373 // to scratch and copies back. Element-wise kernels only ever touch
374 // their own index, so the copy is the only cost.
375 const aliased = tensors.has(stmt.cName);
376 const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
378 const entries: GPUBindGroupEntry[] = [
379 { binding: 0, resource: { buffer: target.buffer } },
380 ];
381 for (const [cName, i] of tensors) {
382 const s = slots.get(cName);
383 if (!s) {
384 throw new UnsupportedOnGpu(
385 `'${stmt.name}' reads a value with no buffer`,
386 stmt.span,
387 );
388 }
389 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
390 }
391 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
393 ops.push({
394 kind: 'kernel',
395 pipeline,
396 bindGroup: device.createBindGroup({
397 layout: bindGroupLayout,
398 entries,
399 }),
400 count,
401 label,
402 copyBack: aliased
403 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
404 : undefined,
405 });
406 }
407 }
409 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
410 setParams(values: Record<string, number>): void {
411 this.paramNames.forEach((name, i) => {
412 const v = values[name];
413 this.#paramData[i] = Number.isFinite(v) ? v : 0;
414 });
415 this.#device.queue.writeBuffer(
416 this.#paramBuf,
417 0,
418 this.#paramData as Float32Array<ArrayBuffer>,
419 );
420 }
422 /** Buffer holding the named value, or undefined if the .m never binds it. */
423 buffer(name: string): GPUBuffer | undefined {
424 return this.#byName.get(name)?.buffer;
425 }
427 elementCount(name: string): number | undefined {
428 return this.#byName.get(name)?.count;
429 }
431 /**
432 * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
433 * ops share one compute pass, which WebGPU executes in submission order
434 * with a barrier between dispatches.
435 */
436 encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
437 for (let s = 0; s < steps; s++) {
438 let pass: GPUComputePassEncoder | null = null;
439 const inPass = (): GPUComputePassEncoder => {
440 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
441 return pass;
442 };
443 const endPass = (): void => {
444 if (pass) {
445 pass.end();
446 pass = null;
447 }
448 };
449 for (const op of this.#ops) {
450 switch (op.kind) {
451 case 'kernel': {
452 const p = inPass();
453 p.setPipeline(op.pipeline);
454 p.setBindGroup(0, op.bindGroup);
455 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
456 if (op.copyBack) {
457 endPass();
458 encoder.copyBufferToBuffer(
459 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
460 );
461 }
462 break;
463 }
464 case 'synth':
465 this.#shtInto(inPass(), op);
466 break;
467 case 'analys':
468 this.#shtInto(inPass(), op);
469 break;
470 case 'copy':
471 endPass();
472 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
473 break;
474 }
475 }
476 endPass();
477 }
478 }
480 #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
481 if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
482 else this.#sht.encodeAnalysInto(pass, op.binding);
483 }
485 /** Human-readable op sequence — what the .m actually compiled to. */
486 describe(): string[] {
487 return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
488 }
490 destroy(): void {
491 for (const b of this.#owned) b.destroy();
492 this.#paramBuf.destroy();
493 this.#owned.length = 0;
494 }
495}
497/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
498function externalCall(
499 stmt: Assign,
500): { name: string; argCName: string; argName: string } | null {
501 const e = stmt.expr;
502 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
503 if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
504 throw new UnsupportedOnGpu(
505 `'${e.name}' must be applied to a single variable`,
506 stmt.span,
507 );
508 }
509 const arg = e.args[0];
510 return { name: e.name, argCName: arg.cName, argName: arg.name };
511}
513function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
514 const walk = (x: IRExpr): void => {
515 switch (x.kind) {
516 case 'Var':
517 if (isTensor(x.ty)) visit(x.cName);
518 return;
519 case 'Binary':
520 walk(x.left);
521 walk(x.right);
522 return;
523 case 'Unary':
524 walk(x.operand);
525 return;
526 case 'Call':
527 x.args.forEach(walk);
528 return;
529 default:
530 return;
531 }
532 };
533 walk(e);
534}