1/**
2 * Sandbox-side fusion: fold the single-use ANF temps numbl's inline pass
3 * left behind.
4 *
5 * numbl's pass (`inlinePass`) only folds producers its C backend can fuse —
6 * no tensor-producing Calls (`sin(x)`), no logical results, no ranges. The
7 * WGSL emitter fuses all of those, so without this pass a line like
8 * `y = y + 0.1*sin(x + k) .* exp(-x)` becomes five kernels instead of one.
9 *
10 * Same shape as numbl's pass, deliberately narrower where it matters:
11 * - only compiler temps (`_mtoc2_*`) are folded, so every user-named
12 * variable still materializes — one source line stays one kernel, and
13 * anything the user might echo or reuse keeps its buffer;
14 * - the producer's RHS must be fusable on the GPU (isGpuFusableExpr);
15 * - the single use must be in a following Assign/ExprStmt at the same body
16 * level, with no intervening write to the producer's operands and no
17 * control-flow statement in between.
18 *
19 * `for` bodies are processed as their own levels; nothing folds across the
20 * loop boundary.
21 */
22import type { Assign, ExprStmt, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
23import { isMultiElement, type Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
24import { isGpuFusableExpr } from './wgsl.ts';
26const isMulti = (t: Type): boolean => t.kind === 'Numeric' && isMultiElement(t);
28export function fuseTemps(stmts: IRStmt[]): void {
29 for (const s of stmts) {
30 if (s.kind === 'For') fuseTemps(s.body);
31 }
32 for (let iter = 0; iter < 32; iter++) {
33 if (!fuseOnePass(stmts)) break;
34 }
35}
37const isTemp = (cName: string): boolean => cName.startsWith('_mtoc2_');
39function fuseOnePass(stmts: IRStmt[]): boolean {
40 const uses = useCounts(stmts);
41 for (let i = 0; i < stmts.length; i++) {
42 const p = stmts[i];
43 if (p.kind !== 'Assign') continue;
44 if (!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' && c.kind !== 'ExprStmt') break; // control flow
54 const holder = c as Assign | ExprStmt;
55 if (c.kind === 'ExprStmt' && isMulti(p.ty) && countIn(holder.expr, p.cName) === 1) {
56 break; // disp/fprintf read tensors by name, not as expressions
57 }
58 if (countIn(holder.expr, p.cName) === 1) {
59 holder.expr = substitute(holder.expr, p.cName, p.expr);
60 stmts.splice(i, 1);
61 return true;
62 }
63 if (c.kind === 'Assign' && (c.cName === p.cName || reads.has(c.cName))) {
64 break; // intervening write invalidates the fold window
65 }
66 }
67 }
68 return false;
69}
71/** Var-read counts by cName, across this level AND nested bodies (a use
72 * inside a nested loop must keep the temp alive at this level). */
73function useCounts(stmts: IRStmt[]): Map<string, number> {
74 const counts = new Map<string, number>();
75 const bump = (c: string): void => {
76 counts.set(c, (counts.get(c) ?? 0) + 1);
77 };
78 const walkStmts = (list: IRStmt[]): void => {
79 for (const s of list) {
80 if (s.kind === 'Assign' || s.kind === 'ExprStmt') walkVars(s.expr, bump);
81 else if (s.kind === 'For') {
82 walkVars(s.start, bump);
83 walkVars(s.end, bump);
84 walkStmts(s.body);
85 }
86 }
87 };
88 walkStmts(stmts);
89 return counts;
90}
92function walkVars(e: IRExpr, visit: (cName: string) => void): void {
93 const walk = (x: IRExpr): void => {
94 switch (x.kind) {
95 case 'Var':
96 visit(x.cName);
97 return;
98 case 'Binary':
99 walk(x.left);
100 walk(x.right);
101 return;
102 case 'Unary':
103 walk(x.operand);
104 return;
105 case 'Call':
106 x.args.forEach(walk);
107 return;
108 case 'IndexSlice':
109 walk(x.base);
110 x.index.forEach((a) => {
111 const inner = (a as { expr?: IRExpr }).expr;
112 if (inner) walk(inner);
113 });
114 return;
115 case 'MakeRange':
116 walk(x.start);
117 walk(x.step);
118 walk(x.end);
119 return;
120 default:
121 return;
122 }
123 };
124 walk(e);
125}
127function countIn(e: IRExpr, cName: string): number {
128 let n = 0;
129 walkVars(e, (c) => {
130 if (c === cName) n++;
131 });
132 return n;
133}
135/** Replace the (single) `Var` read of `cName` with `replacement`. */
136function substitute(e: IRExpr, cName: string, replacement: IRExpr): IRExpr {
137 const sub = (x: IRExpr): IRExpr => {
138 if (x.kind === 'Var' && x.cName === cName) return replacement;
139 switch (x.kind) {
140 case 'Binary':
141 x.left = sub(x.left);
142 x.right = sub(x.right);
143 return x;
144 case 'Unary':
145 x.operand = sub(x.operand);
146 return x;
147 case 'Call':
148 for (let i = 0; i < x.args.length; i++) x.args[i] = sub(x.args[i]);
149 return x;
150 case 'IndexSlice':
151 x.base = sub(x.base);
152 return x;
153 case 'MakeRange':
154 x.start = sub(x.start);
155 x.step = sub(x.step);
156 x.end = sub(x.end);
157 return x;
158 default:
159 return x;
160 }
161 };
162 return sub(e);
163}