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