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 { EXTERNAL_OPS } from './externals.ts';
16import { OpPlan, type OpKind } from './ops.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: 'external';
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 external: OpPlan,
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 planExternal(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 = lapw(p, wall)`: one dispatch, sources and dst distinct. */
464 async function planExternal(
465 stmt: Assign,
466 ext: { name: OpKind; args: (IRExpr & { kind: 'Var' })[] },
467 dest: Slot,
468 ): Promise<void> {
469 const contract = external.spec(ext.name);
470 const argSlots = ext.args.map((arg, i) => {
471 const s = slots.get(arg.cName);
472 if (!s) {
473 throw new UnsupportedOnGpu(
474 `'${ext.name}' reads '${arg.name}', which has no buffer`,
475 stmt.span,
476 );
477 }
478 if (s.count !== contract.argCounts[i]) {
479 throw new UnsupportedOnGpu(
480 `'${ext.name}' wants ${contract.argCounts[i]} elements for its ` +
481 `argument ${i + 1}, but '${arg.name}' holds ${s.count}`,
482 stmt.span,
483 );
484 }
485 // An op reads its neighbours, so unlike an element-wise kernel it
486 // cannot be routed through scratch and copied back — the neighbours
487 // would already have been overwritten. WebGPU forbids the aliasing
488 // outright anyway; refuse rather than silently reroute.
489 if (s.buffer === dest.buffer) {
490 throw new UnsupportedOnGpu(
491 `'${stmt.name} = ${ext.name}(...)' reads and writes the same ` +
492 `buffer through '${arg.name}'; assign to a new name instead`,
493 stmt.span,
494 );
495 }
496 return s;
497 });
498 if (dest.count !== contract.outCount) {
499 throw new UnsupportedOnGpu(
500 `'${ext.name}' produces a ${contract.outCount}-point field, but ` +
501 `'${stmt.name}' holds ${dest.count}`,
502 stmt.span,
503 );
504 }
505 const built = await external.pipeline(ext.name);
506 ops.push({
507 kind: 'external',
508 pipeline: built.pipeline,
509 bindGroup: external.createBinding(
510 ext.name,
511 argSlots.map((s) => s.buffer),
512 dest.buffer,
513 ),
514 workgroups: external.workgroups(contract.outCount),
515 label: `${stmt.name} = ${ext.name}(${ext.args.map((a) => a.name).join(', ')})`,
516 });
517 }
519 /**
520 * Feed declared outputs back into the argument buffers they replace, so
521 * the next call reads what this one produced.
522 *
523 * The copies are not independent: a model whose new history field is the
524 * old current one (`function [pn, pold] = step(p, pm, ...)`) has an output
525 * whose *source* is another output's *destination*. Doing them in order
526 * would then copy the new value where the old one was wanted, silently.
527 * So any source that a previous copy overwrites is staged through scratch
528 * first — normally none, since a model that writes `pold = p;` gets its
529 * own buffer from the copy kernel that line plans to.
530 */
531 function planFeedback(): void {
532 const copies: { from: Slot; to: Slot; label: string }[] = [];
533 fn.outputs.forEach((out, i) => {
534 const to = spec.feedback[i];
535 if (!to) return;
536 const src = slots.get(out.cName);
537 const dst = host.get(to);
538 if (!src) {
539 throw new UnsupportedOnGpu(
540 `'${fn.name}' declares the output '${out.name}' but never assigns it`,
541 );
542 }
543 if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
544 if (src.count !== dst.count) {
545 throw new UnsupportedOnGpu(
546 `'${out.name}' (${src.count} elements) cannot feed '${to}' (${dst.count})`,
547 );
548 }
549 copies.push({ from: src, to: dst, label: `${out.name} -> ${to}` });
550 });
552 const written = new Set<GPUBuffer>();
553 for (const c of copies) written.add(c.to.buffer);
554 for (const c of copies) {
555 // Only a source another copy overwrites needs staging, and only if it
556 // is not that same copy's own destination (which is a no-op anyway).
557 if (c.from.buffer !== c.to.buffer && written.has(c.from.buffer)) {
558 const scratch = alloc(`mgpu-feedback-scratch`, c.from.count);
559 ops.push({
560 kind: 'copy',
561 from: c.from.buffer,
562 to: scratch.buffer,
563 bytes: 4 * c.from.count,
564 label: `${c.label} (staged)`,
565 });
566 c.from = scratch;
567 }
568 }
569 for (const c of copies) {
570 if (c.from.buffer === c.to.buffer) continue; // already in place
571 ops.push({
572 kind: 'copy',
573 from: c.from.buffer,
574 to: c.to.buffer,
575 bytes: 4 * c.from.count,
576 label: c.label,
577 });
578 }
579 }
580 }
582 /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
583 setParams(values: Record<string, number>): void {
584 this.paramNames.forEach((name, i) => {
585 const v = values[name];
586 this.#paramData[i] = Number.isFinite(v) ? v : 0;
587 });
588 this.#device.queue.writeBuffer(
589 this.#paramBuf,
590 0,
591 this.#paramData as Float32Array<ArrayBuffer>,
592 );
593 }
595 /** Buffer holding the named value, or undefined if the .m never binds it. */
596 buffer(name: string): GPUBuffer | undefined {
597 return this.#byName.get(name)?.buffer;
598 }
600 elementCount(name: string): number | undefined {
601 return this.#byName.get(name)?.count;
602 }
604 /**
605 * Record `steps` passes of this plan. Synchronous: no awaits, no readback.
606 * The dispatches share one compute pass, which WebGPU executes in submission
607 * order with a barrier between them.
608 *
609 * `after` runs once per step, inside the same submission — which is what
610 * lets the microphone sample every timestep rather than every frame.
611 */
612 encodeSteps(
613 encoder: GPUCommandEncoder,
614 steps: number,
615 after?: (encoder: GPUCommandEncoder) => void,
616 ): void {
617 for (let s = 0; s < steps; s++) {
618 this.#encodeOps(encoder);
619 after?.(encoder);
620 }
621 }
623 /** Record one pass over the op sequence into `encoder`. */
624 #encodeOps(encoder: GPUCommandEncoder): void {
625 let pass: GPUComputePassEncoder | null = null;
626 const inPass = (): GPUComputePassEncoder => {
627 if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
628 return pass;
629 };
630 const endPass = (): void => {
631 if (pass) {
632 pass.end();
633 pass = null;
634 }
635 };
636 for (const op of this.#ops) {
637 switch (op.kind) {
638 case 'kernel': {
639 const p = inPass();
640 p.setPipeline(op.pipeline);
641 p.setBindGroup(0, op.bindGroup);
642 p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
643 if (op.copyBack) {
644 endPass();
645 encoder.copyBufferToBuffer(
646 op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
647 );
648 }
649 break;
650 }
651 case 'external': {
652 const p = inPass();
653 p.setPipeline(op.pipeline);
654 p.setBindGroup(0, op.bindGroup);
655 p.dispatchWorkgroups(op.workgroups);
656 break;
657 }
658 case 'copy':
659 endPass();
660 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
661 break;
662 }
663 }
664 endPass();
665 }
667 /** Human-readable op sequence — what the .m actually compiled to. */
668 describe(): string[] {
669 return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
670 }
672 destroy(): void {
673 for (const b of this.#owned) b.destroy();
674 this.#paramBuf.destroy();
675 this.#owned.length = 0;
676 }
677}
679/** `lp = lapw(p, wall)` -> the op's name and its arguments. */
680function externalCall(
681 stmt: Assign,
682): { name: OpKind; args: (IRExpr & { kind: 'Var' })[] } | null {
683 const e = stmt.expr;
684 if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
685 const spec = EXTERNAL_OPS.get(e.name)!;
686 if (e.args.length !== spec.args.length) {
687 throw new UnsupportedOnGpu(
688 `'${e.name}' takes ${spec.args.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 as OpKind, 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}