1/**
2 * Fold the single-use ANF temps numbl's inline pass left behind.
3 *
4 * numbl's own pass (`inlinePass`) only folds producers its C backend can fuse,
5 * which leaves out tensor-producing calls — `sin(x)`, `exp(x)`, `tanh(x)`. The
6 * WGSL emitter fuses all of those happily, so without this pass a source line
7 * like
8 *
9 * s = (amp * env) .* sin(2*pi*f*(t - t0)) .* exp(-(gx + gy));
10 *
11 * plans as half a dozen kernels rather than one, each writing a whole grid to
12 * memory for the next one to read straight back.
13 *
14 * Same shape as numbl's pass, deliberately narrower where it matters: only
15 * compiler temps (`_mtoc2_*`) are folded, so every variable the .m names keeps
16 * its own buffer and one source line stays one kernel. The producer's RHS must
17 * be something the emitter can evaluate per element, its result must be used
18 * exactly once, and nothing between the two statements may write to anything
19 * it reads.
20 *
21 * Adapted from math-webgpu-sandbox's src/mgpu/fuse.ts, trimmed to the
22 * straight-line bodies this project compiles.
23 */
24import type { IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
25import { isGpuFusableExpr } from './wgsl.ts';
27const isTemp = (cName: string): boolean => cName.startsWith('_mtoc2_');
29export function fuseTemps(stmts: IRStmt[]): IRStmt[] {
30 let cur = stmts;
31 for (let iter = 0; iter < 32; iter++) {
32 const next = fuseOnePass(cur);
33 if (next === cur) break;
34 cur = next;
35 }
36 return cur;
37}
39/** One sweep, at most one fold. Returns the same array when nothing fired. */
40function fuseOnePass(stmts: IRStmt[]): IRStmt[] {
41 const uses = useCounts(stmts);
42 for (let i = 0; i < stmts.length; i++) {
43 const p = stmts[i];
44 if (p.kind !== 'Assign' || !isTemp(p.cName)) continue;
45 if (uses.get(p.cName) !== 1) continue;
46 if (!isGpuFusableExpr(p.expr)) continue;
48 const reads = new Set<string>();
49 walkVars(p.expr, (c) => reads.add(c));
51 for (let j = i + 1; j < stmts.length; j++) {
52 const c = stmts[j];
53 if (c.kind !== 'Assign') break; // anything else ends the safe window
54 if (countIn(c.expr, p.cName) > 0) {
55 // The one use. Fold into it if it is a kernel; if it is a stencil
56 // call, or anything else that wants its argument in a buffer of its
57 // own, leave the producer alone.
58 if (countIn(c.expr, p.cName) === 1 && isGpuFusableExpr(c.expr)) {
59 c.expr = substitute(c.expr, p.cName, p.expr);
60 const out = stmts.slice();
61 out.splice(i, 1);
62 return out;
63 }
64 break;
65 }
66 // An intervening write to the temp itself, or to one of its operands,
67 // invalidates the fold window. Statements that do neither are simply
68 // skipped over — the stencil dispatch between the source term and the
69 // update it feeds is exactly that case.
70 if (c.cName === p.cName || reads.has(c.cName)) break;
71 }
72 }
73 return stmts;
74}
76function useCounts(stmts: IRStmt[]): Map<string, number> {
77 const counts = new Map<string, number>();
78 const bump = (c: string): void => {
79 counts.set(c, (counts.get(c) ?? 0) + 1);
80 };
81 for (const s of stmts) if (s.kind === 'Assign') walkVars(s.expr, bump);
82 return counts;
83}
85function walkVars(e: IRExpr, visit: (cName: string) => void): void {
86 const walk = (x: IRExpr): void => {
87 switch (x.kind) {
88 case 'Var':
89 visit(x.cName);
90 return;
91 case 'Binary':
92 walk(x.left);
93 walk(x.right);
94 return;
95 case 'Unary':
96 walk(x.operand);
97 return;
98 case 'Call':
99 x.args.forEach(walk);
100 return;
101 default:
102 return;
103 }
104 };
105 walk(e);
106}
108function countIn(e: IRExpr, cName: string): number {
109 let n = 0;
110 walkVars(e, (c) => {
111 if (c === cName) n++;
112 });
113 return n;
114}
116/** Replace the (single) `Var` read of `cName` with `replacement`. */
117function substitute(e: IRExpr, cName: string, replacement: IRExpr): IRExpr {
118 const sub = (x: IRExpr): IRExpr => {
119 if (x.kind === 'Var' && x.cName === cName) return replacement;
120 switch (x.kind) {
121 case 'Binary':
122 x.left = sub(x.left);
123 x.right = sub(x.right);
124 return x;
125 case 'Unary':
126 x.operand = sub(x.operand);
127 return x;
128 case 'Call':
129 for (let i = 0; i < x.args.length; i++) x.args[i] = sub(x.args[i]);
130 return x;
131 default:
132 return x;
133 }
134 };
135 return sub(e);
136}