2 * MATLAB source -> numbl's JIT IR, ready for the WGSL backend.
3 *
4 * A model file defines ordinary MATLAB functions; the host specializes the ones
5 * it needs (`init`, `step`) for the concrete argument types of the current grid.
6 * This is exactly how numbl drives its own JIT — the caller supplies argument
7 * types, and lowering fixes every type and shape from there.
8 *
9 * Driving it through function signatures rather than injected scope means the
10 * .m declares what it needs: each parameter name is matched against what the
11 * host offers, and a name the host does not provide is a compile error rather
12 * than a silently undefined variable.
13 *
14 * Two numbl passes matter here:
15 * - `specializeUserFunction` lowers one function to IR, one statement per
16 * operation (ANF), with every node's type fixed.
17 * - `inlinePass` then folds single-use temps back into their consumer, so a
18 * source line like `fu = a - u + u.*u.*v` becomes ONE statement whose RHS is
19 * an expression tree — i.e. one GPU kernel instead of four.
20 */
21import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
22import { Workspace, Lowerer, tensorDouble, scalarDouble } from 'numbl-src/numbl-core/jit/index.ts';
23import { specializeUserFunction } from 'numbl-src/numbl-core/jit/lowering/specialize.ts';
24import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
25import type { For, IRExpr, IRFunc, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
26import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
27import { externalOpFiles, type GridSizes } from './externals.ts';
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 28import { expandUserCalls } from './inlineCalls.ts';
29import { inFunction, ModelCompileError } from './errors.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 30
31/** What the host can supply for an argument the .m declares. */
32export type Binding =
33 /** An array, passed in a GPU buffer. */
34 | { kind: 'tensor'; shape: number[] }
35 /** A tunable scalar. Deliberately carries no exact value: an exact scalar
36 * would be constant-folded into the kernels, so moving a slider would force
37 * a recompile instead of just rewriting a uniform. */
38 | { kind: 'param' }
39 /** A fixed scalar, exact so array constructors reading it keep static
40 * shapes. */
41 | { kind: 'const'; value: number };
43const typeOf = (b: Binding): Type => {
44 switch (b.kind) {
45 case 'tensor':
46 return tensorDouble(b.shape);
47 case 'param':
48 return scalarDouble('unknown');
49 case 'const':
50 // Carry the sign too: numbl's sign lattice decides, for instance,
51 // whether sqrt() of a value can go complex.
52 return scalarDouble(
53 b.value > 0 ? 'positive' : b.value < 0 ? 'negative' : 'zero',
54 b.value,
55 );
56 }
57};
59/** One specialized function, as the planner consumes it. */
60export interface CompiledFunction {
61 name: string;
62 /** Declared arguments, in order, with the cName each lowered to. */
63 params: { name: string; cName: string; binding: Binding }[];
64 /** Requested outputs, in order, with the cName holding each result. */
65 outputs: { name: string; cName: string; ty: Type }[];
66 /** The lowered body. Read this only after `finish()`: the inline pass
67 * REPLACES the statement array rather than mutating it, so this is a live
68 * view of the function rather than a snapshot. */
69 readonly body: IRStmt[];
70}
72/** The shape of a `function` statement in numbl's AST. */
73interface FunctionDecl {
74 type: 'Function';
75 name: string;
76 params: string[];
77 outputs: string[];
78}
80/**
81 * A parsed model. Specialize the functions you need, then call `finish()` once
82 * — the inline pass rewrites every specialization together.
83 */
84export class CompiledModel {
85 #lowerer: Lowerer;
86 #decls: Map<string, FunctionDecl>;
87 #bindings: Record<string, Binding>;
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 88 /** Functions handed out by specialize(), in order — the ones whose bodies
89 * the planner will execute, and so the ones finish() expands. */
90 #entries: IRFunc[] = [];
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 91
92 constructor(
93 source: string,
94 bindings: Record<string, Binding>,
95 grid: GridSizes,
96 fileName = 'model.m',
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 97 /** Shared .m files compiled alongside the model — the operator and solver
98 * library. Only each file's namesake function is visible to the model,
99 * as in MATLAB; a model function of the same name shadows it. */
100 libs: { name: string; source: string }[] = [],
102 const ast = parseMFile(source, fileName);
103 const ws = new Workspace(fileName, []);
104 ws.addFile({ name: fileName, source, ast });
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 105 for (const lib of libs) {
106 ws.addFile({ name: lib.name, source: lib.source, ast: parseMFile(lib.source, lib.name) });
107 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 108 // synth / analys become resolvable, with their type rules.
109 for (const f of externalOpFiles(grid)) ws.addFile(f);
110 ws.finalize();
112 this.#bindings = bindings;
113 this.#lowerer = new Lowerer(ws);
114 this.#decls = new Map();
115 for (const stmt of ast.body as { type: string }[]) {
116 if (stmt.type === 'Function') {
117 const fn = stmt as unknown as FunctionDecl;
118 this.#decls.set(fn.name, fn);
119 }
120 }
121 }
123 /** Names of the functions the file defines. */
124 functionNames(): string[] {
125 return [...this.#decls.keys()];
126 }
128 /**
129 * Lower `name` for the current bindings, requesting `nargout` outputs.
130 * Every declared parameter must name something the host provides.
131 */
132 specialize(name: string, nargout: number): CompiledFunction {
133 const decl = this.#decls.get(name);
134 if (!decl) {
135 const defined = this.functionNames();
136 throw new ModelCompileError(
137 `the model must define a function named '${name}'` +
138 (defined.length
139 ? ` (it defines ${defined.map((n) => `'${n}'`).join(', ')})`
140 : ' (it defines no functions)'),
141 );
142 }
143 if (decl.outputs.length < nargout) {
144 throw new ModelCompileError(
145 `'${name}' must return ${nargout} value${nargout === 1 ? '' : 's'}, ` +
146 `but declares ${decl.outputs.length}`,
147 );
148 }
150 const bindings = decl.params.map((p) => {
151 const b = this.#bindings[p];
152 if (!b) {
153 const offered = Object.keys(this.#bindings).join(', ');
154 throw new ModelCompileError(
155 `'${name}' takes an argument named '${p}', which this app does not ` +
156 `provide. Available: ${offered}.`,
157 );
158 }
159 return b;
160 });
162 const fn: IRFunc = specializeUserFunction.call(
163 this.#lowerer,
164 decl,
165 bindings.map(typeOf),
166 undefined,
167 undefined,
168 undefined,
169 nargout,
170 undefined,
171 );
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 172 this.#entries.push(fn);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 173
174 return {
175 name,
176 params: fn.params.map((p, i) => ({
177 name: p,
178 cName: fn.cParams[i],
179 binding: bindings[i],
180 })),
181 outputs: fn.outputs.slice(0, nargout).map((o, i) => ({
182 name: o,
183 cName: fn.cOutputs[i],
184 ty: fn.outputTypes[i],
185 })),
186 // A getter, not a snapshot: `finish()` runs after every specialization
187 // and swaps in a rewritten statement array.
188 get body() {
189 return fn.body;
190 },
191 };
192 }
194 /**
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 195 * Expand user-function calls, then run the inline pass. Both rewrite the
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 196 * function bodies in place, so `CompiledFunction`s handed out earlier are
197 * updated too.
199 * Expansion comes first: with every call spliced into its caller, each
200 * entry function is one flat body, and the inline pass fuses it exactly as
201 * it would the same code written out by hand — a call boundary neither
202 * blocks fusion nor changes what compiles. Only the entry functions are
203 * rewritten; the callee specializations they were cloned from are no longer
204 * referenced.
206 finish(): void {
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 207 for (const fn of this.#entries) {
208 inFunction(fn.name, () =>
209 expandUserCalls(fn, (cName) => this.#lowerer.specializations.get(cName)),
210 );
211 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 213 // Snapshot what each loop body assigns, before the pass can rewrite it.
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 214 const entries = new Map(this.#entries.map((fn) => [fn.cName, fn]));
215 const loops = [...entries.values()].flatMap((fn) =>
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 216 forLoops(fn.body).map((loop) => ({
217 fn,
218 loop,
219 assignedBefore: assignedCNames(loop.body),
220 })),
221 );
59f3e22Factor the solver out of the models; add BiCGSTAB and GMRESJeremy Magland 223 inlinePass({ topLevelStmts: [], functions: entries });
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 224
225 for (const { fn, loop, assignedBefore } of loops) checkLoopEscapes(fn, loop, assignedBefore);
226 }
227}
229/** Every `for` in a statement list, including nested ones. */
230function forLoops(stmts: IRStmt[]): For[] {
231 const out: For[] = [];
232 const walk = (list: IRStmt[]): void => {
233 for (const s of list) {
234 if (s.kind === 'For') {
235 out.push(s);
236 walk(s.body);
237 }
238 }
239 };
240 walk(stmts);
241 return out;
242}
244/** cNames assigned anywhere in a statement list, including inside loops. */
245function assignedCNames(stmts: IRStmt[]): Set<string> {
246 const out = new Set<string>();
247 const walk = (list: IRStmt[]): void => {
248 for (const s of list) {
249 if (s.kind === 'Assign') out.add(s.cName);
250 else if (s.kind === 'For') walk(s.body);
251 }
252 };
253 walk(stmts);
254 return out;
255}
257/** Call `visit` for every variable read in an expression. */
258function forEachVarRead(e: IRExpr, visit: (cName: string) => void): void {
259 const walk = (x: IRExpr): void => {
260 switch (x.kind) {
261 case 'Var':
262 return visit(x.cName);
263 case 'Binary':
264 walk(x.left);
265 walk(x.right);
266 return;
267 case 'Unary':
268 walk(x.operand);
269 return;
270 case 'Call':
271 x.args.forEach(walk);
272 return;
273 default:
274 return;
275 }
276 };
277 walk(e);
278}
280/**
281 * Refuse a loop whose result the inline pass folded away.
282 *
283 * numbl's inline pass substitutes a single-use producer into its consumer and
284 * drops the producer. Inside a loop body it runs with no protected names — it
285 * counts uses within that body alone — so an assignment whose only *visible*
286 * use is later in the same body can be elided even though something outside
287 * the loop still wants the value.
288 *
289 * That elision is correct for a body-local temp, which is what makes fusion
290 * work inside the loop, and it is caught downstream in the two cases where the
291 * value has no buffer at all: the planner already refuses a declared output
292 * that is never assigned, and a read of a name it never allocated. The case it
293 * would not catch is a variable assigned *before* the loop as well — there the
294 * buffer exists, holding the pre-loop value, and the loop would silently
295 * contribute nothing. So check all three here, in one place, against what the
296 * body assigned before the pass ran.
297 */
298function checkLoopEscapes(fn: IRFunc, loop: For, assignedBefore: Set<string>): void {
299 const assignedAfter = assignedCNames(loop.body);
300 const elided = [...assignedBefore].filter((c) => !assignedAfter.has(c));
301 if (elided.length === 0) return;
303 // Reads anywhere in the function outside this loop's own body.
304 const readOutside = new Set<string>();
305 const walk = (list: IRStmt[]): void => {
306 for (const s of list) {
307 if (s === (loop as IRStmt)) continue; // the loop's own body is not "outside"
308 if (s.kind === 'Assign') forEachVarRead(s.expr, (c) => readOutside.add(c));
309 else if (s.kind === 'For') walk(s.body);
310 }
311 };
312 walk(fn.body);
314 const outputs = new Set(fn.cOutputs);
315 const escaping = elided.filter((c) => readOutside.has(c) || outputs.has(c));
316 if (escaping.length === 0) return;
318 const names = [...new Set(escaping)].map((c) => `'${c}'`).join(', ');
319 throw new ModelCompileError(
320 `inside the 'for' loop, ${names} is assigned but only read later in the ` +
321 `same iteration, so the compiler folded the assignment into its reader — ` +
322 `yet the value is also wanted outside the loop. Read it once outside the ` +
323 `loop instead, or use it more than once inside it.`,
324 );
325}