2 * Expand calls to user-defined MATLAB functions into the caller's body.
3 *
4 * numbl lowers `dL = dlap(X, ...)` to a single `Call` statement whose cName
5 * names the callee's specialization — one IRFunc per distinct argument-type
6 * signature, shared across call sites. The WGSL planner, though, executes a
7 * flat statement list: every value is a buffer, every statement a dispatch,
8 * and a call boundary has no runtime meaning on the GPU. So this pass gives a
9 * call the only meaning it can have there: the callee's lowered body, spliced
10 * in at the call site.
11 *
12 * Each call site gets its own clone of the callee body, with every
13 * callee-local cName made unique to the site — so two calls to the same
14 * function get separate buffers, while the *same* site re-planned per
15 * unrolled loop iteration keeps reusing its buffers (the planner keys buffers
16 * by cName, and the clone is made once, here, not per iteration).
17 *
18 * Arguments bind by renaming, not copying, wherever that is sound: a callee
19 * parameter whose argument is a plain variable — after the inline pass they
20 * all are, since it never folds an expression into a call argument — reads
21 * the caller's buffer directly, unless the callee reassigns the parameter, in
22 * which case a copy is materialized so the caller's value is not clobbered.
23 * Outputs are the same in reverse: assignments to a callee output become
24 * assignments to the caller's target variable, which is what makes a
25 * loop-carried output (a solver iterating its result) work unchanged.
26 *
27 * The pass runs before numbl's inline (fusion) pass, which works per function
28 * and treats a user call as an opaque producer. With every call already
29 * spliced away, fusion sees one flat body and folds it exactly as it would
30 * the same code written out by hand — a callee's final assignment can fuse
31 * into its consumer, and the compiled plan is identical either way.
32 */
33import type {
34 Assign,
35 For,
36 IRExpr,
37 IRFunc,
38 IRStmt,
39 Span,
40} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
41import { ModelCompileError } from './errors.ts';
43/** A located compile failure. A span from a shared lib file still carries its
44 * offsets; they are only ever mapped onto the model source for display, so a
45 * failure inside a lib mislocates but still names the construct. */
46const fail = (message: string, span: Span): ModelCompileError =>
47 new ModelCompileError(message, { start: span.start, end: span.end });
49/** Looks up a callee's specialization by its mangled cName. */
50export type ResolveFn = (cName: string) => IRFunc | undefined;
52/**
53 * Expand every user-function call in `fn`, recursively, mutating `fn.body`.
54 * Returns true if anything was expanded.
55 */
56export function expandUserCalls(fn: IRFunc, resolve: ResolveFn): boolean {
57 const ctx = { resolve, site: 0, expanded: false };
58 fn.body = expandBody(fn.body, ctx, [fn.cName]);
59 return ctx.expanded;
60}
62interface Ctx {
63 resolve: ResolveFn;
64 /** Call-site counter, for unique cNames. */
65 site: number;
66 expanded: boolean;
67}
69function expandBody(stmts: IRStmt[], ctx: Ctx, stack: string[]): IRStmt[] {
70 const out: IRStmt[] = [];
71 for (const stmt of stmts) {
72 if (stmt.kind === 'Assign' && stmt.expr.kind === 'Call') {
73 const callee = ctx.resolve(stmt.expr.cName);
74 if (callee) {
75 out.push(
76 ...expandCall(
77 callee,
78 stmt.expr.args,
79 [{ name: stmt.name, cName: stmt.cName }],
80 ctx,
81 stack,
82 stmt,
83 ),
84 );
85 continue;
86 }
87 }
88 if (stmt.kind === 'MultiAssignCall' && !stmt.isBuiltin) {
89 const callee = ctx.resolve(stmt.cName);
90 if (!callee) {
91 throw fail(
92 `call to '${stmt.name}' does not resolve to a function in this model`,
93 stmt.span,
94 );
95 }
96 out.push(
97 ...expandCall(
98 callee,
99 stmt.args,
100 stmt.outputs.map((o) => o.binding),
101 ctx,
102 stack,
103 stmt,
104 ),
105 );
106 continue;
107 }
108 if (stmt.kind === 'For') {
109 stmt.body = expandBody(stmt.body, ctx, stack);
110 out.push(stmt);
111 continue;
112 }
113 out.push(stmt);
114 }
115 return out;
116}
118/**
119 * One call site: the callee's body, cloned and renamed into the caller's
120 * namespace, preceded by whatever argument bindings need materializing.
121 */
122function expandCall(
123 callee: IRFunc,
124 args: IRExpr[],
125 outs: ({ name: string; cName: string } | null)[],
126 ctx: Ctx,
127 stack: string[],
128 at: IRStmt,
129): IRStmt[] {
130 if (stack.includes(callee.cName)) {
131 throw fail(
132 `'${callee.name}' calls itself (perhaps through another function); ` +
133 `recursion cannot be compiled to a fixed sequence of GPU operations`,
134 at.span,
135 );
136 }
137 ctx.expanded = true;
138 const site = ++ctx.site;
139 /** Site-unique cName for a callee-local. */
140 const local = (cName: string): string => `${callee.name}$${site}$${cName}`;
141 /** Display name for a callee-local — shows up in buffer labels and in
142 * describe(), so a solver's internals read as `richardson#1.X`. */
143 const display = (name: string): string => `${callee.name}#${site}.${name}`;
145 const assigned = assignedCNames(callee.body);
146 const rename = new Map<string, string>();
147 const names = new Map<string, string>();
148 const prelude: IRStmt[] = [];
150 // Outputs first: assignments to a callee output become assignments to the
151 // caller's target. An ignored output (`~`, or trailing outputs the caller
152 // did not ask for) stays a site-local.
153 callee.cOutputs.forEach((c, j) => {
154 const target = j < outs.length ? outs[j] : null;
155 if (target) {
156 rename.set(c, target.cName);
157 names.set(c, target.name);
158 } else {
159 rename.set(c, local(c));
160 names.set(c, display(callee.outputs[j]));
161 }
162 });
164 // Parameters: rename onto the argument where sound, else materialize.
165 callee.cParams.forEach((p, i) => {
166 const arg = args[i];
167 if (arg === undefined) {
168 throw fail(
169 `'${callee.name}' takes ${callee.cParams.length} arguments, ` +
170 `but this call passes ${args.length}`,
171 at.span,
172 );
173 }
174 if (rename.has(p)) {
175 // The parameter is also an output (`function X = f(X)`): seed the
176 // caller's target with the argument, and let the body update it there.
177 prelude.push(makeAssign(names.get(p)!, rename.get(p)!, arg, callee.paramTypes[i], at));
178 } else if (arg.kind === 'Var' && !assigned.has(p)) {
179 rename.set(p, arg.cName);
180 names.set(p, arg.name);
181 } else {
182 // A non-variable argument, or a parameter the callee reassigns: bind it
183 // to a site-local first. For a scalar this is a free derived-scalar
184 // binding; for a tensor it is one copy kernel.
185 const c = local(p);
186 rename.set(p, c);
187 names.set(p, display(callee.params[i]));
188 prelude.push(makeAssign(names.get(p)!, c, arg, callee.paramTypes[i], at));
189 }
190 });
192 const body = cloneBody(callee.body, { rename, names, local, display, callee, at });
193 return [...prelude, ...expandBody(body, ctx, [...stack, callee.cName])];
194}
196const makeAssign = (
197 name: string,
198 cName: string,
199 expr: IRExpr,
200 ty: IRFunc['paramTypes'][number],
201 at: IRStmt,
202): Assign => ({ kind: 'Assign', name, cName, ty, expr, span: at.span });
204interface CloneCtx {
205 /** Callee cName -> caller-namespace cName. Filled for outputs and params up
206 * front; locals are added on first sight. */
207 rename: Map<string, string>;
208 names: Map<string, string>;
209 local: (cName: string) => string;
210 display: (name: string) => string;
211 callee: IRFunc;
212 at: IRStmt;
213}
215function cloneBody(stmts: IRStmt[], c: CloneCtx): IRStmt[] {
216 const out: IRStmt[] = [];
217 for (const s of stmts) {
218 switch (s.kind) {
219 case 'Assign':
220 out.push({
221 ...s,
222 name: mapName(s.name, s.cName, c),
223 cName: mapCName(s.cName, s.name, c),
224 expr: cloneExpr(s.expr, c),
225 });
226 break;
227 case 'MultiAssignCall':
228 out.push({
229 ...s,
230 args: s.args.map((a) => cloneExpr(a, c)),
231 outputs: s.outputs.map((o) =>
232 o.binding
233 ? {
234 ...o,
235 binding: {
236 name: mapName(o.binding.name, o.binding.cName, c),
237 cName: mapCName(o.binding.cName, o.binding.name, c),
238 },
239 }
240 : o,
241 ),
242 });
243 break;
244 case 'For': {
245 const loop: For = {
246 ...s,
247 cVar: mapCName(s.cVar, s.varName, c),
248 start: cloneExpr(s.start, c),
249 end: cloneExpr(s.end, c),
250 body: [],
251 };
252 loop.body = cloneBody(s.body, c);
253 out.push(loop);
254 break;
255 }
256 case 'ReturnFromFunction':
257 // The callee's return has no meaning at the splice point; statements
258 // never follow it (numbl refuses an early return at lowering).
259 break;
260 default:
261 throw fail(
262 `'${c.callee.name}' contains a '${s.kind}' statement, which cannot ` +
263 `be compiled to a fixed sequence of GPU operations`,
264 s.span,
265 );
266 }
267 }
268 return out;
269}
271/** Caller-namespace cName for a callee-side cName, minting one for a local
272 * seen for the first time. */
273function mapCName(cName: string, name: string, c: CloneCtx): string {
274 const existing = c.rename.get(cName);
275 if (existing) return existing;
276 const fresh = c.local(cName);
277 c.rename.set(cName, fresh);
278 c.names.set(cName, c.display(name));
279 return fresh;
280}
282function mapName(name: string, cName: string, c: CloneCtx): string {
283 return c.names.get(cName) ?? c.display(name);
284}
286function cloneExpr(e: IRExpr, c: CloneCtx): IRExpr {
287 switch (e.kind) {
288 case 'Var':
289 return {
290 ...e,
291 name: mapName(e.name, e.cName, c),
292 cName: mapCName(e.cName, e.name, c),
293 };
294 case 'Binary':
295 return { ...e, left: cloneExpr(e.left, c), right: cloneExpr(e.right, c) };
296 case 'Unary':
297 return { ...e, operand: cloneExpr(e.operand, c) };
298 case 'Call':
299 return { ...e, args: e.args.map((a) => cloneExpr(a, c)) };
300 default:
301 // Literals and anything else without variable reads: share as-is (the
302 // planner treats expressions as read-only).
303 return e;
304 }
305}
307/** cNames assigned anywhere in a statement list, including loop variables. */
308function assignedCNames(stmts: IRStmt[]): Set<string> {
309 const out = new Set<string>();
310 const walk = (list: IRStmt[]): void => {
311 for (const s of list) {
312 if (s.kind === 'Assign') out.add(s.cName);
313 else if (s.kind === 'MultiAssignCall') {
314 for (const o of s.outputs) if (o.binding) out.add(o.binding.cName);
315 } else if (s.kind === 'For') {
316 out.add(s.cVar);
317 walk(s.body);
318 }
319 }
320 };
321 walk(stmts);
322 return out;
323}