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 * `encodeSteps` is pure synchronous command recording, with no allocation, no
8 * pipeline lookup and no readback. That is what lets a whole batch of
9 * timesteps be 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 type { CompiledFunction } from './compile.ts';
15import { EXTERNAL_OPS } from './externals.ts';
16import { StencilPlan, type StencilKind } from './stencil.ts';
17import { kernelOperandBudget } from '../device.ts';
18import {
19 buildKernel,
20 checkShapes,
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);
30interface Slot {
31 buffer: GPUBuffer;
32 count: number;
33}
35const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
36 device.createBuffer({
37 label,
38 size: 4 * count,
39 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
40 });
42/**
43 * Buffers for host-bound variables, shared across plans.
44 *
45 * A model is two programs — `init` and `step` — compiled separately but
46 * operating on the same state. `p` in the step must be the very buffer `init`
47 * wrote, so the buffers for host bindings live here rather than inside either
48 * plan.
49 */
50export class HostBuffers {
51 #device: GPUDevice;
52 #slots = new Map<string, Slot>();
54 constructor(device: GPUDevice) {
55 this.#device = device;
56 }
58 ensure(name: string, count: number): Slot {
59 const existing = this.#slots.get(name);
60 if (existing) {
61 if (existing.count !== count) {
62 throw new UnsupportedOnGpu(
63 `'${name}' is ${existing.count} elements in one program and ` +
64 `${count} in another`,
65 );
66 }
67 return existing;
68 }
69 const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
70 this.#slots.set(name, slot);
71 return slot;
72 }
74 get(name: string): Slot | undefined {
75 return this.#slots.get(name);
76 }
78 /** Upload initial data for a host binding. */
79 upload(name: string, data: Float32Array): void {
80 const slot = this.#slots.get(name);
81 if (!slot) throw new Error(`upload: no buffer named '${name}'`);
82 if (data.length !== slot.count) {
83 throw new Error(
84 `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
85 );
86 }
87 this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
88 }
90 destroy(): void {
91 for (const s of this.#slots.values()) s.buffer.destroy();
92 this.#slots.clear();
93 }
94}
96type Op =
97 | {
98 kind: 'kernel';
99 pipeline: GPUComputePipeline;
100 bindGroup: GPUBindGroup;
101 count: number;
102 label: string;
103 /** Set when the kernel had to write to scratch because its output
104 * aliases one of its inputs; copied back after the dispatch. */
105 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
106 }
107 | {
108 kind: 'stencil';
109 pipeline: GPUComputePipeline;
110 bindGroup: GPUBindGroup;
111 workgroups: number;
112 label: string;
113 }
114 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
116export interface PlanSpec {
117 /** The specialized function this plan executes. */
118 fn: CompiledFunction;
119 /** Output index -> host binding name to copy the result into after the run,
120 * so the next call reads it (the new field feeds the old). */
121 feedback: (string | null)[];
122}
124/**
125 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
126 * buffers after it, then the params buffer.
127 *
128 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
129 * only contains the bindings the shader actually references — so a kernel that
130 * happens to use no parameters would drop the params binding and no longer
131 * match the bind group. An explicit layout may carry bindings the shader
132 * ignores.
133 */
134function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
135 const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
136 binding,
137 visibility: GPUShaderStage.COMPUTE,
138 buffer: { type: 'read-only-storage' },
139 });
140 return device.createBindGroupLayout({
141 entries: [
142 {
143 binding: 0,
144 visibility: GPUShaderStage.COMPUTE,
145 buffer: { type: 'storage' },
146 },
147 ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
148 readOnly(inputs + 1),
149 ],
150 });
151}
153/**
154 * Compile one shader into a pipeline.
155 *
156 * No validation error scope around it: `createComputePipelineAsync` already
157 * rejects on a shader that will not compile or a layout that does not match,
158 * which is the whole of what a scope here would have caught, and the scope
159 * costs an extra device round trip per pipeline. `getCompilationInfo`, which
160 * has the line and column within the generated WGSL, is asked for only once
161 * something has gone wrong, and defensively even then.
162 *
163 * In practice the WGSL here is generated, so a shader that fails to compile is
164 * this project's bug rather than the user's; a mistake in a .m is caught
165 * earlier, by the emitter, with a position in the MATLAB source.
166 */
167async function makePipeline(
168 device: GPUDevice,
169 code: string,
170 label: string,
171 bindGroupLayout: GPUBindGroupLayout,
172): Promise<GPUComputePipeline> {
173 const module = device.createShaderModule({ code, label });
174 try {
175 return await device.createComputePipelineAsync({
176 layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
177 compute: { module, entryPoint: 'main' },
178 label,
179 });
180 } catch (e) {
181 throw new UnsupportedOnGpu(
182 `generated WGSL failed to compile for '${label}':\n` +
183 `${await shaderErrors(module, e)}\n--- shader ---\n${code}`,
184 );
185 }
186}
188/** Per-line compile errors, if the browser will hand them over. */
189async function shaderErrors(module: GPUShaderModule, cause: unknown): Promise<string> {
190 const fallback = cause instanceof Error ? cause.message : String(cause);
191 try {
192 const info = await module.getCompilationInfo();
193 const errors = info.messages.filter((m) => m.type === 'error');
194 if (!errors.length) return fallback;
195 return errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n');
196 } catch {
197 return fallback;
198 }
199}
201/** A compiled .m function, ready to run on the GPU. */
202export class ModelPlan {
203 /** Scalar parameter names, in the order the params buffer expects them. */
204 readonly paramNames: string[];
206 #device: GPUDevice;
207 #ops: Op[];
208 #owned: GPUBuffer[];
209 #paramBuf: GPUBuffer;
210 #paramData: Float32Array;
211 /** Public name -> buffer, for uploading initial state and reading results. */
212 #byName: Map<string, Slot>;
214 private constructor(init: {
215 device: GPUDevice;
216 ops: Op[];
217 byName: Map<string, Slot>;
218 owned: GPUBuffer[];
219 paramBuf: GPUBuffer;
220 paramData: Float32Array;
221 paramNames: string[];
222 }) {
223 this.#device = init.device;
224 this.#ops = init.ops;
225 this.#byName = init.byName;
226 this.#owned = init.owned;
227 this.#paramBuf = init.paramBuf;
228 this.#paramData = init.paramData;
229 this.paramNames = init.paramNames;
230 }
232 static async create(
233 device: GPUDevice,
234 stencil: StencilPlan,
235 spec: PlanSpec,
236 host: HostBuffers,
237 /** Overrides what the device allows; only tests pass it. */
238 operandBudget?: number,
239 ): Promise<ModelPlan> {
240 const { fn } = spec;
242 const slots = new Map<string, Slot>();
243 const byName = new Map<string, Slot>();
244 const owned: GPUBuffer[] = [];
245 /** Scalars the .m computes from its parameters, by cName. */
246 const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
247 /** Grid fields one kernel may read on this device (see fitToBudget). */
248 const budget = operandBudget ?? kernelOperandBudget(device);
249 /** How many kernels a line has been split into, for naming the pieces. */
250 let splits = 0;
252 const alloc = (label: string, count: number): Slot => {
253 const buffer = makeBuffer(device, label, count);
254 owned.push(buffer);
255 return { buffer, count };
256 };
258 // Arguments, bound by what the function's signature declares. Array
259 // arguments come from the shared pool, so a value one function returns is
260 // the same buffer the next one reads. Scalar parameters share one small
261 // storage buffer, in signature order.
262 const paramNames: string[] = [];
263 const paramSlots = new Map<string, number>();
264 for (const p of fn.params) {
265 if (p.binding.kind === 'tensor') {
266 const count = p.binding.shape.reduce((x, y) => x * y, 1);
267 const slot = host.ensure(p.name, count);
268 slots.set(p.cName, slot);
269 byName.set(p.name, slot);
270 } else if (p.binding.kind === 'param') {
271 paramSlots.set(p.cName, paramNames.length);
272 paramNames.push(p.name);
273 }
274 // `const` arguments are exact in the IR and fold into the kernels.
275 }
276 const paramData = new Float32Array(Math.max(1, paramNames.length));
277 const paramBuf = device.createBuffer({
278 label: 'mgpu-params',
279 size: 4 * paramData.length,
280 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
281 });
283 const ops: Op[] = [];
284 for (const stmt of fn.body) {
285 await planStatement(stmt);
286 }
288 planFeedback();
290 return new ModelPlan({ device, ops, byName, owned, paramBuf, paramData, paramNames });
292 async function planStatement(stmt: IRStmt): Promise<void> {
293 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
294 if (stmt.kind !== 'Assign') {
295 throw new UnsupportedOnGpu(
296 `a model function body may only contain assignments ` +
297 `(found '${stmt.kind}')`,
298 stmt.span,
299 );
300 }
301 if (!isNumeric(stmt.ty)) {
302 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric value`, stmt.span);
303 }
304 if (!isTensor(stmt.ty)) {
305 // A scalar the model derives from its parameters (`om = 2*pi*f`). It
306 // gets no buffer and no dispatch: the kernels that read it bind it as
307 // a `let` in their prologue.
308 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
309 return;
310 }
311 const count = numel(stmt.ty);
313 // Reuse the destination buffer across steps: the same cName always maps
314 // to the same buffer, so a step allocates nothing.
315 let dest = slots.get(stmt.cName);
316 if (!dest) {
317 dest = alloc(`mgpu-${stmt.name}`, count);
318 slots.set(stmt.cName, dest);
319 } else if (dest.count !== count) {
320 throw new UnsupportedOnGpu(
321 `'${stmt.name}' changes size between assignments`,
322 stmt.span,
323 );
324 }
325 byName.set(stmt.name, dest);
327 const ext = externalCall(stmt);
328 if (ext) return planStencil(stmt, ext, dest);
330 // Element-wise. Checked against the whole line first, so a broadcasting
331 // mistake is reported against what was written rather than against a
332 // fragment of it.
333 checkShapes(stmt.expr, stmt.ty, stmt.name);
334 const expr = await fitToBudget(stmt.expr, stmt.name, count, stmt.span);
335 await emitElementwise(stmt.name, expr, stmt.ty, stmt.span, dest, count, stmt.cName);
336 }
338 /**
339 * Emit one element-wise kernel: `expr` evaluated at every index into
340 * `dest`. `selfCName` is the variable being assigned, if any, so an
341 * in-place update can be spotted.
342 */
343 async function emitElementwise(
344 name: string,
345 expr: IRExpr,
346 ty: NumericType,
347 span: unknown,
348 dest: Slot,
349 count: number,
350 selfCName?: string,
351 ): Promise<void> {
352 // Collect the distinct tensor operands and give them dense binding slots.
353 const tensors = new Map<string, number>();
354 collectTensorVars(expr, (cName) => {
355 if (!tensors.has(cName)) tensors.set(cName, tensors.size);
356 });
358 const label = `${name} = <${count} elements, element-wise>`;
359 const kernel = buildKernel(
360 { kind: 'Assign', name, ty, expr, span } as unknown as Assign,
361 {
362 tensors,
363 params: paramSlots,
364 scalars: derivedScalars,
365 } satisfies KernelInputs,
366 count,
367 label,
368 );
370 const bindGroupLayout = kernelLayout(device, tensors.size);
371 const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
373 // WebGPU forbids aliasing a writable storage binding with another
374 // binding in the same group, so an in-place update (`p = p + 1`) writes
375 // to scratch and copies back. Element-wise kernels only ever touch
376 // their own index, so the copy is the only cost.
377 const aliased = selfCName !== undefined && tensors.has(selfCName);
378 const target = aliased ? alloc(`mgpu-${name}-scratch`, count) : dest;
380 const entries: GPUBindGroupEntry[] = [
381 { binding: 0, resource: { buffer: target.buffer } },
382 ];
383 for (const [cName, i] of tensors) {
384 const s = slots.get(cName);
385 if (!s) {
386 throw new UnsupportedOnGpu(`'${name}' reads a value with no buffer`, span);
387 }
388 entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
389 }
390 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
392 ops.push({
393 kind: 'kernel',
394 pipeline,
395 bindGroup: device.createBindGroup({ layout: bindGroupLayout, entries }),
396 count,
397 label,
398 copyBack: aliased
399 ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
400 : undefined,
401 });
402 }
404 /**
405 * Split an expression that reads more grid fields than one kernel may bind.
406 *
407 * A kernel binds one storage buffer per distinct field it reads, plus its
408 * output and the parameter block, and WebGPU guarantees only eight per
409 * compute stage — fewer in compatibility mode. numbl's inline pass, which
410 * is what makes one source line become one kernel, does not know about
411 * that limit, and a model has no way to ask it for less: a temporary used
412 * once is exactly what it folds away.
413 *
414 * So the budget is enforced here instead. Any child subtree that reads
415 * more than one field is evaluated into its own buffer and replaced by a
416 * reference to it, which leaves the parent reading at most one field per
417 * child. The result is the same arithmetic in a few more passes over
418 * memory, and it only happens on a line that would not otherwise compile.
419 */
420 async function fitToBudget(
421 expr: IRExpr,
422 hint: string,
423 count: number,
424 span: unknown,
425 ): Promise<IRExpr> {
426 if (tensorCount(expr) <= budget) return expr;
428 const fit = async (e: IRExpr): Promise<IRExpr> => {
429 if (tensorCount(e) <= budget) return e;
430 const kids = children(e);
431 if (!kids.length) return e;
432 const out: IRExpr[] = [];
433 for (const kid of kids) {
434 const fitted = await fit(kid);
435 out.push(tensorCount(fitted) > 1 ? await hoist(fitted) : fitted);
436 }
437 return withChildren(e, out);
438 };
440 /** Evaluate a subtree into its own buffer and hand back a reference. */
441 const hoist = async (e: IRExpr): Promise<IRExpr> => {
442 if (!isNumeric(e.ty) || !isTensor(e.ty)) return e;
443 const name = `${hint}_part${++splits}`;
444 const cName = `mgpu_split_${splits}`;
445 const slot = alloc(`mgpu-${name}`, count);
446 slots.set(cName, slot);
447 await emitElementwise(name, e, e.ty, e.span, slot, count);
448 return { kind: 'Var', name, cName, ty: e.ty, span: e.span } as IRExpr;
449 };
451 const fitted = await fit(expr);
452 if (tensorCount(fitted) > budget) {
453 throw new UnsupportedOnGpu(
454 `'${hint}' reads ${tensorCount(fitted)} grid fields at once, and this ` +
455 `device allows ${budget} per kernel. Compute part of it into a ` +
456 `named field on a line of its own.`,
457 span,
458 );
459 }
460 return fitted;
461 }
463 /** `lp = lap2(p)`: one stencil dispatch, src and dst distinct. */
464 async function planStencil(
465 stmt: Assign,
466 ext: { name: StencilKind; arg: IRExpr & { kind: 'Var' } },
467 dest: Slot,
468 ): Promise<void> {
469 const argSlot = slots.get(ext.arg.cName);
470 if (!argSlot) {
471 throw new UnsupportedOnGpu(
472 `'${ext.name}' reads '${ext.arg.name}', which has no buffer`,
473 stmt.span,
474 );
475 }
476 // A stencil reads its neighbours, so unlike an element-wise kernel it
477 // cannot be routed through scratch and copied back — the neighbours
478 // would already have been overwritten. WebGPU forbids the aliasing
479 // outright anyway; refuse rather than silently reroute.
480 if (argSlot.buffer === dest.buffer) {
481 throw new UnsupportedOnGpu(
482 `'${stmt.name} = ${ext.name}(${ext.arg.name})' reads and writes the ` +
483 `same buffer; assign to a new name instead`,
484 stmt.span,
485 );
486 }
487 if (dest.count !== stencil.npts) {
488 throw new UnsupportedOnGpu(
489 `'${ext.name}' produces a grid field (${stencil.npts} points), but ` +
490 `'${stmt.name}' holds ${dest.count}`,
491 stmt.span,
492 );
493 }
494 ops.push({
495 kind: 'stencil',
496 pipeline: await stencil.pipeline(ext.name),
497 bindGroup: stencil.createBinding(argSlot.buffer, dest.buffer),
498 workgroups: stencil.workgroups,
499 label: `${stmt.name} = ${ext.name}(${ext.arg.name})`,
500 });
501 }
503 /**
504 * Feed declared outputs back into the argument buffers they replace, so
505 * the next call reads what this one produced.
506 *
507 * The copies are not independent: a model whose new history field is the
508 * old current one (`function [pn, pold] = step(p, pm, ...)`) has an output
509 * whose *source* is another output's *destination*. Doing them in order
510 * would then copy the new value where the old one was wanted, silently.
511 * So any source that a previous copy overwrites is staged through scratch
512 * first — normally none, since a model that writes `pold = p;` gets its
513 * own buffer from the copy kernel that line plans to.
514 */
515 function planFeedback(): void {
516 const copies: { from: Slot; to: Slot; label: string }[] = [];
517 fn.outputs.forEach((out, i) => {
518 const to = spec.feedback[i];
519 if (!to) return;
520 const src = slots.get(out.cName);
521 const dst = host.get(to);
522 if (!src) {
523 throw new UnsupportedOnGpu(
524 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
525 );
526 }
527 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
528 if (src.count !== dst.count) {
529 throw new UnsupportedOnGpu(
530 `'${out.name}' (${src.count} elements) cannot feed '${to}' (${dst.count})`,
531 );
532 }
533 copies.push({ from: src, to: dst, label: `${out.name} -> ${to}` });
534 });
536 const written = new Set<GPUBuffer>();
537 for (const c of copies) written.add(c.to.buffer);
538 for (const c of copies) {
539 // Only a source another copy overwrites needs staging, and only if it
540 // is not that same copy's own destination (which is a no-op anyway).
541 if (c.from.buffer !== c.to.buffer && written.has(c.from.buffer)) {
542 const scratch = alloc(`mgpu-feedback-scratch`, c.from.count);
543 ops.push({
544 kind: 'copy',
545 from: c.from.buffer,
546 to: scratch.buffer,
547 bytes: 4 * c.from.count,
548 label: `${c.label} (staged)`,
549 });
550 c.from = scratch;
551 }
552 }
553 for (const c of copies) {
554 if (c.from.buffer === c.to.buffer) continue; // already in place
555 ops.push({
556 kind: 'copy',
557 from: c.from.buffer,
558 to: c.to.buffer,
559 bytes: 4 * c.from.count,
560 label: c.label,
561 });
562 }
563 }
564 }
566 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
567 setParams(values: Record<string, number>): void {
568 this.paramNames.forEach((name, i) => {
569 const v = values[name];
570 this.#paramData[i] = Number.isFinite(v) ? v : 0;
571 });
572 this.#device.queue.writeBuffer(
573 this.#paramBuf,
574 0,
575 this.#paramData as Float32Array<ArrayBuffer>,
576 );
577 }
579 /** Buffer holding the named value, or undefined if the .m never binds it. */
580 buffer(name: string): GPUBuffer | undefined {
581 return this.#byName.get(name)?.buffer;
582 }
584 elementCount(name: string): number | undefined {
585 return this.#byName.get(name)?.count;
586 }
588 /**
589 * Record `steps` passes of this plan. Synchronous: no awaits, no readback.
590 * The dispatches share one compute pass, which WebGPU executes in submission
591 * order with a barrier between them.
592 *
593 * `after` runs once per step, inside the same submission — which is what
594 * lets the microphone sample every timestep rather than every frame.
595 */
596 encodeSteps(
597 encoder: GPUCommandEncoder,
598 steps: number,
599 after?: (encoder: GPUCommandEncoder) => void,
600 ): void {
601 for (let s = 0; s < steps; s++) {
602 this.#encodeOps(encoder);
603 after?.(encoder);
604 }
605 }
607 /** Record one pass over the op sequence into `encoder`. */
608 #encodeOps(encoder: GPUCommandEncoder): void {
609 let pass: GPUComputePassEncoder | null = null;
610 const inPass = (): GPUComputePassEncoder => {
611 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
612 return pass;
613 };
614 const endPass = (): void => {
615 if (pass) {
616 pass.end();
617 pass = null;
618 }
619 };
620 for (const op of this.#ops) {
621 switch (op.kind) {
622 case 'kernel': {
623 const p = inPass();
624 p.setPipeline(op.pipeline);
625 p.setBindGroup(0, op.bindGroup);
626 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
627 if (op.copyBack) {
628 endPass();
629 encoder.copyBufferToBuffer(
630 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
631 );
632 }
633 break;
634 }
635 case 'stencil': {
636 const p = inPass();
637 p.setPipeline(op.pipeline);
638 p.setBindGroup(0, op.bindGroup);
639 p.dispatchWorkgroups(op.workgroups);
640 break;
641 }
642 case 'copy':
643 endPass();
644 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
645 break;
646 }
647 }
648 endPass();
649 }
651 /** Human-readable op sequence — what the .m actually compiled to. */
652 describe(): string[] {
653 return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
654 }
656 destroy(): void {
657 for (const b of this.#owned) b.destroy();
658 this.#paramBuf.destroy();
659 this.#owned.length = 0;
660 }
661}
663/** `lp = lap2(p)` -> the stencil's name and its argument. */
664function externalCall(
665 stmt: Assign,
666): { name: StencilKind; arg: IRExpr & { kind: 'Var' } } | null {
667 const e = stmt.expr;
668 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
669 if (e.args.length !== 1) {
670 throw new UnsupportedOnGpu(
671 `'${e.name}' must be applied to a single variable`,
672 stmt.span,
673 );
674 }
675 const arg = e.args[0];
676 if (arg.kind !== 'Var') {
677 throw new UnsupportedOnGpu(
678 `'${e.name}' must be applied to a variable, not an expression — ` +
679 `name the field first`,
680 stmt.span,
681 );
682 }
683 return { name: e.name as StencilKind, arg };
684}
686/** Distinct grid fields an expression reads — its storage-buffer cost. */
687function tensorCount(e: IRExpr): number {
688 const seen = new Set<string>();
689 collectTensorVars(e, (c) => seen.add(c));
690 return seen.size;
691}
693/** The subexpressions of a node, in evaluation order. Leaves have none. */
694function children(e: IRExpr): IRExpr[] {
695 switch (e.kind) {
696 case 'Binary':
697 return [e.left, e.right];
698 case 'Unary':
699 return [e.operand];
700 case 'Call':
701 return e.args;
702 default:
703 return [];
704 }
705}
707/** The same node with its subexpressions replaced. */
708function withChildren(e: IRExpr, kids: IRExpr[]): IRExpr {
709 switch (e.kind) {
710 case 'Binary':
711 return { ...e, left: kids[0], right: kids[1] };
712 case 'Unary':
713 return { ...e, operand: kids[0] };
714 case 'Call':
715 return { ...e, args: kids };
716 default:
717 return e;
718 }
719}
721function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
722 const walk = (x: IRExpr): void => {
723 switch (x.kind) {
724 case 'Var':
725 if (isTensor(x.ty)) visit(x.cName);
726 return;
727 case 'Binary':
728 walk(x.left);
729 walk(x.right);
730 return;
731 case 'Unary':
732 walk(x.operand);
733 return;
734 case 'Call':
735 x.args.forEach(walk);
736 return;
737 default:
738 return;
739 }
740 };
741 walk(e);
742}