e01ddf1MATLAB-syntax scripts on WebGPU: fused kernels, tic/toc timing, CPU comparisonJeremy Magland 1/**
2 * MATLAB script source -> numbl's JIT IR, ready for the WGSL planner.
3 *
4 * Unlike turing-surface — which specializes named functions against
5 * host-supplied argument types — the sandbox lowers a whole *script*: shapes
6 * come from the script itself (`n = 2048; A = rand(n);`), pinned static by
7 * numbl's exact-value propagation through the type lattice.
8 *
9 * Two numbl passes matter here:
10 * - `lowerProgram` lowers the top-level statements to IR, one statement per
11 * operation (ANF), with every node's type fixed.
12 * - `inlinePass` then folds single-use temps back into their consumer, so a
13 * source line like `y = a - u + u.*u.*v` becomes ONE statement whose RHS is
14 * an expression tree — i.e. one fused GPU kernel instead of four.
15 */
16import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
17import { Workspace, Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
18import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
19import type { IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
20import { applyBuiltinPatches } from './patches.ts';
21import { fuseTemps } from './fuse.ts';
22import { inScript, ScriptCompileError } from './errors.ts';
24export interface CompiledScript {
25 /** The lowered, inline-folded top-level statements. */
26 stmts: IRStmt[];
27 /**
28 * Should the statement covering source offset `at` echo its result,
29 * MATLAB-style? True exactly when the source statement has no trailing
30 * semicolon. Compiler temps sit inside their source statement's span, so
31 * the caller must additionally skip `_mtoc2_*` names.
32 */
33 isEchoed(at: number): boolean;
34}
36export function compileScript(source: string, fileName = 'script.m'): CompiledScript {
37 applyBuiltinPatches();
39 const ast = inScript(() => parseMFile(source, fileName));
41 // Statements the parser marked unsuppressed (no `;`), by source range.
42 const echoed: { start: number; end: number }[] = [];
43 for (const s of ast.body) {
44 if (s.suppressed === false) echoed.push({ start: s.span.start, end: s.span.end });
45 }
47 if (ast.body.some((s) => s.type === 'Function')) {
48 // A script may syntactically end with local functions, but nothing here
49 // compiles calls to them — say so up front rather than at the call site.
50 throw new ScriptCompileError(
51 'local functions are not supported in the sandbox yet — inline their bodies',
52 );
53 }
55 const prog = inScript(() => {
56 const ws = new Workspace(fileName, []);
57 ws.addFile({ name: fileName, source, ast });
58 ws.finalize();
59 const lowered = new Lowerer(ws).lowerProgram(ast);
60 inlinePass(lowered);
61 // numbl's pass stops at what its C backend fuses; fold the rest of the
62 // single-use temps the WGSL emitter can absorb (sin/exp, comparisons,
63 // logicals, generators) so one source line is one kernel.
64 fuseTemps(lowered.topLevelStmts);
65 return lowered;
66 });
68 return {
69 stmts: prog.topLevelStmts,
70 isEchoed: (at) => echoed.some((r) => at >= r.start && at <= r.end),
71 };
72}