4f822e1turing-surface-cache: reaction-diffusion solutions at a chosen end time, shared through a cloud cacheJeremy Magland 1/**
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';
28import { ModelCompileError } from './errors.ts';
30/** What the host can supply for an argument the .m declares. */
31export type Binding =
32 /** An array, passed in a GPU buffer. */
33 | { kind: 'tensor'; shape: number[] }
34 /** A tunable scalar. Deliberately carries no exact value: an exact scalar
35 * would be constant-folded into the kernels, so moving a slider would force
36 * a recompile instead of just rewriting a uniform. */
37 | { kind: 'param' }
38 /** A fixed scalar, exact so array constructors reading it keep static
39 * shapes. */
40 | { kind: 'const'; value: number };
42const typeOf = (b: Binding): Type => {
43 switch (b.kind) {
44 case 'tensor':
45 return tensorDouble(b.shape);
46 case 'param':
47 return scalarDouble('unknown');
48 case 'const':
49 // Carry the sign too: numbl's sign lattice decides, for instance,
50 // whether sqrt() of a value can go complex.
51 return scalarDouble(
52 b.value > 0 ? 'positive' : b.value < 0 ? 'negative' : 'zero',
53 b.value,
54 );
55 }
56};
58/** One specialized function, as the planner consumes it. */
59export interface CompiledFunction {
60 name: string;
61 /** Declared arguments, in order, with the cName each lowered to. */
62 params: { name: string; cName: string; binding: Binding }[];
63 /** Requested outputs, in order, with the cName holding each result. */
64 outputs: { name: string; cName: string; ty: Type }[];
65 /** The lowered body. Read this only after `finish()`: the inline pass
66 * REPLACES the statement array rather than mutating it, so this is a live
67 * view of the function rather than a snapshot. */
68 readonly body: IRStmt[];
69}
71/** The shape of a `function` statement in numbl's AST. */
72interface FunctionDecl {
73 type: 'Function';
74 name: string;
75 params: string[];
76 outputs: string[];
77}
79/**
80 * A parsed model. Specialize the functions you need, then call `finish()` once
81 * — the inline pass rewrites every specialization together.
82 */
83export class CompiledModel {
84 #lowerer: Lowerer;
85 #decls: Map<string, FunctionDecl>;
86 #bindings: Record<string, Binding>;
88 constructor(
89 source: string,
90 bindings: Record<string, Binding>,
91 grid: GridSizes,
92 fileName = 'model.m',
93 ) {
94 const ast = parseMFile(source, fileName);
95 const ws = new Workspace(fileName, []);
96 ws.addFile({ name: fileName, source, ast });
97 // synth / analys become resolvable, with their type rules.
98 for (const f of externalOpFiles(grid)) ws.addFile(f);
99 ws.finalize();
101 this.#bindings = bindings;
102 this.#lowerer = new Lowerer(ws);
103 this.#decls = new Map();
104 for (const stmt of ast.body as { type: string }[]) {
105 if (stmt.type === 'Function') {
106 const fn = stmt as unknown as FunctionDecl;
107 this.#decls.set(fn.name, fn);
108 }
109 }
110 }
112 /** Names of the functions the file defines. */
113 functionNames(): string[] {
114 return [...this.#decls.keys()];
115 }
117 /**
118 * Lower `name` for the current bindings, requesting `nargout` outputs.
119 * Every declared parameter must name something the host provides.
120 */
121 specialize(name: string, nargout: number): CompiledFunction {
122 const decl = this.#decls.get(name);
123 if (!decl) {
124 const defined = this.functionNames();
125 throw new ModelCompileError(
126 `the model must define a function named '${name}'` +
127 (defined.length
128 ? ` (it defines ${defined.map((n) => `'${n}'`).join(', ')})`
129 : ' (it defines no functions)'),
130 );
131 }
132 if (decl.outputs.length < nargout) {
133 throw new ModelCompileError(
134 `'${name}' must return ${nargout} value${nargout === 1 ? '' : 's'}, ` +
135 `but declares ${decl.outputs.length}`,
136 );
137 }
139 const bindings = decl.params.map((p) => {
140 const b = this.#bindings[p];
141 if (!b) {
142 const offered = Object.keys(this.#bindings).join(', ');
143 throw new ModelCompileError(
144 `'${name}' takes an argument named '${p}', which this app does not ` +
145 `provide. Available: ${offered}.`,
146 );
147 }
148 return b;
149 });
151 const fn: IRFunc = specializeUserFunction.call(
152 this.#lowerer,
153 decl,
154 bindings.map(typeOf),
155 undefined,
156 undefined,
157 undefined,
158 nargout,
159 undefined,
160 );
162 return {
163 name,
164 params: fn.params.map((p, i) => ({
165 name: p,
166 cName: fn.cParams[i],
167 binding: bindings[i],
168 })),
169 outputs: fn.outputs.slice(0, nargout).map((o, i) => ({
170 name: o,
171 cName: fn.cOutputs[i],
172 ty: fn.outputTypes[i],
173 })),
174 // A getter, not a snapshot: `finish()` runs after every specialization
175 // and swaps in a rewritten statement array.
176 get body() {
177 return fn.body;
178 },
179 };
180 }
182 /**
183 * Run the inline pass over everything specialized so far. It rewrites the
184 * function bodies in place, so `CompiledFunction`s handed out earlier are
185 * updated too.
186 */
187 finish(): void {
188 // Snapshot what each loop body assigns, before the pass can rewrite it.
189 const loops = [...this.#lowerer.specializations.values()].flatMap((fn) =>
190 forLoops(fn.body).map((loop) => ({
191 fn,
192 loop,
193 assignedBefore: assignedCNames(loop.body),
194 })),
195 );
197 inlinePass({ topLevelStmts: [], functions: this.#lowerer.specializations });
199 for (const { fn, loop, assignedBefore } of loops) checkLoopEscapes(fn, loop, assignedBefore);
200 }
201}
203/** Every `for` in a statement list, including nested ones. */
204function forLoops(stmts: IRStmt[]): For[] {
205 const out: For[] = [];
206 const walk = (list: IRStmt[]): void => {
207 for (const s of list) {
208 if (s.kind === 'For') {
209 out.push(s);
210 walk(s.body);
211 }
212 }
213 };
214 walk(stmts);
215 return out;
216}
218/** cNames assigned anywhere in a statement list, including inside loops.
219 * A MultiAssignCall assigns every bound output slot. */
220function assignedCNames(stmts: IRStmt[]): Set<string> {
221 const out = new Set<string>();
222 const walk = (list: IRStmt[]): void => {
223 for (const s of list) {
224 if (s.kind === 'Assign') out.add(s.cName);
225 else if (s.kind === 'MultiAssignCall') {
226 for (const o of s.outputs) if (o.binding) out.add(o.binding.cName);
227 } else if (s.kind === 'For') walk(s.body);
228 }
229 };
230 walk(stmts);
231 return out;
232}
234/** Call `visit` for every variable read in an expression. */
235function forEachVarRead(e: IRExpr, visit: (cName: string) => void): void {
236 const walk = (x: IRExpr): void => {
237 switch (x.kind) {
238 case 'Var':
239 return visit(x.cName);
240 case 'Binary':
241 walk(x.left);
242 walk(x.right);
243 return;
244 case 'Unary':
245 walk(x.operand);
246 return;
247 case 'Call':
248 x.args.forEach(walk);
249 return;
250 default:
251 return;
252 }
253 };
254 walk(e);
255}
257/**
258 * Refuse a loop whose result the inline pass folded away.
259 *
260 * numbl's inline pass substitutes a single-use producer into its consumer and
261 * drops the producer. Inside a loop body it runs with no protected names — it
262 * counts uses within that body alone — so an assignment whose only *visible*
263 * use is later in the same body can be elided even though something outside
264 * the loop still wants the value.
265 *
266 * That elision is correct for a body-local temp, which is what makes fusion
267 * work inside the loop, and it is caught downstream in the two cases where the
268 * value has no buffer at all: the planner already refuses a declared output
269 * that is never assigned, and a read of a name it never allocated. The case it
270 * would not catch is a variable assigned *before* the loop as well — there the
271 * buffer exists, holding the pre-loop value, and the loop would silently
272 * contribute nothing. So check all three here, in one place, against what the
273 * body assigned before the pass ran.
274 */
275function checkLoopEscapes(fn: IRFunc, loop: For, assignedBefore: Set<string>): void {
276 const assignedAfter = assignedCNames(loop.body);
277 const elided = [...assignedBefore].filter((c) => !assignedAfter.has(c));
278 if (elided.length === 0) return;
280 // Reads anywhere in the function outside this loop's own body.
281 const readOutside = new Set<string>();
282 const walk = (list: IRStmt[]): void => {
283 for (const s of list) {
284 if (s === (loop as IRStmt)) continue; // the loop's own body is not "outside"
285 if (s.kind === 'Assign') forEachVarRead(s.expr, (c) => readOutside.add(c));
286 else if (s.kind === 'MultiAssignCall') {
287 for (const a of s.args) forEachVarRead(a, (c) => readOutside.add(c));
288 } else if (s.kind === 'For') walk(s.body);
289 }
290 };
291 walk(fn.body);
293 const outputs = new Set(fn.cOutputs);
294 const escaping = elided.filter((c) => readOutside.has(c) || outputs.has(c));
295 if (escaping.length === 0) return;
297 const names = [...new Set(escaping)].map((c) => `'${c}'`).join(', ');
298 throw new ModelCompileError(
299 `inside the 'for' loop, ${names} is assigned but only read later in the ` +
300 `same iteration, so the compiler folded the assignment into its reader — ` +
301 `yet the value is also wanted outside the loop. Read it once outside the ` +
302 `loop instead, or use it more than once inside it.`,
303 );
304}