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
6 * grid. This is exactly how numbl drives its own JIT — the caller supplies
7 * argument 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 `pn = 2*p - pm + cdt2 .* lap` becomes ONE statement whose
19 * RHS is 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 { IRFunc, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
26import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
27import { fuseTemps } from './fuse.ts';
28import { externalOpFiles } from './externals.ts';
29import type { OpInstance } from './ops.ts';
30import { ModelCompileError } from './errors.ts';
32/** What the host can supply for an argument the .m declares. */
33export type Binding =
34 /** An array, passed in a GPU buffer. */
35 | { kind: 'tensor'; shape: number[] }
36 /** A tunable scalar. Deliberately carries no exact value: an exact scalar
37 * would be constant-folded into the kernels, so moving a slider would force
38 * a recompile instead of just rewriting a uniform. */
39 | { kind: 'param' }
40 /** A fixed scalar, exact so array constructors reading it keep static
41 * shapes. */
42 | { kind: 'const'; value: number };
44const typeOf = (b: Binding): Type => {
45 switch (b.kind) {
46 case 'tensor':
47 return tensorDouble(b.shape);
48 case 'param':
49 return scalarDouble('unknown');
50 case 'const':
51 // Carry the sign too: numbl's sign lattice decides, for instance,
52 // whether sqrt() of a value can go complex.
53 return scalarDouble(
54 b.value > 0 ? 'positive' : b.value < 0 ? 'negative' : 'zero',
55 b.value,
56 );
57 }
58};
60/** One specialized function, as the planner consumes it. */
61export interface CompiledFunction {
62 name: string;
63 /** Declared arguments, in order, with the cName each lowered to. */
64 params: { name: string; cName: string; binding: Binding }[];
65 /** Requested outputs, in order, with the cName holding each result. */
66 outputs: { name: string; cName: string; ty: Type }[];
67 /** The lowered body. Read this only after `finish()`: the inline pass
68 * REPLACES the statement array rather than mutating it, so this is a live
69 * view of the function rather than a snapshot. */
70 readonly body: IRStmt[];
71}
73/** The shape of a `function` statement in numbl's AST. */
74interface FunctionDecl {
75 type: 'Function';
76 name: string;
77 params: string[];
78 outputs: string[];
79}
81/**
82 * A parsed model. Specialize the functions you need, then call `finish()` once
83 * — the inline pass rewrites every specialization together.
84 */
85export class CompiledModel {
86 #lowerer: Lowerer;
87 #decls: Map<string, FunctionDecl>;
88 #bindings: Record<string, Binding>;
90 constructor(
91 source: string,
92 bindings: Record<string, Binding>,
93 ops: Pick<OpInstance, 'name' | 'argCounts' | 'outCount'>[],
94 fileName = 'model.m',
95 ) {
96 const ast = parseMFile(source, fileName);
97 const ws = new Workspace(fileName, []);
98 ws.addFile({ name: fileName, source, ast });
99 // The scene's op instances (dxx_s1, lapw, ...) become resolvable, with
100 // their type rules.
101 for (const f of externalOpFiles(ops)) ws.addFile(f);
102 ws.finalize();
104 this.#bindings = bindings;
105 this.#lowerer = new Lowerer(ws);
106 this.#decls = new Map();
107 for (const stmt of ast.body as { type: string }[]) {
108 if (stmt.type === 'Function') {
109 const fn = stmt as unknown as FunctionDecl;
110 this.#decls.set(fn.name, fn);
111 }
112 }
113 }
115 /** Names of the functions the file defines. */
116 functionNames(): string[] {
117 return [...this.#decls.keys()];
118 }
120 /**
121 * Lower `name` for the current bindings, requesting `nargout` outputs.
122 * Every declared parameter must name something the host provides.
123 */
124 specialize(name: string, nargout: number): CompiledFunction {
125 const decl = this.#decls.get(name);
126 if (!decl) {
127 const defined = this.functionNames();
128 throw new ModelCompileError(
129 `the model must define a function named '${name}'` +
130 (defined.length
131 ? ` (it defines ${defined.map((n) => `'${n}'`).join(', ')})`
132 : ' (it defines no functions)'),
133 );
134 }
135 if (decl.outputs.length < nargout) {
136 throw new ModelCompileError(
137 `'${name}' must return ${nargout} value${nargout === 1 ? '' : 's'}, ` +
138 `but declares ${decl.outputs.length}`,
139 );
140 }
142 const bindings = decl.params.map((p) => {
143 const b = this.#bindings[p];
144 if (!b) {
145 const offered = Object.keys(this.#bindings).join(', ');
146 throw new ModelCompileError(
147 `'${name}' takes an argument named '${p}', which this app does not ` +
148 `provide. Available: ${offered}.`,
149 );
150 }
151 return b;
152 });
154 const fn: IRFunc = specializeUserFunction.call(
155 this.#lowerer,
156 decl,
157 bindings.map(typeOf),
158 undefined,
159 undefined,
160 undefined,
161 nargout,
162 undefined,
163 );
165 return {
166 name,
167 params: fn.params.map((p, i) => ({
168 name: p,
169 cName: fn.cParams[i],
170 binding: bindings[i],
171 })),
172 outputs: fn.outputs.slice(0, nargout).map((o, i) => ({
173 name: o,
174 cName: fn.cOutputs[i],
175 ty: fn.outputTypes[i],
176 })),
177 // A getter, not a snapshot: `finish()` runs after every specialization
178 // and swaps in a rewritten statement array.
179 get body() {
180 return fn.body;
181 },
182 };
183 }
185 /**
186 * Run the fusion passes over everything specialized so far. They rewrite the
187 * function bodies in place, so `CompiledFunction`s handed out earlier are
188 * updated too.
189 *
190 * Two of them: numbl's, which folds the temps its own C backend would fuse,
191 * and this project's (src/mgpu/fuse.ts), which folds the ones it declines —
192 * `sin`, `exp`, `tanh` and the rest, which WGSL evaluates per element just
193 * as happily as it does a multiply.
194 *
195 * Neither touches a variable the .m names. In particular a bare `pold = p;`
196 * — the line that turns this step's field into the next step's history —
197 * survives as its own statement, and plans as the copy it is: numbl's pass
198 * gives every declared output a protective use count, and ours only folds
199 * compiler temps.
200 */
201 finish(): void {
202 inlinePass({ topLevelStmts: [], functions: this.#lowerer.specializations });
203 for (const fn of this.#lowerer.specializations.values()) {
204 fn.body = fuseTemps(fn.body);
205 }
206 }
207}