concept-collection / turing-sphere
Write the solver in MATLAB and compile it to WebGPU
The IMEX reaction-diffusion loop moves out of TypeScript and into .m files under models/, one per model. numbl parses and lowers them in the browser; each element-wise line becomes a WebGPU compute kernel. The MATLAB is editable on the page, so changing the reaction changes the pattern live. How it fits together (src/mgpu/): - compile.ts specializes init() and step() for the current grid's argument types, via the same specializeUserFunction entry point numbl's own JIT uses, then runs numbl's inline pass. Lowering emits one statement per operator (ANF); the inline pass folds those back into per-line expression trees, so `uuv = u .* u .* v` is one kernel rather than three. - wgsl.ts emits one kernel per element-wise statement, the WebGPU counterpart of numbl's C-side fused emitter. Anything it cannot express is refused at compile time with a position in the file. - plan.ts builds every pipeline, buffer and bind group once. Because numbl fixes types and shapes at lowering time, the op sequence is static, so a timestep is pure synchronous command recording: one submit per batch, and the only await in the loop is the single readback per rendered frame. - externals.ts gives numbl the type rules for synth/analys through a .mtoc2.js workspace file, its sanctioned JS-builtin extension point; the backend maps those calls onto the existing transform pipelines. Argument binding goes through the MATLAB signatures, so a model file declares what it needs and an unknown parameter name is a compile error rather than a silently undefined variable. Tunable scalars are lowered without exact values so they live in a uniform buffer: editing the MATLAB recompiles, moving a slider does not. src/solver/ stays as an independent implementation of the same scheme and becomes the test oracle. test/mgpuChecks.ts runs both from the same seeded perturbation through the same fp32 transforms: the three models agree to 5.5e-8 - 3.2e-7 relative L2 after ten steps, and match exactly at init. It also asserts the per-model kernel count, which catches fusion silently not happening - a regression that leaves results correct but makes every operator its own dispatch. WebGPU is now required; there is no CPU fallback in the app. The editor gets MATLAB syntax highlighting, by overlaying a highlighted <pre> behind a transparent textarea.
Jeremy Magland <jmagland@flatironinstitute.org> committed commit 61e12f115438 parent 182fa98 Browse files
29 changed files+3077−179
.github/workflows/ci.ymlmodified+19−1View file
@@ -13,7 +13,25 @@ jobs:
1313 with:
1414 node-version: 24
1515 cache: npm
16- - run: npm ci
16+ # numbl is a `file:../../numbl` dependency: we use its compiler internals
17+ # (parser, lowerer, IR, inline pass), which its published package `exports`
18+ # do not expose. Clone it where that relative path expects it. Pinned so a
19+ # change to those internals cannot silently break the build — the surface we
20+ # rely on is written down in src/mgpu/numbl.d.ts.
21+ #
22+ # numbl's own dependencies are NOT needed: the slice we import is
23+ # self-contained TypeScript, verified by building against a checkout with no
24+ # node_modules.
25+ - name: Check out numbl (sibling dependency)
26+ env:
27+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
28+ run: |
29+ git clone --filter=blob:none --no-checkout \
30+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
31+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
32+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
33+ # numbl's is husky, which is not installed here.
34+ - run: npm ci --ignore-scripts
1735 - run: npm run test:node
1836 # headless Chrome + SwiftShader software WebGPU
1937 - run: npm run test:gpu
.github/workflows/deploy.ymlmodified+22−1View file
@@ -25,10 +25,31 @@ jobs:
2525 with:
2626 node-version: 24
2727 cache: npm
28- - run: npm ci
28+ # numbl is a `file:../../numbl` dependency: we use its compiler internals
29+ # (parser, lowerer, IR, inline pass), which its published package `exports`
30+ # do not expose. Clone it where that relative path expects it. Pinned so a
31+ # change to those internals cannot silently break the build — the surface we
32+ # rely on is written down in src/mgpu/numbl.d.ts.
33+ #
34+ # numbl's own dependencies are NOT needed: the slice we import is
35+ # self-contained TypeScript, verified by building against a checkout with no
36+ # node_modules.
37+ - name: Check out numbl (sibling dependency)
38+ env:
39+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
40+ run: |
41+ git clone --filter=blob:none --no-checkout \
42+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
43+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
44+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
45+ # numbl's is husky, which is not installed here.
46+ - run: npm ci --ignore-scripts
2947 - run: npm run test:node
3048 - run: npm run build
3149 - uses: actions/configure-pages@v5
50+ # turn Pages on for a fresh repo without a manual visit to Settings
51+ with:
52+ enablement: true
3253 - uses: actions/upload-pages-artifact@v3
3354 with:
3455 path: dist
README.mdmodified+128−19View file
@@ -4,6 +4,11 @@ Reaction–diffusion systems (Turing patterns) solved **live in the browser on t
44 surface of a sphere**, using a spectral spherical-harmonic method with the
55 transforms running on the GPU via WebGPU.
66
7+The solver itself is **MATLAB**. The `.m` files under [`models/`](models/) are the
8+algorithm — [numbl](https://numbl.org) parses and lowers them in the browser, and
9+each element-wise line becomes a WebGPU compute kernel. You can edit the MATLAB
10+on the page and watch the pattern change.
11+
712 **Live demo:** <https://concept-collection.github.io/turing-sphere/>
813
914 ## What it does
@@ -29,11 +34,66 @@ You watch the patterns emerge in real time on orbitable 3D spheres (one per
2934 species, cameras synced), with pause/resume, re-seeding, live parameter editing,
3035 and colormap selection.
3136
32-Three presets are included:
37+Three models are included, one `.m` file each:
38+
39+- **[Schnakenberg](models/schnakenberg.m)** — Turing spots (unstable band
40+ 14 ≤ l ≤ 40, peak l = 24)
41+- **[Brusselator](models/brusselator.m)** — stripes and spots from a stiffer reaction
42+- **[Allen–Cahn](models/allencahn.m)** — a single species whose interfaces form
43+ and coarsen
44+
45+## MATLAB, compiled to WebGPU
3346
34-- **Schnakenberg** — Turing spots (unstable band 14 ≤ l ≤ 40, peak l = 24)
35-- **Brusselator** — stripes and spots from a stiffer reaction
36-- **Allen–Cahn** — a single species whose interfaces form and coarsen
47+A model file is ordinary MATLAB defining two functions — `init` builds the initial
48+spectral state, `step` advances it one timestep:
49+
50+```matlab
51+function [Un, Vn, u, v] = step(U, V, lam, a, b, D1, D2, dt)
52+ u = synth(U);
53+ v = synth(V);
54+ uuv = u .* u .* v;
55+ Un = (U + dt * analys(a - u + uuv)) ./ (1 + (dt * D1) * lam);
56+ Vn = (V + dt * analys(b - uuv)) ./ (1 + (dt * D2) * lam);
57+end
58+```
59+
60+Getting from there to the GPU uses numbl for everything up to the IR, and this
61+repo only for the backend:
62+
63+1. **numbl parses and lowers.** Each function is specialized for the concrete
64+ argument types of the current grid, via the same `specializeUserFunction`
65+ entry point numbl's own JIT uses. Types and array shapes are fixed at this
66+ point, so the backend never has to re-decide what an operation means.
67+2. **numbl's inline pass fuses.** Lowering emits one statement per *operator*
68+ (ANF); `inlinePass` folds single-use temps back into their consumer, so one
69+ line of MATLAB becomes one expression tree. `uuv = u .* u .* v` arrives as a
70+ single statement, not three.
71+3. **This repo emits WGSL** ([`src/mgpu/wgsl.ts`](src/mgpu/wgsl.ts)). Each
72+ element-wise statement becomes one compute kernel that computes one output
73+ element per invocation — the WebGPU counterpart of numbl's own C-side fused
74+ emitter. Anything it cannot express is refused at compile time with a source
75+ position, never silently mis-compiled.
76+4. **`synth` / `analys` are external operations.** numbl learns their type rules
77+ from a `.mtoc2.js` workspace file — its sanctioned extension point for a
78+ JS-defined builtin — and the backend maps each call onto the existing
79+ spherical-harmonic compute pipelines.
80+
81+The Schnakenberg step above compiles to 11 GPU operations: 4 transforms, 5
82+generated kernels, and 2 buffer copies feeding the new state back.
83+
84+Two consequences worth noting:
85+
86+- **The step is synchronous.** WebGPU's encode path (`writeBuffer`, dispatch,
87+ `submit`) is all synchronous; only readback and pipeline creation are async, and
88+ every pipeline is built once at compile time. So a timestep is pure command
89+ recording — the whole batch goes out in one submit, and the only `await` in the
90+ loop is the single readback per rendered frame. numbl's own execution being
91+ synchronous is therefore not an obstacle: nothing about the algorithm needs to
92+ block.
93+- **Parameters are uniforms, not constants.** Tunable scalars are deliberately
94+ lowered without exact values, so moving a slider rewrites a small buffer
95+ instead of triggering a recompile. Editing the MATLAB recompiles; changing `dt`
96+ does not.
3797
3898 ## Provenance
3999
@@ -47,14 +107,23 @@ compute, so this port swaps in:
47107 fp32 spherical harmonic transforms in WGSL compute shaders, modeled on
48108 [SHTNS](https://nschaeff.bitbucket.io/shtns/). Its source is vendored under
49109 [`src/sht/`](src/sht/) (CECILL-2.1), including the f64 CPU reference
50- transform used for testing and as a no-WebGPU fallback.
110+ transform used for testing.
51111 - **Rendering:** three.js spheres with per-vertex colormaps, adapted from the
52112 `SphereEmbedding` view in
53113 [figpack](https://github.com/flatironinstitute/figpack)'s experimental
54114 extension package ([`src/render/`](src/render/)).
55-- **Solver:** [`src/solver/simulation.ts`](src/solver/simulation.ts), a direct
56- TypeScript port of the MATLAB IMEX loop, in f64 on the coefficients with the
57- transforms in fp32 on the GPU.
115+- **Solver:** the MATLAB stayed MATLAB. [`models/`](models/) holds the IMEX loop
116+ as `.m` files, executed on the GPU by [`src/mgpu/`](src/mgpu/).
117+
118+[`src/solver/`](src/solver/) still holds the TypeScript port of the same loop. The
119+app no longer runs it, but it is an *independent* implementation of the scheme,
120+which makes it the test oracle: `npm run test:gpu` runs both from the same seeded
121+perturbation through the same transforms and compares. It is also where parameter
122+metadata (names, defaults, slider ranges) lives, so the two paths cannot be
123+configured differently.
124+
125+Because the algorithm is now compiled to compute shaders, **WebGPU is required** —
126+there is no CPU fallback in the app (the f64 CPU transform remains, for tests).
58127
59128 ## Numerics
60129
@@ -71,10 +140,9 @@ compute, so this port swaps in:
71140 ## Desktop vs browser
72141
73142 How much does running this in a browser cost? [`scripts/bench.ts`](scripts/bench.ts)
74-answers that by running the *same* code — same `Simulation`, same WGSL
75-transforms, same parameters — from Node on desktop WebGPU (Google Dawn), and
76-the app prints the command line that reproduces whatever it is currently
77-simulating:
143+runs the reference solver — same WGSL transforms, same parameters — from Node on
144+desktop WebGPU (Google Dawn), and the app prints the command line that
145+reproduces whatever it is currently simulating:
78146
79147 ```
80148 node scripts/bench.mjs --preset schnak-spots --lmax 63 --backend webgpu --steps 2000 \
@@ -102,12 +170,21 @@ compare against Dawn's own software adapter.
102170
103171 What the comparison does and does not control for:
104172
105-- the benchmark is **solver only**; the app's `ms/step` excludes `draw()` but is
106- still measured on a page that renders two spheres between steps. For a browser
107- number with no rendering at all, open `test.html?soak=2000&lmax=63`.
108-- each step is four transforms, each ending in a buffer readback, so both sides
109- are dominated by submit-and-map latency rather than arithmetic — this measures
110- a driver round-trip more than it measures a GPU.
173+- **it is not the same solver.** The benchmark runs the TypeScript reference; the
174+ app runs the `.m` compiled to WGSL. Node cannot load numbl's TypeScript sources
175+ (its internal imports are extensionless-`.js`, which needs a bundler's
176+ resolution), so the `.m` path is browser-only for now. Same scheme, same
177+ transforms, same parameters — but the reaction and the IMEX update happen in f64
178+ on the CPU there and in fp32 on the GPU here.
179+- the benchmark is **solver only**; the app's `ms/step` includes the per-frame
180+ readback amortized over the step batch. For a browser number with no rendering,
181+ open `test.html?soak=2000&lmax=63` (that soak also runs the reference solver).
182+- the reference pays a buffer readback on *every* transform — four driver
183+ round-trips per step — so it measures submit-and-map latency more than
184+ arithmetic. The `.m` path keeps everything in GPU buffers and submits once per
185+ batch, which is where its advantage should come from. On the software rasterizer
186+ in CI the transforms dominate and the two come out within ~10% of each other;
187+ the gap on real hardware is untested.
111188 - the browser adds its own GPU-process boundary and, for a page that is not
112189 cross-origin isolated, coarser timers.
113190
@@ -119,7 +196,10 @@ What the comparison does and does not control for:
119196 linear recurrence, exact uniform-state reaction ODE, and the linearized
120197 Turing-mode 2×2 IMEX recurrence (all at ~1e-12).
121198 - `npm run test:gpu` — builds and drives headless Chrome: GPU-vs-CPU transform
122- and solver cross-checks, plus a 100-step stability run.
199+ and solver cross-checks, a 100-step stability run, and then for every `.m`
200+ model: that it compiles, that its element-wise lines each fuse into exactly one
201+ kernel, and that 10 steps agree with the reference solver from the same seed
202+ (they agree to ~1e-7 relative L2 — fp32 round-off).
123203 - `node scripts/longrun-node.ts` — CPU run to t = 100 confirming pattern
124204 saturation.
125205 - `node scripts/soak.mjs [steps] [lmax] [backend]` — drive the demo for many
@@ -149,6 +229,35 @@ npm run dev # local dev server
149229 npm run build # type-check + production build to dist/
150230 ```
151231
232+### The numbl dependency
233+
234+numbl is a local `file:../../numbl` dependency, so a sibling checkout of
235+[numbl](https://github.com/flatironinstitute/numbl) is required. We use its
236+compiler internals — parser, lowerer, IR, inline pass — which its package
237+`exports` map does not publish, so they are reached through the `numbl-src` path
238+alias in [`vite.config.ts`](vite.config.ts).
239+
240+The exact surface we depend on is written down in
241+[`src/mgpu/numbl.d.ts`](src/mgpu/numbl.d.ts) and TypeScript checks against
242+*that*, not against numbl's sources. This keeps this project's compiler settings
243+independent of numbl's (its sources do not type-check under the stricter options
244+used here), and means a change to one of those shapes upstream breaks the build
245+here with a clear diff rather than deep inside numbl's tree.
246+
247+The compiler is ~395 kB gzipped and lands in its own chunk. That is the cost of
248+compiling MATLAB in the page; a build-time lowering step could remove it at the
249+price of no longer being editable live.
250+
251+CI clones numbl to the sibling path that `file:` dependency expects, pinned to a
252+commit. Two details make that work, both verified by building against a checkout
253+that had none of numbl's own dependencies installed:
254+
255+- **numbl's `node_modules` are not needed.** The slice we import — parser,
256+ lowering, IR, inline pass — is self-contained TypeScript. (Other parts of numbl
257+ do import `three`, `react` and `fflate`; we never reach them.)
258+- **the install must pass `--ignore-scripts`.** npm runs a linked package's
259+ `prepare` script, and numbl's is `husky`, which is not installed in CI.
260+
152261 The `.ts` entry points under `scripts/` are run by Node directly, which strips
153262 types without being asked only from Node 22.18 / 23.6 / 24 on. Everything here
154263 works back to 22.6, where stripping exists but is flagged: the npm scripts pass
index.htmlmodified+89−8View file
@@ -13,6 +13,11 @@
1313 --line: #d0d7de;
1414 --accent: #0969da;
1515 --sphere-bg: #f4f6f8;
16+ --tok-com: #6e7781;
17+ --tok-str: #0a3069;
18+ --tok-num: #0550ae;
19+ --tok-kw: #cf222e;
20+ --tok-ext: #8250df;
1621 color-scheme: light dark;
1722 }
1823 @media (prefers-color-scheme: dark) {
@@ -23,6 +28,11 @@
2328 --line: #333b44;
2429 --accent: #58a6ff;
2530 --sphere-bg: #14161c;
31+ --tok-com: #8b949e;
32+ --tok-str: #a5d6ff;
33+ --tok-num: #79c0ff;
34+ --tok-kw: #ff7b72;
35+ --tok-ext: #d2a8ff;
2636 }
2737 }
2838 body {
@@ -92,6 +102,59 @@
92102 }
93103 #blurb { margin-top: 4px; font-size: 13px; color: var(--ink-2); }
94104 #err { color: #b35900; white-space: pre-wrap; font-size: 13px; }
105+ .editor {
106+ margin-top: 12px; border: 1px solid var(--line); border-radius: 8px;
107+ overflow: hidden;
108+ }
109+ .editor-head {
110+ display: flex; gap: 10px; align-items: center; justify-content: space-between;
111+ padding: 6px 10px; font-size: 12px; color: var(--ink-2);
112+ background: var(--sphere-bg); border-bottom: 1px solid var(--line);
113+ }
114+ .editor-head button { padding: 2px 10px; font-size: 12px; }
115+ /* Source and compiled-op list side by side, so the editor gets the
116+ height rather than sharing it with the list below. */
117+ .editor-body { display: flex; align-items: stretch; }
118+ .editor-code { position: relative; flex: 1 1 62%; min-width: 0; height: 34em; }
119+ /* The overlay and the textarea must agree on every metric that affects
120+ where a character lands. Keep these two rules together. */
121+ .editor-code > pre,
122+ .editor-code > textarea {
123+ margin: 0; padding: 10px 12px; border: 0;
124+ box-sizing: border-box; width: 100%; height: 100%;
125+ font: 12.5px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
126+ tab-size: 2;
127+ white-space: pre; overflow-wrap: normal;
128+ }
129+ #highlight {
130+ position: absolute; inset: 0; overflow: hidden;
131+ pointer-events: none; background: var(--bg); color: var(--ink);
132+ }
133+ #source {
134+ position: relative; z-index: 1; display: block;
135+ resize: none; overflow: auto;
136+ background: transparent; color: transparent; caret-color: var(--ink);
137+ }
138+ #source:focus { outline: none; }
139+ /* Transparent text means the selection must be see-through, or selected
140+ code would be invisible. */
141+ #source::selection { background: color-mix(in srgb, var(--accent) 28%, transparent); }
142+ #compiled {
143+ flex: 1 1 38%; min-width: 0; margin: 0; padding: 10px 12px; overflow: auto;
144+ border-left: 1px solid var(--line); background: var(--sphere-bg);
145+ font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
146+ color: var(--ink-2); white-space: pre;
147+ }
148+ @media (max-width: 860px) {
149+ .editor-body { flex-direction: column; }
150+ .editor-code { flex: none; height: 26em; }
151+ #compiled { border-left: 0; border-top: 1px solid var(--line); max-height: 12em; }
152+ }
153+ .tok-com { color: var(--tok-com); }
154+ .tok-str { color: var(--tok-str); }
155+ .tok-num { color: var(--tok-num); }
156+ .tok-kw { color: var(--tok-kw); font-weight: 600; }
157+ .tok-ext { color: var(--tok-ext); }
95158 </style>
96159 </head>
97160 <body>
@@ -101,7 +164,9 @@
101164 Reaction-diffusion on the sphere, solved live with spherical harmonics:
102165 implicit spectral diffusion + explicit reaction (IMEX Euler), transforms on WebGPU via
103166 <a href="https://github.com/concept-collection/shtns-webgpu">shtns-webgpu</a>.
104- Drag to rotate.
167+ The solver itself is the MATLAB below — <a href="https://numbl.org">numbl</a>
168+ parses and lowers it in your browser, and each line becomes a WebGPU
169+ compute kernel. Edit it and the pattern changes. Drag to rotate.
105170 </p>
106171 <div class="controls">
107172 <label>preset
@@ -118,12 +183,6 @@
118183 <label>colormap
119184 <select id="colormap"></select>
120185 </label>
121- <label>backend
122- <select id="backend">
123- <option value="webgpu" selected>WebGPU (fp32)</option>
124- <option value="cpu">CPU (f64)</option>
125- </select>
126- </label>
127186 <button id="runpause" class="primary">Run</button>
128187 <button id="reseed">Re-seed</button>
129188 <button id="resetview">Reset view</button>
@@ -131,9 +190,31 @@
131190 <div class="controls" id="params"></div>
132191 <div id="panels"></div>
133192 <p class="stats" id="stats"></p>
193+ <div class="editor">
194+ <div class="editor-head">
195+ <span id="editor-title">the solver, in MATLAB</span>
196+ <span>
197+ <button id="recompile" type="button">Recompile</button>
198+ <button id="revert" type="button">Revert</button>
199+ </span>
200+ </div>
201+ <div class="editor-body">
202+ <div class="editor-code">
203+ <pre id="highlight" aria-hidden="true"></pre>
204+ <textarea
205+ id="source"
206+ spellcheck="false"
207+ autocomplete="off"
208+ autocapitalize="off"
209+ aria-label="model source (MATLAB)"
210+ ></textarea>
211+ </div>
212+ <pre id="compiled"></pre>
213+ </div>
214+ </div>
134215 <div class="cli">
135216 <div class="cli-head">
136- <span>Same run on the desktop, solver only — compare its ms/step with the one above</span>
217+ <span>The same parameters on the desktop, via the reference TypeScript solver</span>
137218 <button id="copycmd" type="button">Copy</button>
138219 </div>
139220 <code id="cmd"></code>
models/allencahn.madded+22−0View file
@@ -0,0 +1,22 @@
1+% Allen-Cahn on the unit sphere. One species: interfaces form and then coarsen
2+% until one domain swallows the sphere.
3+%
4+% du/dt = eps2*lap(u) + u - u^3
5+%
6+% Diffusion is implicit in spherical-harmonic space, where lap is diagonal with
7+% eigenvalues -l(l+1); the reaction is explicit on the grid, giving one IMEX
8+% Euler step. Provided by the caller: synth/analys (the transforms), lam =
9+% l(l+1) per coefficient, noise (the seeded perturbation), and the parameters.
10+% Grid fields are npts x 1; spectral fields are real 2 x nlm (row 1 real part,
11+% row 2 imaginary), so no complex arithmetic is needed. Each function returns
12+% the new spectral state followed by the grid fields to display.
13+
14+function [U, u] = init(noise)
15+ U = analys(noise);
16+ u = synth(U);
17+end
18+
19+function [Un, u] = step(U, lam, eps2, dt)
20+ u = synth(U);
21+ Un = (U + dt * analys(u - u.^3)) ./ (1 + (dt * eps2) * lam);
22+end
models/brusselator.madded+28−0View file
@@ -0,0 +1,28 @@
1+% Brusselator reaction-diffusion on the unit sphere. Turing stripes and spots,
2+% from a smaller diffusivity contrast than Schnakenberg but a stiffer reaction.
3+%
4+% du/dt = D1*lap(u) + A - (B+1)*u + u^2*v
5+% dv/dt = D2*lap(v) + B*u - u^2*v
6+%
7+% Diffusion is implicit in spherical-harmonic space, where lap is diagonal with
8+% eigenvalues -l(l+1); the reaction is explicit on the grid, giving one IMEX
9+% Euler step. Provided by the caller: synth/analys (the transforms), lam =
10+% l(l+1) per coefficient, noise (the seeded perturbation), and the parameters.
11+% Grid fields are npts x 1; spectral fields are real 2 x nlm (row 1 real part,
12+% row 2 imaginary), so no complex arithmetic is needed. Each function returns
13+% the new spectral state followed by the grid fields to display.
14+
15+function [U, V, u, v] = init(noise, A, B)
16+ U = analys(A + noise);
17+ V = analys((B / A) * ones(numel(noise), 1));
18+ u = synth(U);
19+ v = synth(V);
20+end
21+
22+function [Un, Vn, u, v] = step(U, V, lam, A, B, D1, D2, dt)
23+ u = synth(U);
24+ v = synth(V);
25+ uuv = u .* u .* v;
26+ Un = (U + dt * analys(A - (B + 1) * u + uuv)) ./ (1 + (dt * D1) * lam);
27+ Vn = (V + dt * analys(B * u - uuv)) ./ (1 + (dt * D2) * lam);
28+end
models/schnakenberg.madded+29−0View file
@@ -0,0 +1,29 @@
1+% Schnakenberg reaction-diffusion on the unit sphere.
2+%
3+% du/dt = D1*lap(u) + a - u + u^2*v
4+% dv/dt = D2*lap(v) + b - u^2*v
5+%
6+% Diffusion is implicit in spherical-harmonic space, where lap is diagonal with
7+% eigenvalues -l(l+1); the reaction is explicit on the grid, giving one IMEX
8+% Euler step. Provided by the caller: synth/analys (the transforms), lam =
9+% l(l+1) per coefficient, noise (the seeded perturbation), and the parameters.
10+% Grid fields are npts x 1; spectral fields are real 2 x nlm (row 1 real part,
11+% row 2 imaginary), so no complex arithmetic is needed. Each function returns
12+% the new spectral state followed by the grid fields to display.
13+
14+function [U, V, u, v] = init(noise, a, b)
15+ us = a + b;
16+ vs = b / (us * us);
17+ U = analys(us + noise);
18+ V = analys(vs * ones(numel(noise), 1));
19+ u = synth(U);
20+ v = synth(V);
21+end
22+
23+function [Un, Vn, u, v] = step(U, V, lam, a, b, D1, D2, dt)
24+ u = synth(U);
25+ v = synth(V);
26+ uuv = u .* u .* v;
27+ Un = (U + dt * analys(a - u + uuv)) ./ (1 + (dt * D1) * lam);
28+ Vn = (V + dt * analys(b - uuv)) ./ (1 + (dt * D2) * lam);
29+end
package-lock.jsonmodified+84−0View file
@@ -9,6 +9,7 @@
99 "version": "0.1.0",
1010 "license": "CECILL-2.1",
1111 "dependencies": {
12+ "numbl": "file:../../numbl",
1213 "three": "^0.183.0"
1314 },
1415 "devDependencies": {
@@ -19,10 +20,89 @@
1920 "typescript": "^5.5.0",
2021 "vite": "^5.4.0"
2122 },
23+ "engines": {
24+ "node": ">=22.6"
25+ },
2226 "optionalDependencies": {
2327 "webgpu": "^0.4.0"
2428 }
2529 },
30+ "../../numbl": {
31+ "version": "0.4.18",
32+ "hasInstallScript": true,
33+ "license": "Apache-2.0",
34+ "dependencies": {
35+ "fflate": "^0.8.2",
36+ "h5wasm": "^0.10.3",
37+ "node-addon-api": "^8.3.0",
38+ "pako": "^2.1.0",
39+ "qhull-wasm": "^0.0.1",
40+ "react-markdown": "^10.1.0",
41+ "react-syntax-highlighter": "^16.1.1",
42+ "remark-gfm": "^4.0.1",
43+ "three": "^0.183.2"
44+ },
45+ "bin": {
46+ "numbl": "dist-cli/cli.js"
47+ },
48+ "devDependencies": {
49+ "@emotion/react": "^11.14.0",
50+ "@emotion/styled": "^11.14.1",
51+ "@eslint/js": "^9.39.1",
52+ "@monaco-editor/react": "^4.7.0",
53+ "@mui/icons-material": "^7.3.7",
54+ "@mui/material": "^7.3.7",
55+ "@playwright/test": "^1.59.1",
56+ "@types/node": "^24.10.1",
57+ "@types/pako": "^2.0.4",
58+ "@types/react": "^19.2.5",
59+ "@types/react-dom": "^19.2.3",
60+ "@types/react-syntax-highlighter": "^15.5.13",
61+ "@types/three": "^0.183.1",
62+ "@vitejs/plugin-react": "^5.1.1",
63+ "@vitest/coverage-v8": "^4.0.18",
64+ "@xterm/addon-fit": "^0.11.0",
65+ "@xterm/xterm": "^6.0.0",
66+ "dexie": "^4.3.0",
67+ "esbuild": "^0.27.3",
68+ "eslint": "^9.39.1",
69+ "eslint-config-prettier": "^10.1.8",
70+ "eslint-plugin-react-hooks": "^7.0.1",
71+ "eslint-plugin-react-refresh": "^0.4.24",
72+ "globals": "^16.5.0",
73+ "husky": "^9.1.7",
74+ "lint-staged": "^16.2.7",
75+ "node-gyp": "^11.2.0",
76+ "prettier": "^3.8.1",
77+ "react": "^19.2.0",
78+ "react-dom": "^19.2.0",
79+ "react-router-dom": "^7.13.0",
80+ "tsx": "^4.21.0",
81+ "typescript": "~5.9.3",
82+ "typescript-eslint": "^8.46.4",
83+ "vite": "^7.2.4",
84+ "vite-plugin-wasm": "^3.5.0",
85+ "vitest": "^4.0.18"
86+ },
87+ "engines": {
88+ "node": ">=18"
89+ },
90+ "optionalDependencies": {
91+ "koffi": "^2.15.2"
92+ },
93+ "peerDependencies": {
94+ "react": ">=18",
95+ "react-dom": ">=18"
96+ },
97+ "peerDependenciesMeta": {
98+ "react": {
99+ "optional": true
100+ },
101+ "react-dom": {
102+ "optional": true
103+ }
104+ }
105+ },
26106 "node_modules/@dimforge/rapier3d-compat": {
27107 "version": "0.12.0",
28108 "resolved": "https://registry.npmjs.org/@dimforge/rapier3d-compat/-/rapier3d-compat-0.12.0.tgz",
@@ -1587,6 +1667,10 @@
15871667 "node": ">= 0.4.0"
15881668 }
15891669 },
1670+ "node_modules/numbl": {
1671+ "resolved": "../../numbl",
1672+ "link": true
1673+ },
15901674 "node_modules/once": {
15911675 "version": "1.4.0",
15921676 "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz",
package.jsonmodified+1−0View file
@@ -16,6 +16,7 @@
1616 "bench": "node scripts/bench.mjs"
1717 },
1818 "dependencies": {
19+ "numbl": "file:../../numbl",
1920 "three": "^0.183.0"
2021 },
2122 "optionalDependencies": {
scripts/bench.tsmodified+2−63View file
@@ -34,6 +34,7 @@ import {
3434 DEFAULT_BACKEND,
3535 type RunSpec,
3636 } from '../src/bench/runSpec.ts';
37+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
3738
3839 const USAGE = `usage: ${BENCH_COMMAND} [options]
3940
@@ -55,7 +56,6 @@ function fail(msg: string, code = 1): never {
5556 console.error(`bench: ${msg}`);
5657 process.exit(code);
5758 }
58-const errMsg = (e: unknown): string => (e instanceof Error ? e.message : String(e));
5959
6060 // ---------------------------------------------------------------- arguments
6161 const argv = process.argv.slice(2);
@@ -71,64 +71,6 @@ try {
7171 fail(`${errMsg(e)}\n\n${USAGE}`, 2);
7272 }
7373
74-// ---------------------------------------------------------------- WebGPU
75-/**
76- * Install Dawn under the globals the transform code expects (navigator.gpu,
77- * GPUBufferUsage, ...), so src/ runs here unchanged — including
78- * requestShtDevice(), which is the same device request the browser makes.
79- * The specifier is indirect so that typechecking does not require the
80- * optional package to be installed.
81- */
82-async function installWebGpu(): Promise<string> {
83- const specifier = 'webgpu';
84- let mod: {
85- create: (flags: string[]) => GPU;
86- globals: Record<string, unknown>;
87- };
88- try {
89- mod = await import(specifier);
90- } catch (e) {
91- // Distinguish "not installed" from "installed but the prebuilt Dawn binary
92- // will not load" — the second is what a machine missing a system library
93- // looks like, and reporting it as the first sends people in circles.
94- const detail = errMsg(e);
95- if (/Cannot find (package|module) '?webgpu'?/.test(detail)) {
96- throw new Error(
97- 'desktop WebGPU needs the optional `webgpu` package (prebuilt Google Dawn):\n' +
98- ' npm install webgpu\n' +
99- 'It is an optionalDependency, so npm can skip it silently — `npm ls webgpu`\n' +
100- 'says whether it is there. Or run with --backend cpu.',
101- );
102- }
103- const glibc = /GLIBC_([0-9.]+)/.exec(detail);
104- throw new Error(
105- `the \`webgpu\` package is installed but did not load:\n ${detail}\n` +
106- (glibc
107- ? `Dawn's prebuilt binary wants glibc ${glibc[1]} or newer and this host is older\n` +
108- '(`ldd --version` says how old). No flag bridges that — use a container with a\n' +
109- 'newer base image, or a newer host.\n'
110- : 'That is usually the prebuilt Dawn binary missing a system library.\n') +
111- 'Run with --backend cpu for the f64 CPU reference instead.',
112- );
113- }
114- Object.assign(globalThis, mod.globals);
115- // DAWN_FLAGS is ';'-separated because individual Dawn options take
116- // comma-separated lists, e.g. 'enable-dawn-features=allow_unsafe_apis,timestamp_quantization'
117- const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
118- Object.defineProperty(globalThis, 'navigator', {
119- value: { gpu: mod.create(dawnFlags) },
120- configurable: true,
121- writable: true,
122- });
123- const { version } = await import(`${specifier}/package.json`, {
124- with: { type: 'json' },
125- }).then(
126- (m) => m.default as { version: string },
127- () => ({ version: '?' }),
128- );
129- return `node-webgpu ${version} (Google Dawn)`;
130-}
131-
13274 // ---------------------------------------------------------------- statistics
13375 interface Timing {
13476 meanMs: number;
@@ -183,10 +125,7 @@ try {
183125 runtime = await installWebGpu();
184126 device = await requestShtDevice().catch((e: unknown) => {
185127 throw new Error(
186- `${errMsg(e)}\n` +
187- ' Dawn reaches the GPU through Vulkan on Linux and Windows, Metal on macOS,\n' +
188- " so a headless box may have no adapter at all. DAWN_FLAGS='backend=vulkan'\n" +
189- ' makes it explain itself; --backend cpu always works.',
128+ `${errMsg(e)}\n${NO_ADAPTER_HINT}\n --backend cpu always works.`,
190129 );
191130 });
192131 adapter = await describeAdapter(device);
scripts/nodeWebGpu.tsadded+71−0View file
@@ -0,0 +1,71 @@
1+/**
2+ * Desktop WebGPU for the command-line scripts, via the optional `webgpu`
3+ * package (prebuilt Google Dawn).
4+ *
5+ * Installs Dawn under the globals the transform code expects (navigator.gpu,
6+ * GPUBufferUsage, ...) so everything under src/ runs here unchanged —
7+ * including requestShtDevice(), which makes the same device request the
8+ * browser makes.
9+ */
10+
11+export const errMsg = (e: unknown): string =>
12+ e instanceof Error ? e.message : String(e);
13+
14+/**
15+ * Returns a human-readable runtime description. The import specifier is
16+ * indirect so typechecking does not require the optional package.
17+ */
18+export async function installWebGpu(): Promise<string> {
19+ const specifier = 'webgpu';
20+ let mod: {
21+ create: (flags: string[]) => GPU;
22+ globals: Record<string, unknown>;
23+ };
24+ try {
25+ mod = await import(specifier);
26+ } catch (e) {
27+ // Distinguish "not installed" from "installed but the prebuilt Dawn binary
28+ // will not load" — the second is what a machine missing a system library
29+ // looks like, and reporting it as the first sends people in circles.
30+ const detail = errMsg(e);
31+ if (/Cannot find (package|module) '?webgpu'?/.test(detail)) {
32+ throw new Error(
33+ 'desktop WebGPU needs the optional `webgpu` package (prebuilt Google Dawn):\n' +
34+ ' npm install webgpu\n' +
35+ 'It is an optionalDependency, so npm can skip it silently — `npm ls webgpu`\n' +
36+ 'says whether it is there.',
37+ );
38+ }
39+ const glibc = /GLIBC_([0-9.]+)/.exec(detail);
40+ throw new Error(
41+ `the \`webgpu\` package is installed but did not load:\n ${detail}\n` +
42+ (glibc
43+ ? `Dawn's prebuilt binary wants glibc ${glibc[1]} or newer and this host is older\n` +
44+ '(`ldd --version` says how old). No flag bridges that — use a container with a\n' +
45+ 'newer base image, or a newer host.\n'
46+ : 'That is usually the prebuilt Dawn binary missing a system library.\n'),
47+ );
48+ }
49+ Object.assign(globalThis, mod.globals);
50+ // DAWN_FLAGS is ';'-separated because individual Dawn options take
51+ // comma-separated lists, e.g. 'enable-dawn-features=allow_unsafe_apis,...'
52+ const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
53+ Object.defineProperty(globalThis, 'navigator', {
54+ value: { gpu: mod.create(dawnFlags) },
55+ configurable: true,
56+ writable: true,
57+ });
58+ const { version } = await import(`${specifier}/package.json`, {
59+ with: { type: 'json' },
60+ }).then(
61+ (m) => m.default as { version: string },
62+ () => ({ version: '?' }),
63+ );
64+ return `node-webgpu ${version} (Google Dawn)`;
65+}
66+
67+/** The hint to print when Dawn loads but finds no adapter. */
68+export const NO_ADAPTER_HINT =
69+ ' Dawn reaches the GPU through Vulkan on Linux and Windows, Metal on macOS,\n' +
70+ " so a headless box may have no adapter at all. DAWN_FLAGS='backend=vulkan'\n" +
71+ ' makes it explain itself.';
src/editor/codeEditor.tsadded+92−0View file
@@ -0,0 +1,92 @@
1+/**
2+ * A textarea with syntax highlighting, by overlay.
3+ *
4+ * A textarea cannot colour its own text, so the highlighted source is rendered
5+ * into a <pre> underneath and the textarea sits on top with transparent text and
6+ * a visible caret. The two must agree on every metric that affects layout —
7+ * font, line height, padding, tab size, wrapping — and their scroll offsets are
8+ * kept in sync, or the colours drift away from the characters.
9+ */
10+import { highlightMatlab } from './matlab.ts';
11+
12+export interface CodeEditorOptions {
13+ textarea: HTMLTextAreaElement;
14+ /** The <pre> behind it, holding the highlighted copy. */
15+ overlay: HTMLElement;
16+ /** Names to mark as host-provided operations. */
17+ external?: ReadonlySet<string>;
18+ /** Called on every edit. */
19+ onInput?: (value: string) => void;
20+}
21+
22+export class CodeEditor {
23+ #textarea: HTMLTextAreaElement;
24+ #overlay: HTMLElement;
25+ #external: ReadonlySet<string>;
26+
27+ constructor(opts: CodeEditorOptions) {
28+ this.#textarea = opts.textarea;
29+ this.#overlay = opts.overlay;
30+ this.#external = opts.external ?? new Set();
31+
32+ this.#textarea.addEventListener('input', () => {
33+ this.#repaint();
34+ opts.onInput?.(this.#textarea.value);
35+ });
36+ // Keep the colours under the characters while scrolling.
37+ this.#textarea.addEventListener('scroll', () => this.#syncScroll());
38+ // Tab should indent rather than leave the editor.
39+ this.#textarea.addEventListener('keydown', (e) => this.#onKeyDown(e));
40+ this.#repaint();
41+ }
42+
43+ get value(): string {
44+ return this.#textarea.value;
45+ }
46+
47+ set value(next: string) {
48+ this.#textarea.value = next;
49+ this.#repaint();
50+ }
51+
52+ focus(): void {
53+ this.#textarea.focus();
54+ }
55+
56+ /** Select a character range, scrolling it into view. */
57+ select(start: number, end: number): void {
58+ this.#textarea.focus();
59+ this.#textarea.setSelectionRange(start, end);
60+ // setSelectionRange does not always scroll; nudge the line into view.
61+ const line = this.#textarea.value.slice(0, start).split('\n').length - 1;
62+ const lineHeight = this.#textarea.scrollHeight / Math.max(1, this.#lineCount());
63+ const target = line * lineHeight - this.#textarea.clientHeight / 2;
64+ this.#textarea.scrollTop = Math.max(0, target);
65+ this.#syncScroll();
66+ }
67+
68+ #lineCount(): number {
69+ return this.#textarea.value.split('\n').length + 1; // +1 for the trailing line
70+ }
71+
72+ #onKeyDown(e: KeyboardEvent): void {
73+ if (e.key !== 'Tab' || e.ctrlKey || e.metaKey || e.altKey) return;
74+ e.preventDefault();
75+ const el = this.#textarea;
76+ const { selectionStart: s, selectionEnd: t, value } = el;
77+ el.value = `${value.slice(0, s)} ${value.slice(t)}`;
78+ el.selectionStart = el.selectionEnd = s + 2;
79+ // Let the input listener repaint and notify, as for any other edit.
80+ el.dispatchEvent(new Event('input'));
81+ }
82+
83+ #repaint(): void {
84+ this.#overlay.innerHTML = highlightMatlab(this.#textarea.value, this.#external);
85+ this.#syncScroll();
86+ }
87+
88+ #syncScroll(): void {
89+ this.#overlay.scrollTop = this.#textarea.scrollTop;
90+ this.#overlay.scrollLeft = this.#textarea.scrollLeft;
91+ }
92+}
src/editor/matlab.tsadded+213−0View file
@@ -0,0 +1,213 @@
1+/**
2+ * A small MATLAB tokenizer, for syntax highlighting the model editor.
3+ *
4+ * Only what highlighting needs — comments, literals, numbers, keywords — and
5+ * deliberately not a parser: numbl does the real parsing, and reports errors
6+ * with positions. Tokens preserve the source text exactly, character for
7+ * character, because the highlighted output is overlaid on a textarea and any
8+ * dropped or added character would shift the two out of alignment.
9+ */
10+
11+export type TokenClass = 'com' | 'str' | 'num' | 'kw' | 'ext';
12+
13+export interface Token {
14+ text: string;
15+ cls: TokenClass | null;
16+}
17+
18+const KEYWORDS = new Set([
19+ 'break', 'case', 'catch', 'classdef', 'continue', 'else', 'elseif', 'end',
20+ 'for', 'function', 'global', 'if', 'otherwise', 'parfor', 'persistent',
21+ 'return', 'spmd', 'switch', 'try', 'while',
22+]);
23+
24+const isIdentStart = (c: string): boolean => /[A-Za-z_]/.test(c);
25+const isIdent = (c: string): boolean => /[A-Za-z0-9_]/.test(c);
26+const isDigit = (c: string): boolean => c >= '0' && c <= '9';
27+
28+/**
29+ * In MATLAB `'` is both the transpose operator and the char-literal delimiter.
30+ * It opens a literal unless it directly follows something that can be
31+ * transposed — a value, a closing bracket, or another transpose.
32+ */
33+function quoteIsTranspose(src: string, at: number): boolean {
34+ for (let i = at - 1; i >= 0; i--) {
35+ const c = src[i];
36+ if (c === ' ' || c === '\t') continue;
37+ return isIdent(c) || c === ')' || c === ']' || c === '}' || c === '.' || c === "'";
38+ }
39+ return false;
40+}
41+
42+/**
43+ * Tokenize `src`. `external` names (the operations the host provides, e.g.
44+ * `synth` / `analys`) get their own class so the boundary between the model and
45+ * what it is given is visible in the editor.
46+ */
47+export function tokenizeMatlab(
48+ src: string,
49+ external: ReadonlySet<string> = new Set(),
50+): Token[] {
51+ const out: Token[] = [];
52+ const push = (text: string, cls: TokenClass | null): void => {
53+ if (!text) return;
54+ const last = out[out.length - 1];
55+ if (last && last.cls === cls) last.text += text;
56+ else out.push({ text, cls });
57+ };
58+
59+ let i = 0;
60+ let atLineStart = true;
61+ let inBlockComment = false;
62+
63+ while (i < src.length) {
64+ const c = src[i];
65+
66+ // Block comments: `%{` and `%}` each alone on their line.
67+ if (atLineStart) {
68+ const eol = src.indexOf('\n', i);
69+ const lineEnd = eol === -1 ? src.length : eol;
70+ const line = src.slice(i, lineEnd);
71+ const trimmed = line.trim();
72+ if (!inBlockComment && trimmed === '%{') inBlockComment = true;
73+ else if (inBlockComment && trimmed === '%}') {
74+ push(line, 'com');
75+ i = lineEnd;
76+ inBlockComment = false;
77+ atLineStart = false;
78+ continue;
79+ }
80+ if (inBlockComment) {
81+ push(line, 'com');
82+ i = lineEnd;
83+ atLineStart = false;
84+ continue;
85+ }
86+ }
87+
88+ if (c === '\n') {
89+ push(c, null);
90+ i++;
91+ atLineStart = true;
92+ continue;
93+ }
94+ if (c === ' ' || c === '\t') {
95+ push(c, null);
96+ i++;
97+ continue;
98+ }
99+ atLineStart = false;
100+
101+ // Line comment, including MATLAB's `%%` section markers.
102+ if (c === '%') {
103+ const eol = src.indexOf('\n', i);
104+ const end = eol === -1 ? src.length : eol;
105+ push(src.slice(i, end), 'com');
106+ i = end;
107+ continue;
108+ }
109+
110+ // Line continuation is an operator, but any trailing text is a comment.
111+ if (c === '.' && src.startsWith('...', i)) {
112+ const eol = src.indexOf('\n', i);
113+ const end = eol === -1 ? src.length : eol;
114+ push('...', null);
115+ push(src.slice(i + 3, end), 'com');
116+ i = end;
117+ continue;
118+ }
119+
120+ // Char literal (or transpose).
121+ if (c === "'") {
122+ if (quoteIsTranspose(src, i)) {
123+ push("'", null);
124+ i++;
125+ continue;
126+ }
127+ let j = i + 1;
128+ while (j < src.length && src[j] !== '\n') {
129+ if (src[j] === "'") {
130+ if (src[j + 1] === "'") j += 2; // escaped quote
131+ else {
132+ j++;
133+ break;
134+ }
135+ } else j++;
136+ }
137+ push(src.slice(i, j), 'str');
138+ i = j;
139+ continue;
140+ }
141+
142+ // Double-quoted string.
143+ if (c === '"') {
144+ let j = i + 1;
145+ while (j < src.length && src[j] !== '\n') {
146+ if (src[j] === '"') {
147+ if (src[j + 1] === '"') j += 2;
148+ else {
149+ j++;
150+ break;
151+ }
152+ } else j++;
153+ }
154+ push(src.slice(i, j), 'str');
155+ i = j;
156+ continue;
157+ }
158+
159+ // Number: 12, 1.5, .5, 1e-3, 2i
160+ if (isDigit(c) || (c === '.' && isDigit(src[i + 1]))) {
161+ let j = i;
162+ while (j < src.length && isDigit(src[j])) j++;
163+ if (src[j] === '.') {
164+ j++;
165+ while (j < src.length && isDigit(src[j])) j++;
166+ }
167+ if (src[j] === 'e' || src[j] === 'E') {
168+ let k = j + 1;
169+ if (src[k] === '+' || src[k] === '-') k++;
170+ if (isDigit(src[k])) {
171+ k++;
172+ while (k < src.length && isDigit(src[k])) k++;
173+ j = k;
174+ }
175+ }
176+ if (src[j] === 'i' || src[j] === 'j') j++;
177+ push(src.slice(i, j), 'num');
178+ i = j;
179+ continue;
180+ }
181+
182+ // Identifier / keyword / external operation.
183+ if (isIdentStart(c)) {
184+ let j = i;
185+ while (j < src.length && isIdent(src[j])) j++;
186+ const word = src.slice(i, j);
187+ push(word, KEYWORDS.has(word) ? 'kw' : external.has(word) ? 'ext' : null);
188+ i = j;
189+ continue;
190+ }
191+
192+ push(c, null);
193+ i++;
194+ }
195+
196+ return out;
197+}
198+
199+const escapeHtml = (s: string): string =>
200+ s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
201+
202+/** Highlighted HTML for `src`, safe to assign to innerHTML. */
203+export function highlightMatlab(
204+ src: string,
205+ external: ReadonlySet<string> = new Set(),
206+): string {
207+ const html = tokenizeMatlab(src, external)
208+ .map((t) => (t.cls ? `<span class="tok-${t.cls}">${escapeHtml(t.text)}</span>` : escapeHtml(t.text)))
209+ .join('');
210+ // A trailing newline keeps the last line's box height stable, so the overlay
211+ // and the textarea scroll to the same extent.
212+ return `${html}\n`;
213+}
src/main.tsmodified+164−74View file
@@ -1,18 +1,17 @@
1-import {
2- GpuBackend,
3- CpuBackend,
4- requestShtDevice,
5- describeAdapter,
6- type ShtBackend,
7-} from './solver/backend.ts';
8-import { Simulation, gridForLmax } from './solver/simulation.ts';
9-import { presets, type ModelSpec, type Params } from './solver/models.ts';
1+import { requestShtDevice, ShtPlan } from './sht/sht.ts';
2+import { describeAdapter } from './solver/backend.ts';
3+import { gridForLmax, makeRandn } from './solver/simulation.ts';
4+import { presets, type Params } from './solver/models.ts';
5+import { GpuModel } from './mgpu/model.ts';
6+import { mModelByKey, type MModel } from './mgpu/registry.ts';
7+import { ModelCompileError, formatFailure } from './mgpu/errors.ts';
8+import { EXTERNAL_OPS } from './mgpu/externals.ts';
9+import { CodeEditor } from './editor/codeEditor.ts';
1010 import {
1111 formatCommand,
1212 resolvePreset,
1313 DEFAULT_STEPS,
1414 DEFAULT_WARMUP,
15- type BackendKind,
1615 type RunSpec,
1716 } from './bench/runSpec.ts';
1817 import {
@@ -31,7 +30,6 @@ const $ = <T extends HTMLElement>(id: string): T =>
3130 const elModel = $<HTMLSelectElement>('model');
3231 const elLmax = $<HTMLSelectElement>('lmax');
3332 const elColormap = $<HTMLSelectElement>('colormap');
34-const elBackend = $<HTMLSelectElement>('backend');
3533 const elRunPause = $<HTMLButtonElement>('runpause');
3634 const elReseed = $<HTMLButtonElement>('reseed');
3735 const elResetView = $<HTMLButtonElement>('resetview');
@@ -42,6 +40,12 @@ const elCmd = $('cmd');
4240 const elCopyCmd = $<HTMLButtonElement>('copycmd');
4341 const elBlurb = $('blurb');
4442 const elErr = $('err');
43+const elSource = $<HTMLTextAreaElement>('source');
44+const elHighlight = $('highlight');
45+const elCompiled = $('compiled');
46+const elEditorTitle = $('editor-title');
47+const elRecompile = $<HTMLButtonElement>('recompile');
48+const elRevert = $<HTMLButtonElement>('revert');
4549
4650 for (const p of presets) {
4751 const o = document.createElement('option');
@@ -57,10 +61,26 @@ for (const name of colormapNames) {
5761 }
5862 elColormap.value = 'jet';
5963
64+/** The model source, with MATLAB highlighting. The host-provided operations are
65+ * marked so the boundary between the model and what it is given is visible. */
66+const editor = new CodeEditor({
67+ textarea: elSource,
68+ overlay: elHighlight,
69+ external: EXTERNAL_OPS,
70+ onInput: (value) => {
71+ editedSource = value;
72+ elRecompile.textContent = 'Recompile *';
73+ },
74+});
75+
76+/** Timesteps submitted per rendered frame. Nothing is read back between them,
77+ * so the batch costs one submit and one readback regardless of size. */
78+const STEPS_PER_FRAME = 4;
79+
6080 // ---------------------------------------------------------------- state
6181 let device: GPUDevice | null = null;
62-let backend: ShtBackend | null = null;
63-let sim: Simulation | null = null;
82+let sht: ShtPlan | null = null;
83+let gpu: GpuModel | null = null;
6484 let topo: SphereMeshTopology | null = null;
6585 let scenes: SphereScene[] = [];
6686 let colorbars: Colorbar[] = [];
@@ -70,15 +90,21 @@ let ranges: { lo: number; hi: number }[] = [];
7090 let resizeObs: ResizeObserver | null = null;
7191
7292 const initial = resolvePreset(presets[0].key);
73-let model: ModelSpec = initial.model;
93+let model: MModel = mModelByKey(initial.model.key)!;
7494 let params: Params = initial.params;
95+/** The .m as edited in the page; `null` while it matches the file. */
96+let editedSource: string | null = null;
7597 let seed = 1;
7698 let running = false;
7799 let adapterName = '';
78100 let pumping = false;
79101 let stepMs = 0;
102+let simTime = 0;
103+let stepCount = 0;
80104 let generation = 0; // bumped on every rebuild to cancel stale pumps
81105
106+const source = (): string => editedSource ?? model.source;
107+
82108 // ---------------------------------------------------------------- UI wiring
83109 function buildParamInputs(): void {
84110 elParams.replaceChildren();
@@ -94,6 +120,9 @@ function buildParamInputs(): void {
94120 input.addEventListener('change', () => {
95121 const v = Number(input.value);
96122 if (Number.isFinite(v)) params[spec.key] = v;
123+ // Parameters are uniforms, not constants baked into the kernels, so a
124+ // change costs an upload rather than a recompile.
125+ gpu?.setParams(params);
97126 updateCommand();
98127 });
99128 label.append(input);
@@ -103,8 +132,17 @@ function buildParamInputs(): void {
103132
104133 function applyPreset(presetKey: string): void {
105134 const resolved = resolvePreset(presetKey);
106- model = resolved.model;
135+ const next = mModelByKey(resolved.model.key);
136+ if (!next) {
137+ elErr.textContent = `No .m model for '${resolved.model.key}'`;
138+ return;
139+ }
140+ model = next;
107141 params = resolved.params;
142+ editedSource = null;
143+ editor.value = model.source;
144+ elEditorTitle.textContent =
145+ `models/${model.key}.m — init() and step(), compiled to WebGPU`;
108146 buildParamInputs();
109147 elBlurb.textContent = model.blurb;
110148 updateCommand();
@@ -115,7 +153,7 @@ function currentSpec(): RunSpec {
115153 return {
116154 preset: elModel.value,
117155 lmax: Number(elLmax.value),
118- backend: elBackend.value as BackendKind,
156+ backend: 'webgpu',
119157 seed,
120158 steps: DEFAULT_STEPS,
121159 warmup: DEFAULT_WARMUP,
@@ -132,8 +170,8 @@ elModel.addEventListener('change', () => {
132170 void rebuild();
133171 });
134172 elLmax.addEventListener('change', () => void rebuild());
135-elBackend.addEventListener('change', () => void rebuild());
136-elColormap.addEventListener('change', () => draw());
173+elColormap.addEventListener('change', () => void draw());
174+
137175 function setRunning(next: boolean): void {
138176 running = next;
139177 elRunPause.textContent = running ? 'Pause' : 'Run';
@@ -151,8 +189,18 @@ elResetView.addEventListener('click', () => {
151189 for (const s of scenes) s.resetCamera();
152190 });
153191
154-// The command reproduces this exact run on the desktop; keep it selectable
155-// even where the clipboard API is unavailable.
192+elRecompile.addEventListener('click', () => {
193+ editedSource = editor.value;
194+ void rebuild();
195+});
196+elRevert.addEventListener('click', () => {
197+ editedSource = null;
198+ editor.value = model.source;
199+ void rebuild();
200+});
201+
202+// The command reproduces this run's parameters on the desktop; keep it
203+// selectable even where the clipboard API is unavailable.
156204 elCopyCmd.addEventListener('click', () => {
157205 const text = elCmd.textContent ?? '';
158206 const flash = (msg: string): void => {
@@ -181,49 +229,82 @@ function disposeView(): void {
181229 elPanels.replaceChildren();
182230 }
183231
232+/** Seeded perturbation, one normal deviate per grid point. */
233+function makeNoise(npts: number): Float32Array {
234+ const randn = makeRandn(seed);
235+ const noise = new Float32Array(npts);
236+ for (let i = 0; i < npts; i++) noise[i] = model.seedAmp * randn();
237+ return noise;
238+}
239+
240+/** Report a compile failure, and select the offending text in the editor. */
241+function reportCompileError(e: unknown): void {
242+ elErr.textContent = formatFailure(e, source());
243+ elCompiled.textContent = '';
244+ if (e instanceof ModelCompileError && e.start !== undefined) {
245+ editor.select(e.start, e.end ?? e.start);
246+ }
247+}
248+
184249 async function rebuild(): Promise<void> {
185250 generation++;
186251 const gen = generation;
187- // a rebuild restarts from a fresh initial state, so pause like Re-seed does
188252 setRunning(false);
189253 disposeView();
190- backend?.destroy();
191- backend = null;
192- sim = null;
254+ gpu?.destroy();
255+ gpu = null;
256+ sht?.destroy();
257+ sht = null;
193258 stepMs = 0;
259+ simTime = 0;
260+ stepCount = 0;
261+ elErr.textContent = '';
194262 updateCommand();
263+ if (!device) return;
195264
196265 const lmax = Number(elLmax.value);
197266 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
198267 const cfg = { lmax, mmax: lmax, nlat, nphi };
199- const wantGpu = elBackend.value === 'webgpu' && device !== null;
268+
200269 try {
201- backend = wantGpu
202- ? await GpuBackend.create(device!, cfg)
203- : new CpuBackend(cfg);
270+ sht = await ShtPlan.create(device, cfg);
271+ gpu = await GpuModel.create({
272+ device,
273+ sht,
274+ cfg,
275+ source: source(),
276+ paramNames: model.params.map((p) => p.key),
277+ state: model.state,
278+ view: model.species,
279+ });
204280 } catch (e) {
205- elErr.textContent = `Failed to create transform plan: ${e}`;
281+ reportCompileError(e);
282+ gpu?.destroy();
283+ gpu = null;
284+ sht?.destroy();
285+ sht = null;
206286 return;
207287 }
208288 if (gen !== generation) return;
209- elErr.textContent =
210- !wantGpu && lmax > 31
211- ? 'Heads up: the CPU backend is a direct-summation f64 reference — expect well under 10 steps/s at this lmax.'
212- : '';
213289
214- sim = new Simulation(backend, model, params);
215- await sim.init(seed);
216- if (gen !== generation) return;
290+ gpu.setParams(params);
291+ gpu.init(makeNoise(nlat * nphi));
292+
293+ const plan = gpu.describe();
294+ elCompiled.textContent =
295+ `one step compiled to ${plan.step.length} GPU operations:\n` +
296+ plan.step.map((l) => ` ${l}`).join('\n');
297+ elRecompile.textContent = 'Recompile';
217298
218299 // mesh + scenes
219300 const phi = new Float64Array(nphi);
220301 for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
221- topo = buildTopology(backend.cosTheta, phi);
302+ topo = buildTopology(sht.cosTheta, phi);
222303
223304 const sphereBg = getComputedStyle(document.documentElement)
224305 .getPropertyValue('--sphere-bg')
225306 .trim();
226- for (let k = 0; k < sim.nspecies; k++) {
307+ for (let k = 0; k < model.species.length; k++) {
227308 const panel = document.createElement('div');
228309 panel.className = 'panel';
229310 const box = document.createElement('div');
@@ -262,29 +343,44 @@ async function rebuild(): Promise<void> {
262343 .querySelectorAll<HTMLElement>('.sphere-box')
263344 .forEach((box) => resizeObs!.observe(box));
264345
265- draw();
346+ await draw();
266347 updateStats();
267348 void pump();
268349 }
269350
270351 async function reseed(): Promise<void> {
271- if (!sim) return;
352+ if (!gpu || !sht) return;
272353 const gen = generation;
273- await sim.init(seed);
354+ gpu.init(makeNoise(sht.cfg.nlat * sht.cfg.nphi));
355+ simTime = 0;
356+ stepCount = 0;
274357 if (gen !== generation) return;
275358 for (const r of ranges) {
276359 r.lo = NaN;
277360 r.hi = NaN;
278361 }
279- draw();
362+ await draw();
363+ updateStats();
280364 }
281365
282366 // ---------------------------------------------------------------- drawing
283-function draw(): void {
284- if (!sim || !topo) return;
367+async function draw(): Promise<void> {
368+ if (!gpu || !topo) return;
369+ const gen = generation;
285370 const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
286- for (let k = 0; k < sim.nspecies; k++) {
287- fillFieldValues(valueBufs[k], sim.V[k], topo);
371+ for (let k = 0; k < model.species.length; k++) {
372+ // The one readback per frame — the loop is otherwise entirely on the GPU.
373+ // A rebuild can land while this is in flight and destroy the buffer being
374+ // mapped, which rejects the map; that result is stale anyway, so drop it.
375+ let field: Float32Array;
376+ try {
377+ field = await gpu.read(model.species[k]);
378+ } catch (e) {
379+ if (gen !== generation) return;
380+ throw e;
381+ }
382+ if (gen !== generation || !topo) return;
383+ fillFieldValues(valueBufs[k], field, topo);
288384 let lo = Infinity;
289385 let hi = -Infinity;
290386 for (const v of valueBufs[k]) {
@@ -314,17 +410,14 @@ function draw(): void {
314410 }
315411
316412 function updateStats(): void {
317- if (!sim || !backend) return;
318- const { nlat, nphi } = backend.cfg;
319- const kind =
320- backend.kind === 'webgpu'
321- ? `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`
322- : 'CPU f64 (direct summation)';
413+ if (!gpu || !sht) return;
414+ const { nlat, nphi } = sht.cfg;
415+ const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
323416 const rate = stepMs > 0 ? `${(1000 / stepMs).toFixed(1)} steps/s` : '—';
324417 elStats.innerHTML =
325- `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${backend.nlm.toLocaleString()} · ` +
418+ `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${sht.nlm.toLocaleString()} · ` +
326419 `${stepMs > 0 ? stepMs.toFixed(1) : '—'} ms/step · ${rate} · ` +
327- `t = <b>${sim.t.toFixed(2)}</b> (${sim.stepCount} steps)`;
420+ `t = <b>${simTime.toFixed(2)}</b> (${stepCount} steps)`;
328421 }
329422
330423 // ---------------------------------------------------------------- sim loop
@@ -334,24 +427,23 @@ async function pump(): Promise<void> {
334427 if (pumping) return;
335428 pumping = true;
336429 const gen = generation;
337- let lastYield = performance.now();
338430 try {
339- while (running && sim && gen === generation) {
431+ while (running && gpu && gen === generation) {
340432 const t0 = performance.now();
341- await sim.step();
342- const dtMs = performance.now() - t0;
433+ gpu.step(STEPS_PER_FRAME);
434+ // draw() awaits the readback, which also waits for the batch to finish,
435+ // so this measures the real end-to-end cost per step.
436+ await draw();
437+ if (gen !== generation) break;
438+ const dtMs = (performance.now() - t0) / STEPS_PER_FRAME;
343439 stepMs = stepMs === 0 ? dtMs : stepMs + 0.05 * (dtMs - stepMs);
344- const now = performance.now();
345- if (now - lastYield > 25 || backend?.kind === 'cpu') {
346- draw();
347- updateStats();
348- await nextFrame();
349- lastYield = performance.now();
350- }
440+ simTime += STEPS_PER_FRAME * (params.dt ?? 0);
441+ stepCount += STEPS_PER_FRAME;
442+ updateStats();
443+ await nextFrame();
351444 }
352- // final frame after pausing
353445 if (gen === generation) {
354- draw();
446+ await draw();
355447 updateStats();
356448 }
357449 } finally {
@@ -368,15 +460,13 @@ async function boot(): Promise<void> {
368460 adapterName = await describeAdapter(device);
369461 } catch (e) {
370462 device = null;
371- elLmax.value = '31';
372- elBackend.value = 'cpu';
373- elBackend.options[0].disabled = true;
374463 elErr.textContent =
375- `WebGPU is not available (${e instanceof Error ? e.message : e}); ` +
376- `falling back to the slow CPU transform at low resolution. ` +
377- `Use a WebGPU-capable browser (Chrome/Edge 113+) for the full experience.`;
464+ `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
465+ `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +
466+ `needs a WebGPU-capable browser (Chrome/Edge 113+).`;
467+ return;
378468 }
379- device?.lost.then((info) => {
469+ device.lost.then((info) => {
380470 if (info.reason !== 'destroyed') {
381471 elErr.textContent = `WebGPU device lost: ${info.message}`;
382472 }
src/mgpu/compile.tsadded+190−0View file
@@ -0,0 +1,190 @@
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+ */
21+import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
22+import { Workspace, Lowerer, tensorDouble, scalarDouble } from 'numbl-src/numbl-core/jit/index.ts';
23+import { specializeUserFunction } from 'numbl-src/numbl-core/jit/lowering/specialize.ts';
24+import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
25+import type { IRFunc, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
26+import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
27+import { externalOpFiles, type GridSizes } from './externals.ts';
28+import { ModelCompileError } from './errors.ts';
29+
30+/** What the host can supply for an argument the .m declares. */
31+export 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 };
41+
42+const 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+};
57+
58+/** One specialized function, as the planner consumes it. */
59+export 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+}
70+
71+/** The shape of a `function` statement in numbl's AST. */
72+interface FunctionDecl {
73+ type: 'Function';
74+ name: string;
75+ params: string[];
76+ outputs: string[];
77+}
78+
79+/**
80+ * A parsed model. Specialize the functions you need, then call `finish()` once
81+ * — the inline pass rewrites every specialization together.
82+ */
83+export class CompiledModel {
84+ #lowerer: Lowerer;
85+ #decls: Map<string, FunctionDecl>;
86+ #bindings: Record<string, Binding>;
87+
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();
100+
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+ }
111+
112+ /** Names of the functions the file defines. */
113+ functionNames(): string[] {
114+ return [...this.#decls.keys()];
115+ }
116+
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+ }
138+
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+ });
150+
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+ );
161+
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+ }
181+
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+ inlinePass({ topLevelStmts: [], functions: this.#lowerer.specializations });
189+ }
190+}
src/mgpu/errors.tsadded+105−0View file
@@ -0,0 +1,105 @@
1+/**
2+ * Compile failures, reported in coordinates of the model file the user edits.
3+ *
4+ * Failures arrive from three places, each with its own idea of position:
5+ * numbl's parser (a `position` offset), numbl's lowerer (`UnsupportedConstruct`
6+ * / `JitTypeError`, with a `span`), and this project's WGSL emitter
7+ * (`UnsupportedOnGpu`, carrying the numbl span it was given). All of them are
8+ * offsets into the whole model file — the file is parsed once, and each function
9+ * is specialized from that one AST — so they need only be turned into a line and
10+ * column for the editor.
11+ */
12+
13+/** A compile failure located in the full model source. */
14+export class ModelCompileError extends Error {
15+ /** Offset into the whole .m file, when the failure has a position. */
16+ readonly start?: number;
17+ readonly end?: number;
18+ /** Name of the model function being compiled. */
19+ readonly fn?: string;
20+
21+ constructor(
22+ message: string,
23+ opts: { start?: number; end?: number; fn?: string; cause?: unknown } = {},
24+ ) {
25+ super(message, { cause: opts.cause });
26+ this.name = 'ModelCompileError';
27+ this.start = opts.start;
28+ this.end = opts.end;
29+ this.fn = opts.fn;
30+ }
31+}
32+
33+/** Extract whatever position information an error carries. */
34+function positionOf(e: unknown): { start?: number; end?: number } {
35+ const span = (e as { span?: { start?: unknown; end?: unknown } }).span;
36+ if (span && typeof span.start === 'number') {
37+ return {
38+ start: span.start,
39+ end: typeof span.end === 'number' ? span.end : undefined,
40+ };
41+ }
42+ // numbl's parser SyntaxError reports a bare offset.
43+ const position = (e as { position?: unknown }).position;
44+ if (typeof position === 'number') return { start: position };
45+ return {};
46+}
47+
48+/** Normalize any thrown value into a located `ModelCompileError`. */
49+function asCompileError(e: unknown, fn?: string): ModelCompileError {
50+ if (e instanceof ModelCompileError) return e;
51+ const { start, end } = positionOf(e);
52+ const raw = e instanceof Error ? e.message : String(e);
53+ // numbl's parse errors read as bare token complaints out of context.
54+ const message =
55+ (e as Error)?.name === 'SyntaxError' ? `MATLAB syntax error: ${raw}` : raw;
56+ return new ModelCompileError(message, { fn, start, end, cause: e });
57+}
58+
59+/**
60+ * Run `fn`, locating any compile failure in the model file. Use for whole-file
61+ * phases (parsing) that belong to no single function.
62+ */
63+export function inModel<T>(fn: () => T): T {
64+ try {
65+ return fn();
66+ } catch (e) {
67+ throw asCompileError(e);
68+ }
69+}
70+
71+/** Run `fn`, attributing any compile failure to the model function `name`. */
72+export function inFunction<T>(name: string, fn: () => T): T {
73+ try {
74+ return fn();
75+ } catch (e) {
76+ throw asCompileError(e, name);
77+ }
78+}
79+
80+/** Async form of `inFunction`. */
81+export async function inFunctionAsync<T>(
82+ name: string,
83+ fn: () => Promise<T>,
84+): Promise<T> {
85+ try {
86+ return await fn();
87+ } catch (e) {
88+ throw asCompileError(e, name);
89+ }
90+}
91+
92+/** Render a failure for display: message, section, and 1-based line/column. */
93+export function formatFailure(e: unknown, source: string): string {
94+ const message = e instanceof Error ? e.message : String(e);
95+ if (!(e instanceof ModelCompileError)) return message;
96+ const where: string[] = [];
97+ if (e.start !== undefined && e.start <= source.length) {
98+ const before = source.slice(0, e.start);
99+ const line = before.split('\n').length;
100+ const column = e.start - before.lastIndexOf('\n');
101+ where.push(`line ${line}, column ${column}`);
102+ }
103+ if (e.fn) where.push(`in ${e.fn}()`);
104+ return where.length ? `${message} (${where.join(', ')})` : message;
105+}
src/mgpu/externals.tsadded+95−0View file
@@ -0,0 +1,95 @@
1+/**
2+ * The two spherical-harmonic transforms, as external operations the .m can
3+ * call: `synth` (spectral -> grid) and `analys` (grid -> spectral).
4+ *
5+ * numbl needs only their *type rule* in order to lower a call site. It gets
6+ * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
7+ * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
8+ * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
9+ * imports available, so `transfer` builds numbl `Type` objects as plain
10+ * literals, and the grid sizes are baked in by the generator below (a grid
11+ * change recompiles anyway).
12+ *
13+ * The `emit`/`cBody` exports exist only because the loader's contract requires
14+ * them; we never emit C. The actual implementation is supplied by the WGSL
15+ * backend, which turns each of these calls into an ShtPlan encode.
16+ *
17+ * Spectral fields are carried as REAL 2 x nlm arrays (row 0 real part, row 1
18+ * imaginary), matching the interleaved layout the GPU buffers already use.
19+ * The IMEX update is real-linear, so no complex arithmetic is needed.
20+ */
21+
22+export interface GridSizes {
23+ /** Grid points, nlat*nphi. Grid fields are npts x 1 column vectors. */
24+ npts: number;
25+ /** Spectral coefficients. Spectral fields are 2 x nlm. */
26+ nlm: number;
27+}
28+
29+const numericType = (rows: number, cols: number): string =>
30+ `{ kind: "Numeric", elem: "double", isComplex: false, ` +
31+ `dims: [${dim(rows)}, ${dim(cols)}], shape: [${rows}, ${cols}], sign: "unknown" }`;
32+
33+// numbl's tensorDouble() canonicalizes an extent of 1 to its shared DIM_ONE
34+// singleton; mirror that so types compare equal to host-built ones.
35+const dim = (n: number): string =>
36+ n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;
37+
38+/** Source for one transform's `.mtoc2.js`. */
39+function transformSource(
40+ name: string,
41+ inRows: number,
42+ inCols: number,
43+ outRows: number,
44+ outCols: number,
45+): string {
46+ return `
47+exports.name = ${JSON.stringify(name)};
48+
49+exports.transfer = function (argTypes, nargout) {
50+ if (argTypes.length !== 1) {
51+ throw new Error("${name} takes exactly one argument, got " + argTypes.length);
52+ }
53+ if (nargout > 1) {
54+ throw new Error("${name} returns one value, but " + nargout + " were requested");
55+ }
56+ var a = argTypes[0];
57+ if (!a || a.kind !== "Numeric" || a.isComplex) {
58+ throw new Error("${name} requires a real numeric array");
59+ }
60+ var s = a.shape;
61+ if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
62+ throw new Error(
63+ "${name} requires a ${inRows}x${inCols} array, got " +
64+ (s ? s.join("x") : "unknown shape")
65+ );
66+ }
67+ return [${numericType(outRows, outCols)}];
68+};
69+
70+// Never called: this project executes the IR on WebGPU and emits no C.
71+exports.emit = function () {
72+ throw new Error("${name}: no C backend (this transform runs on WebGPU)");
73+};
74+exports.cBody = function () {
75+ return "";
76+};
77+`;
78+}
79+
80+/** Workspace files that make `synth` / `analys` resolvable during lowering. */
81+export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
82+ return [
83+ {
84+ name: 'synth.mtoc2.js',
85+ source: transformSource('synth', 2, g.nlm, g.npts, 1),
86+ },
87+ {
88+ name: 'analys.mtoc2.js',
89+ source: transformSource('analys', g.npts, 1, 2, g.nlm),
90+ },
91+ ];
92+}
93+
94+/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
95+export const EXTERNAL_OPS = new Set(['synth', 'analys']);
src/mgpu/model.tsadded+217−0View file
@@ -0,0 +1,217 @@
1+/**
2+ * A .m model, compiled and running on the GPU.
3+ *
4+ * A model file is ordinary MATLAB: it defines an `init` function that builds the
5+ * initial spectral state and a `step` function that advances it one timestep.
6+ * Each is specialized for the current grid and compiled into a ModelPlan, and
7+ * both operate on the same state buffers (see HostBuffers).
8+ *
9+ * Both functions return the new state followed by the grid fields the app
10+ * renders, so their signatures say exactly what they produce:
11+ *
12+ * function [U, V, u, v] = init(noise, a, b)
13+ * function [U, V, u, v] = step(U, V, lam, a, b, D1, D2, dt)
14+ *
15+ * The host supplies the things that are precomputation rather than algorithm:
16+ * the grid, the Laplace-Beltrami eigenvalues, the seeded initial noise, and the
17+ * parameter values. Each argument is matched to the .m's declared parameter
18+ * name, so the file documents its own interface.
19+ */
20+import { ShtPlan } from '../sht/sht.ts';
21+import { lmIndex, type ShtConfig } from '../sht/layout.ts';
22+import { HostBuffers, ModelPlan } from './plan.ts';
23+import { inFunction, inFunctionAsync, inModel } from './errors.ts';
24+import { CompiledModel, type Binding } from './compile.ts';
25+
26+export interface ModelParams {
27+ [key: string]: number;
28+}
29+
30+export interface GpuModelOptions {
31+ device: GPUDevice;
32+ sht: ShtPlan;
33+ cfg: ShtConfig;
34+ /** Model source (.m text). */
35+ source: string;
36+ /** Parameter names the .m may take as arguments. */
37+ paramNames: string[];
38+ /** Spectral state names, in order (e.g. ['U', 'V']). */
39+ state: string[];
40+ /** Grid fields to render, in order (e.g. ['u', 'v']). */
41+ view: string[];
42+}
43+
44+/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
45+ * matches the 2 x nlm spectral layout element for element. */
46+export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
47+ const lam = new Float32Array(2 * nlm);
48+ for (let m = 0; m <= cfg.mmax; m++) {
49+ for (let l = m; l <= cfg.lmax; l++) {
50+ const i = lmIndex(cfg.lmax, l, m);
51+ lam[2 * i] = l * (l + 1);
52+ lam[2 * i + 1] = l * (l + 1);
53+ }
54+ }
55+ return lam;
56+}
57+
58+export class GpuModel {
59+ readonly paramNames: string[];
60+ readonly state: string[];
61+ readonly view: string[];
62+ readonly npts: number;
63+ readonly nlm: number;
64+
65+ #device: GPUDevice;
66+ #host: HostBuffers;
67+ #initPlan: ModelPlan;
68+ #stepPlan: ModelPlan;
69+ #readback: GPUBuffer;
70+ /** Which function wrote the state most recently; see `read`. */
71+ #lastRan: 'init' | 'step' = 'init';
72+
73+ private constructor(init: {
74+ device: GPUDevice;
75+ host: HostBuffers;
76+ initPlan: ModelPlan;
77+ stepPlan: ModelPlan;
78+ readback: GPUBuffer;
79+ paramNames: string[];
80+ state: string[];
81+ view: string[];
82+ npts: number;
83+ nlm: number;
84+ }) {
85+ this.#device = init.device;
86+ this.#host = init.host;
87+ this.#initPlan = init.initPlan;
88+ this.#stepPlan = init.stepPlan;
89+ this.#readback = init.readback;
90+ this.paramNames = init.paramNames;
91+ this.state = init.state;
92+ this.view = init.view;
93+ this.npts = init.npts;
94+ this.nlm = init.nlm;
95+ }
96+
97+ static async create(opts: GpuModelOptions): Promise<GpuModel> {
98+ const { device, sht, cfg, source, paramNames, state, view } = opts;
99+ const npts = cfg.nlat * cfg.nphi;
100+ const nlm = sht.nlm;
101+
102+ // What the .m may ask for by parameter name. Spectral state and the
103+ // eigenvalues are 2 x nlm; the seeded perturbation is a grid field.
104+ const bindings: Record<string, Binding> = {
105+ lam: { kind: 'tensor', shape: [2, nlm] },
106+ noise: { kind: 'tensor', shape: [npts, 1] },
107+ npts: { kind: 'const', value: npts },
108+ nlm: { kind: 'const', value: nlm },
109+ };
110+ for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
111+ for (const p of paramNames) bindings[p] = { kind: 'param' };
112+
113+ // Parsing belongs to the file, not to either function.
114+ const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
115+ // Both functions return the new state first, then the rendered grid fields.
116+ const nargout = state.length + view.length;
117+ const initFn = inFunction('init', () => compiled.specialize('init', nargout));
118+ const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
119+ compiled.finish();
120+
121+ // Only the state outputs feed back into the argument buffers; the grid
122+ // fields are read for display and then overwritten next call.
123+ const feedback = [...state, ...view.map(() => null)];
124+
125+ const host = new HostBuffers(device);
126+ // The host owns the state and the inputs it uploads, whether or not a given
127+ // function happens to take them as arguments — `init` does not read `U`, but
128+ // it writes it, and `step` reads it back.
129+ for (const s of state) host.ensure(s, 2 * nlm);
130+ host.ensure('lam', 2 * nlm);
131+ host.ensure('noise', npts);
132+
133+ const initPlan = await inFunctionAsync('init', () =>
134+ ModelPlan.create(device, sht, { fn: initFn, feedback }, host),
135+ );
136+ const stepPlan = await inFunctionAsync('step', () =>
137+ ModelPlan.create(device, sht, { fn: stepFn, feedback }, host),
138+ );
139+
140+ host.upload('lam', eigenvalues(cfg, nlm));
141+
142+ const readback = device.createBuffer({
143+ label: 'mgpu-readback',
144+ size: 4 * Math.max(npts, 2 * nlm),
145+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
146+ });
147+
148+ return new GpuModel({
149+ device, host, initPlan, stepPlan, readback,
150+ paramNames, state, view, npts, nlm,
151+ });
152+ }
153+
154+ setParams(params: ModelParams): void {
155+ this.#initPlan.setParams(params);
156+ this.#stepPlan.setParams(params);
157+ }
158+
159+ /** Upload the seeded perturbation and run `init`. */
160+ init(noise: Float32Array): void {
161+ this.#host.upload('noise', noise);
162+ const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
163+ this.#initPlan.encodeSteps(enc, 1);
164+ this.#device.queue.submit([enc.finish()]);
165+ this.#lastRan = 'init';
166+ }
167+
168+ /**
169+ * Advance `steps` timesteps. Synchronous — this only records commands and
170+ * submits them; nothing is read back and nothing is awaited.
171+ */
172+ step(steps = 1): void {
173+ const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
174+ this.#stepPlan.encodeSteps(enc, steps);
175+ this.#device.queue.submit([enc.finish()]);
176+ this.#lastRan = 'step';
177+ }
178+
179+ /**
180+ * Read a named value back to the CPU. The only await in the whole loop.
181+ *
182+ * Grid fields like `u` are produced by both functions, into separate buffers
183+ * (only the spectral state is shared), so this reads from whichever ran most
184+ * recently — which is what makes the first frame show the initial state
185+ * rather than an unwritten buffer.
186+ */
187+ async read(name: string): Promise<Float32Array> {
188+ const [first, second] =
189+ this.#lastRan === 'init'
190+ ? [this.#initPlan, this.#stepPlan]
191+ : [this.#stepPlan, this.#initPlan];
192+ const buffer = first.buffer(name) ?? second.buffer(name);
193+ const count = first.elementCount(name) ?? second.elementCount(name);
194+ if (!buffer || count === undefined) {
195+ throw new Error(`read: the model has no value named '${name}'`);
196+ }
197+ const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
198+ enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
199+ this.#device.queue.submit([enc.finish()]);
200+ await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
201+ const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
202+ this.#readback.unmap();
203+ return out;
204+ }
205+
206+ /** What the .m compiled to, for display. */
207+ describe(): { init: string[]; step: string[] } {
208+ return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
209+ }
210+
211+ destroy(): void {
212+ this.#initPlan.destroy();
213+ this.#stepPlan.destroy();
214+ this.#host.destroy();
215+ this.#readback.destroy();
216+ }
217+}
src/mgpu/numbl.d.tsadded+242−0View file
@@ -0,0 +1,242 @@
1+/**
2+ * The numbl compiler surface this project depends on.
3+ *
4+ * We reach past numbl's published entry points into its JIT internals (parser,
5+ * lowerer, IR, inline pass), which its package `exports` map does not expose.
6+ * Those imports resolve through the `numbl-src` alias in vite.config.ts; these
7+ * declarations are what TypeScript checks against.
8+ *
9+ * Declaring the surface here rather than type-checking numbl's sources
10+ * directly keeps this project's compiler settings independent of numbl's, and
11+ * pins the exact contract we rely on. If numbl changes one of these shapes,
12+ * the build breaks here with a clear diff rather than deep inside its tree.
13+ *
14+ * Only the nodes the WGSL backend actually walks are spelled out; every other
15+ * IR kind is collapsed into a catch-all so that unhandled constructs are
16+ * rejected with a message instead of being silently mis-compiled.
17+ */
18+
19+declare module 'numbl-src/numbl-core/jit/lowering/types.ts' {
20+ export type Sign =
21+ | 'positive' | 'nonneg' | 'negative' | 'nonpositive'
22+ | 'zero' | 'nonzero' | 'unknown';
23+
24+ export type DimInfo = { kind: 'exact'; value: number } | { kind: 'unknown' };
25+
26+ export type NumericExact =
27+ | number
28+ | Float64Array
29+ | { re: number; im: number }
30+ | { re: Float64Array; im: Float64Array };
31+
32+ export interface NumericType {
33+ kind: 'Numeric';
34+ elem: 'double' | 'logical' | 'char' | string;
35+ isComplex: boolean;
36+ dims: DimInfo[];
37+ /** Present iff every dim is exact. */
38+ shape?: number[];
39+ sign: Sign;
40+ exact?: NumericExact;
41+ }
42+
43+ /** Everything the WGSL backend rejects. */
44+ export interface NonNumericType {
45+ kind: 'Void' | 'Unknown' | 'String' | 'Handle' | 'Struct' | 'Class' | 'Cell';
46+ }
47+
48+ export type Type = NumericType | NonNumericType;
49+
50+ export function isMultiElement(t: NumericType): boolean;
51+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
52+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
53+}
54+
55+declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
56+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
57+
58+ export interface Span {
59+ file: string;
60+ start: number;
61+ end: number;
62+ }
63+
64+ export interface NumLit {
65+ kind: 'NumLit';
66+ value: number;
67+ ty: Type;
68+ span: Span;
69+ }
70+ export interface Var {
71+ kind: 'Var';
72+ name: string;
73+ cName: string;
74+ ty: Type;
75+ span: Span;
76+ }
77+ export interface Binary {
78+ kind: 'Binary';
79+ builtin: string;
80+ left: IRExpr;
81+ right: IRExpr;
82+ ty: Type;
83+ span: Span;
84+ }
85+ export interface Unary {
86+ kind: 'Unary';
87+ builtin: string;
88+ operand: IRExpr;
89+ ty: Type;
90+ span: Span;
91+ }
92+ export interface Call {
93+ kind: 'Call';
94+ cName: string;
95+ name: string;
96+ args: IRExpr[];
97+ ty: Type;
98+ span: Span;
99+ }
100+ /** Any other IR expression kind — rejected by the WGSL emitter. */
101+ export interface OtherExpr {
102+ kind:
103+ | 'ImagLit' | 'StringLit' | 'TensorBuild' | 'TensorConcat' | 'CellLit'
104+ | 'CellEmpty' | 'CellIndexLoad' | 'HandleLit' | 'HandleCaptureLoad'
105+ | 'StructLit' | 'MemberLoad' | 'IndexLoad' | 'IndexSlice' | 'EndRef'
106+ | 'MakeRange';
107+ ty: Type;
108+ span: Span;
109+ }
110+
111+ export type IRExpr = NumLit | Var | Binary | Unary | Call | OtherExpr;
112+
113+ export interface Assign {
114+ kind: 'Assign';
115+ name: string;
116+ cName: string;
117+ ty: Type;
118+ expr: IRExpr;
119+ span: Span;
120+ }
121+ /** Any other IR statement kind — rejected by the planner. */
122+ export interface OtherStmt {
123+ kind:
124+ | 'ExprStmt' | 'If' | 'While' | 'For' | 'ReturnFromFunction' | 'Break'
125+ | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
126+ | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
127+ span: Span;
128+ }
129+
130+ export type IRStmt = Assign | OtherStmt;
131+
132+ export interface IRFunc {
133+ name: string;
134+ cName: string;
135+ /** Parameter source names. */
136+ params: string[];
137+ /** Parameter cNames, parallel to `params`. */
138+ cParams: string[];
139+ paramTypes: Type[];
140+ /** Output source names. */
141+ outputs: string[];
142+ /** Output cNames, parallel to `outputs`. */
143+ cOutputs: string[];
144+ outputTypes: Type[];
145+ body: IRStmt[];
146+ span: Span;
147+ }
148+
149+ export interface IRProgram {
150+ topLevelStmts: IRStmt[];
151+ functions: Map<string, IRFunc>;
152+ }
153+}
154+
155+declare module 'numbl-src/numbl-core/parser/index.ts' {
156+ export interface AbstractSyntaxTree {
157+ body: unknown[];
158+ }
159+ export function parseMFile(input: string, fileName?: string): AbstractSyntaxTree;
160+ export class SyntaxError extends Error {}
161+}
162+
163+declare module 'numbl-src/numbl-core/jit/index.ts' {
164+ import type { AbstractSyntaxTree } from 'numbl-src/numbl-core/parser/index.ts';
165+ import type { IRProgram, IRFunc, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
166+ import type { Type, NumericType, Sign } from 'numbl-src/numbl-core/jit/lowering/types.ts';
167+
168+ export interface WorkspaceFile {
169+ name: string;
170+ source: string;
171+ ast?: AbstractSyntaxTree;
172+ }
173+
174+ export class Workspace {
175+ constructor(mainFile: string, searchPaths?: ReadonlyArray<string>);
176+ addFile(file: WorkspaceFile): void;
177+ finalize(): void;
178+ }
179+
180+ export interface EnvEntry {
181+ cName: string;
182+ ty: Type;
183+ maybeUnassigned?: boolean;
184+ }
185+
186+ export class Lowerer {
187+ constructor(workspace: Workspace);
188+ /** Pre-bindable variable scope: seed host-provided values here. */
189+ env: Map<string, EnvEntry>;
190+ specializations: Map<string, IRFunc>;
191+ lowerProgram(ast: AbstractSyntaxTree): IRProgram;
192+ }
193+
194+ /** Thrown for MATLAB the JIT pipeline cannot lower; carries a source span. */
195+ export class UnsupportedConstruct extends Error {
196+ span?: Span;
197+ }
198+ export class JitTypeError extends Error {
199+ span?: Span;
200+ }
201+
202+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
203+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
204+ export function isMultiElement(t: NumericType): boolean;
205+}
206+
207+declare module 'numbl-src/numbl-core/jit/lowering/specialize.ts' {
208+ import type { Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
209+ import type { IRFunc, IRExpr, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
210+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
211+
212+ /**
213+ * Lower one user function for a concrete argument-type signature. Called with
214+ * a `Lowerer` as `this` (numbl's own JIT does the same), so specializations
215+ * accumulate in `lowerer.specializations`.
216+ */
217+ export function specializeUserFunction(
218+ this: Lowerer,
219+ decl: unknown,
220+ argTypes: Type[],
221+ specSource?: string,
222+ definingFile?: string,
223+ preSeedOutput?: { name: string; ty: Type; initExpr: IRExpr },
224+ nargout?: number,
225+ callSiteSpan?: Span,
226+ ): IRFunc;
227+}
228+
229+declare module 'numbl-src/numbl-core/jit/codegen/inlinePass.ts' {
230+ import type { IRProgram } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
231+ /** Folds single-use ANF temps into their consumer, in place. */
232+ export function inlinePass(prog: IRProgram): void;
233+}
234+
235+declare module 'numbl-src/numbl-core/jit/builtins/index.ts' {
236+ export interface Builtin {
237+ name: string;
238+ /** Safe to evaluate one output element from one input element per slot. */
239+ elementwise?: boolean;
240+ }
241+ export function getBuiltin(name: string): Builtin | undefined;
242+}
src/mgpu/plan.tsadded+534−0View file
@@ -0,0 +1,534 @@
1+/**
2+ * Statement list -> a replayable sequence of GPU operations.
3+ *
4+ * Everything expensive happens once, here: pipeline compilation, buffer
5+ * allocation, bind-group construction. Because numbl fixes every type and
6+ * shape at lowering time, the resulting op sequence is fully static — so
7+ * `encodeStep` is pure synchronous command recording, with no allocation, no
8+ * pipeline lookup and no readback. That is what lets the whole timestep be
9+ * encoded into one submit and keeps the CPU out of the loop.
10+ */
11+import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
12+import type { Assign, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
13+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
14+import { ShtPlan, type ShtBinding } from '../sht/sht.ts';
15+import type { CompiledFunction } from './compile.ts';
16+import { EXTERNAL_OPS, type GridSizes } from './externals.ts';
17+import {
18+ buildKernel,
19+ UnsupportedOnGpu,
20+ WORKGROUP_SIZE,
21+ type KernelInputs,
22+} from './wgsl.ts';
23+
24+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
25+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
26+const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
27+
28+interface Slot {
29+ buffer: GPUBuffer;
30+ count: number;
31+}
32+
33+const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
34+ device.createBuffer({
35+ label,
36+ size: 4 * count,
37+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
38+ });
39+
40+/**
41+ * Buffers for host-bound variables, shared across plans.
42+ *
43+ * A model is two programs — `init` and `step` — compiled separately but
44+ * operating on the same state. `U` in the step must be the very buffer `init`
45+ * wrote, so the buffers for host bindings live here rather than inside either
46+ * plan.
47+ */
48+export class HostBuffers {
49+ #device: GPUDevice;
50+ #slots = new Map<string, Slot>();
51+
52+ constructor(device: GPUDevice) {
53+ this.#device = device;
54+ }
55+
56+ ensure(name: string, count: number): Slot {
57+ const existing = this.#slots.get(name);
58+ if (existing) {
59+ if (existing.count !== count) {
60+ throw new UnsupportedOnGpu(
61+ `'${name}' is ${existing.count} elements in one program and ` +
62+ `${count} in another`,
63+ );
64+ }
65+ return existing;
66+ }
67+ const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
68+ this.#slots.set(name, slot);
69+ return slot;
70+ }
71+
72+ get(name: string): Slot | undefined {
73+ return this.#slots.get(name);
74+ }
75+
76+ /** Upload initial data for a host binding. */
77+ upload(name: string, data: Float32Array): void {
78+ const slot = this.#slots.get(name);
79+ if (!slot) throw new Error(`upload: no buffer named '${name}'`);
80+ if (data.length !== slot.count) {
81+ throw new Error(
82+ `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
83+ );
84+ }
85+ this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
86+ }
87+
88+ destroy(): void {
89+ for (const s of this.#slots.values()) s.buffer.destroy();
90+ this.#slots.clear();
91+ }
92+}
93+
94+type Op =
95+ | {
96+ kind: 'kernel';
97+ pipeline: GPUComputePipeline;
98+ bindGroup: GPUBindGroup;
99+ count: number;
100+ label: string;
101+ /** Set when the kernel had to write to scratch because its output
102+ * aliases one of its inputs; copied back after the dispatch. */
103+ copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
104+ }
105+ | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
106+ | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
107+
108+export interface PlanSpec {
109+ /** The specialized function this plan executes. */
110+ fn: CompiledFunction;
111+ /** Output index -> host binding name to copy the result into after the run,
112+ * so the next call reads it (the new spectral state feeds the old). */
113+ feedback: (string | null)[];
114+}
115+
116+/**
117+ * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
118+ * buffers after it, then the params buffer.
119+ *
120+ * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
121+ * only contains the bindings the shader actually references — so a kernel that
122+ * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
123+ * binding and no longer match the bind group. An explicit layout may carry
124+ * bindings the shader ignores.
125+ */
126+function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
127+ const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
128+ binding,
129+ visibility: GPUShaderStage.COMPUTE,
130+ buffer: { type: 'read-only-storage' },
131+ });
132+ return device.createBindGroupLayout({
133+ entries: [
134+ {
135+ binding: 0,
136+ visibility: GPUShaderStage.COMPUTE,
137+ buffer: { type: 'storage' },
138+ },
139+ ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
140+ readOnly(inputs + 1),
141+ ],
142+ });
143+}
144+
145+async function makePipeline(
146+ device: GPUDevice,
147+ code: string,
148+ label: string,
149+ bindGroupLayout: GPUBindGroupLayout,
150+): Promise<GPUComputePipeline> {
151+ device.pushErrorScope('validation');
152+ const module = device.createShaderModule({ code, label });
153+ const info = await module.getCompilationInfo();
154+ const errors = info.messages.filter((m) => m.type === 'error');
155+ if (errors.length) {
156+ throw new UnsupportedOnGpu(
157+ `generated WGSL failed to compile for '${label}':\n` +
158+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
159+ `\n--- shader ---\n${code}`,
160+ );
161+ }
162+ const pipeline = await device.createComputePipelineAsync({
163+ layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
164+ compute: { module, entryPoint: 'main' },
165+ label,
166+ });
167+ const err = await device.popErrorScope();
168+ if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
169+ return pipeline;
170+}
171+
172+/** A compiled .m step, ready to run on the GPU. */
173+export class ModelPlan {
174+ /** Scalar parameter names, in the order the params buffer expects them. */
175+ readonly paramNames: string[];
176+
177+ #device: GPUDevice;
178+ #sht: ShtPlan;
179+ #ops: Op[];
180+ #owned: GPUBuffer[];
181+ #paramBuf: GPUBuffer;
182+ #paramData: Float32Array;
183+ /** Public name -> buffer, for uploading initial state and reading results. */
184+ #byName: Map<string, Slot>;
185+
186+ private constructor(init: {
187+ device: GPUDevice;
188+ sht: ShtPlan;
189+ ops: Op[];
190+ byName: Map<string, Slot>;
191+ owned: GPUBuffer[];
192+ paramBuf: GPUBuffer;
193+ paramData: Float32Array;
194+ paramNames: string[];
195+ }) {
196+ this.#device = init.device;
197+ this.#sht = init.sht;
198+ this.#ops = init.ops;
199+ this.#byName = init.byName;
200+ this.#owned = init.owned;
201+ this.#paramBuf = init.paramBuf;
202+ this.#paramData = init.paramData;
203+ this.paramNames = init.paramNames;
204+ }
205+
206+ static async create(
207+ device: GPUDevice,
208+ sht: ShtPlan,
209+ spec: PlanSpec,
210+ host: HostBuffers,
211+ ): Promise<ModelPlan> {
212+ const { fn } = spec;
213+
214+ const slots = new Map<string, Slot>();
215+ const byName = new Map<string, Slot>();
216+ const owned: GPUBuffer[] = [];
217+ /** Scalars the .m computes from its parameters, by cName. */
218+ const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
219+
220+ const alloc = (label: string, count: number): Slot => {
221+ const buffer = makeBuffer(device, label, count);
222+ owned.push(buffer);
223+ return { buffer, count };
224+ };
225+
226+ // Arguments, bound by what the function's signature declares. Array
227+ // arguments come from the shared pool, so a value one function returns is
228+ // the same buffer the next one reads. Scalar parameters share one small
229+ // storage buffer, in signature order.
230+ const paramNames: string[] = [];
231+ const paramSlots = new Map<string, number>();
232+ for (const p of fn.params) {
233+ if (p.binding.kind === 'tensor') {
234+ const count = p.binding.shape.reduce((x, y) => x * y, 1);
235+ const slot = host.ensure(p.name, count);
236+ slots.set(p.cName, slot);
237+ byName.set(p.name, slot);
238+ } else if (p.binding.kind === 'param') {
239+ paramSlots.set(p.cName, paramNames.length);
240+ paramNames.push(p.name);
241+ }
242+ // `const` arguments are exact in the IR and fold into the kernels.
243+ }
244+ const paramData = new Float32Array(Math.max(1, paramNames.length));
245+ const paramBuf = device.createBuffer({
246+ label: 'mgpu-params',
247+ size: 4 * paramData.length,
248+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
249+ });
250+
251+ const ops: Op[] = [];
252+ for (const stmt of fn.body) {
253+ await planStatement(stmt);
254+ }
255+
256+ // Feed declared outputs back into the argument buffers they replace.
257+ fn.outputs.forEach((out, i) => {
258+ const to = spec.feedback[i];
259+ if (!to) return;
260+ const src = slots.get(out.cName);
261+ const dst = host.get(to);
262+ if (!src) {
263+ throw new UnsupportedOnGpu(
264+ `'${fn.name}' declares the output '${out.name}' but never assigns it`,
265+ );
266+ }
267+ if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
268+ if (src.count !== dst.count) {
269+ throw new UnsupportedOnGpu(
270+ `'${out.name}' (${src.count} elements) cannot feed ` +
271+ `'${to}' (${dst.count})`,
272+ );
273+ }
274+ ops.push({
275+ kind: 'copy',
276+ from: src.buffer,
277+ to: dst.buffer,
278+ bytes: 4 * src.count,
279+ label: `${out.name} -> ${to}`,
280+ });
281+ });
282+
283+ return new ModelPlan({
284+ device, sht, ops, byName, owned, paramBuf, paramData, paramNames,
285+ });
286+
287+ async function planStatement(stmt: IRStmt): Promise<void> {
288+ if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
289+ if (stmt.kind !== 'Assign') {
290+ throw new UnsupportedOnGpu(
291+ `a model function body may only contain assignments ` +
292+ `(found '${stmt.kind}')`,
293+ stmt.span,
294+ );
295+ }
296+ if (!isNumeric(stmt.ty)) {
297+ throw new UnsupportedOnGpu(
298+ `'${stmt.name}' is not a numeric value`,
299+ stmt.span,
300+ );
301+ }
302+ if (!isTensor(stmt.ty)) {
303+ // A scalar the model derives from its parameters (`us = a + b`). It
304+ // gets no buffer and no dispatch: the kernels that read it bind it as
305+ // a `let` in their prologue.
306+ derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
307+ return;
308+ }
309+ const count = numel(stmt.ty);
310+
311+ // Reuse the destination buffer across steps: the same cName always maps
312+ // to the same buffer, so a step allocates nothing.
313+ let dest = slots.get(stmt.cName);
314+ if (!dest) {
315+ dest = alloc(`mgpu-${stmt.name}`, count);
316+ slots.set(stmt.cName, dest);
317+ } else if (dest.count !== count) {
318+ throw new UnsupportedOnGpu(
319+ `'${stmt.name}' changes size between assignments`,
320+ stmt.span,
321+ );
322+ }
323+ byName.set(stmt.name, dest);
324+
325+ const ext = externalCall(stmt);
326+ if (ext) {
327+ const argSlot = slots.get(ext.argCName);
328+ if (!argSlot) {
329+ throw new UnsupportedOnGpu(
330+ `'${ext.name}' reads '${ext.argName}', which has no buffer`,
331+ stmt.span,
332+ );
333+ }
334+ ops.push(
335+ ext.name === 'synth'
336+ ? {
337+ kind: 'synth',
338+ binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
339+ label: `${stmt.name} = synth(${ext.argName})`,
340+ }
341+ : {
342+ kind: 'analys',
343+ binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
344+ label: `${stmt.name} = analys(${ext.argName})`,
345+ },
346+ );
347+ return;
348+ }
349+
350+ // Element-wise kernel. Collect the distinct tensor operands and give
351+ // them dense binding slots.
352+ const tensors = new Map<string, number>();
353+ collectTensorVars(stmt.expr, (cName) => {
354+ if (!tensors.has(cName)) tensors.set(cName, tensors.size);
355+ });
356+
357+ const label = `${stmt.name} = <${count} elements, element-wise>`;
358+ const kernel = buildKernel(
359+ stmt,
360+ {
361+ tensors,
362+ params: paramSlots,
363+ scalars: derivedScalars,
364+ } satisfies KernelInputs,
365+ count,
366+ label,
367+ );
368+ const bindGroupLayout = kernelLayout(device, tensors.size);
369+ const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
370+
371+ // WebGPU forbids aliasing a writable storage binding with another
372+ // binding in the same group, so an in-place update (`u = u + 1`) writes
373+ // to scratch and copies back. Element-wise kernels only ever touch
374+ // their own index, so the copy is the only cost.
375+ const aliased = tensors.has(stmt.cName);
376+ const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
377+
378+ const entries: GPUBindGroupEntry[] = [
379+ { binding: 0, resource: { buffer: target.buffer } },
380+ ];
381+ for (const [cName, i] of tensors) {
382+ const s = slots.get(cName);
383+ if (!s) {
384+ throw new UnsupportedOnGpu(
385+ `'${stmt.name}' reads a value with no buffer`,
386+ stmt.span,
387+ );
388+ }
389+ entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
390+ }
391+ entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
392+
393+ ops.push({
394+ kind: 'kernel',
395+ pipeline,
396+ bindGroup: device.createBindGroup({
397+ layout: bindGroupLayout,
398+ entries,
399+ }),
400+ count,
401+ label,
402+ copyBack: aliased
403+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
404+ : undefined,
405+ });
406+ }
407+ }
408+
409+ /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
410+ setParams(values: Record<string, number>): void {
411+ this.paramNames.forEach((name, i) => {
412+ const v = values[name];
413+ this.#paramData[i] = Number.isFinite(v) ? v : 0;
414+ });
415+ this.#device.queue.writeBuffer(
416+ this.#paramBuf,
417+ 0,
418+ this.#paramData as Float32Array<ArrayBuffer>,
419+ );
420+ }
421+
422+ /** Buffer holding the named value, or undefined if the .m never binds it. */
423+ buffer(name: string): GPUBuffer | undefined {
424+ return this.#byName.get(name)?.buffer;
425+ }
426+
427+ elementCount(name: string): number | undefined {
428+ return this.#byName.get(name)?.count;
429+ }
430+
431+ /**
432+ * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
433+ * ops share one compute pass, which WebGPU executes in submission order
434+ * with a barrier between dispatches.
435+ */
436+ encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
437+ for (let s = 0; s < steps; s++) {
438+ let pass: GPUComputePassEncoder | null = null;
439+ const inPass = (): GPUComputePassEncoder => {
440+ if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
441+ return pass;
442+ };
443+ const endPass = (): void => {
444+ if (pass) {
445+ pass.end();
446+ pass = null;
447+ }
448+ };
449+ for (const op of this.#ops) {
450+ switch (op.kind) {
451+ case 'kernel': {
452+ const p = inPass();
453+ p.setPipeline(op.pipeline);
454+ p.setBindGroup(0, op.bindGroup);
455+ p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
456+ if (op.copyBack) {
457+ endPass();
458+ encoder.copyBufferToBuffer(
459+ op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
460+ );
461+ }
462+ break;
463+ }
464+ case 'synth':
465+ this.#shtInto(inPass(), op);
466+ break;
467+ case 'analys':
468+ this.#shtInto(inPass(), op);
469+ break;
470+ case 'copy':
471+ endPass();
472+ encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
473+ break;
474+ }
475+ }
476+ endPass();
477+ }
478+ }
479+
480+ #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
481+ if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
482+ else this.#sht.encodeAnalysInto(pass, op.binding);
483+ }
484+
485+ /** Human-readable op sequence — what the .m actually compiled to. */
486+ describe(): string[] {
487+ return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
488+ }
489+
490+ destroy(): void {
491+ for (const b of this.#owned) b.destroy();
492+ this.#paramBuf.destroy();
493+ this.#owned.length = 0;
494+ }
495+}
496+
497+/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
498+function externalCall(
499+ stmt: Assign,
500+): { name: string; argCName: string; argName: string } | null {
501+ const e = stmt.expr;
502+ if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
503+ if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
504+ throw new UnsupportedOnGpu(
505+ `'${e.name}' must be applied to a single variable`,
506+ stmt.span,
507+ );
508+ }
509+ const arg = e.args[0];
510+ return { name: e.name, argCName: arg.cName, argName: arg.name };
511+}
512+
513+function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
514+ const walk = (x: IRExpr): void => {
515+ switch (x.kind) {
516+ case 'Var':
517+ if (isTensor(x.ty)) visit(x.cName);
518+ return;
519+ case 'Binary':
520+ walk(x.left);
521+ walk(x.right);
522+ return;
523+ case 'Unary':
524+ walk(x.operand);
525+ return;
526+ case 'Call':
527+ x.args.forEach(walk);
528+ return;
529+ default:
530+ return;
531+ }
532+ };
533+ walk(e);
534+}
src/mgpu/registry.tsadded+59−0View file
@@ -0,0 +1,59 @@
1+/**
2+ * The available .m models.
3+ *
4+ * Parameter metadata (names, defaults, slider ranges), species names and the
5+ * dealiasing degree stay in src/solver/models.ts: the host owns those, and
6+ * sharing them means the reference solver used by the tests and the .m running
7+ * on the GPU are always configured identically. This module only attaches each
8+ * model's MATLAB source.
9+ *
10+ * Naming convention, relied on by the app and documented in each .m:
11+ * `u`, `v`, ... grid fields the model computes and the app renders
12+ * `U`, `V`, ... the corresponding spectral state (uppercase)
13+ */
14+import { models, type ModelSpec, type ParamSpec } from '../solver/models.ts';
15+import schnakenbergSource from '../../models/schnakenberg.m?raw';
16+import brusselatorSource from '../../models/brusselator.m?raw';
17+import allencahnSource from '../../models/allencahn.m?raw';
18+
19+const sources: Record<string, string> = {
20+ schnakenberg: schnakenbergSource,
21+ brusselator: brusselatorSource,
22+ allencahn: allencahnSource,
23+};
24+
25+export interface MModel {
26+ key: string;
27+ label: string;
28+ blurb: string;
29+ /** Grid fields to render, one panel each. */
30+ species: string[];
31+ /** Spectral state names the .m advances. */
32+ state: string[];
33+ params: ParamSpec[];
34+ /** Polynomial degree of the reaction, for grid dealiasing. */
35+ pdeg: number;
36+ /** Amplitude of the seeded perturbation. */
37+ seedAmp: number;
38+ /** MATLAB source — the algorithm itself. */
39+ source: string;
40+}
41+
42+const fromSpec = (m: ModelSpec): MModel => ({
43+ key: m.key,
44+ label: m.label,
45+ blurb: m.blurb,
46+ species: m.species,
47+ state: m.species.map((s) => s.toUpperCase()),
48+ params: m.params,
49+ pdeg: m.pdeg,
50+ seedAmp: m.seedAmp,
51+ source: sources[m.key],
52+});
53+
54+export const mModels: MModel[] = models
55+ .filter((m) => sources[m.key] !== undefined)
56+ .map(fromSpec);
57+
58+export const mModelByKey = (key: string): MModel | undefined =>
59+ mModels.find((m) => m.key === key);
src/mgpu/wgsl.tsadded+364−0View file
@@ -0,0 +1,364 @@
1+/**
2+ * IR expression tree -> one WGSL compute kernel.
3+ *
4+ * This is the WebGPU counterpart of numbl's C-side fused emitter
5+ * (`codegen/emitTensorFused.ts`): for an `Assign` whose right-hand side is
6+ * purely element-wise over operands of the target's shape, emit a single
7+ * kernel that computes one output element per invocation. Because numbl's
8+ * inline pass has already folded the ANF temps back together, one source line
9+ * of MATLAB becomes one kernel.
10+ *
11+ * Everything is f32, matching the existing fp32 WebGPU transform backend.
12+ */
13+import { getBuiltin } from 'numbl-src/numbl-core/jit/builtins/index.ts';
14+import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
15+import type { IRExpr, Assign } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
16+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
17+
18+/** Raised for a construct the WGSL backend cannot express. Mirrors numbl's
19+ * own decline discipline: fail at compile time with a source span, never
20+ * silently produce something that computes the wrong thing. */
21+export class UnsupportedOnGpu extends Error {
22+ readonly span?: unknown;
23+ constructor(message: string, span?: unknown) {
24+ super(message);
25+ this.name = 'UnsupportedOnGpu';
26+ this.span = span;
27+ }
28+}
29+
30+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
31+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
32+
33+/** Element-wise binary builtins -> WGSL infix operator. */
34+const BINARY_OPS: Record<string, string> = {
35+ plus: '+',
36+ minus: '-',
37+ times: '*',
38+ rdivide: '/',
39+ // Degenerate to element-wise when at least one side is a scalar; the
40+ // both-tensor (true matrix) case is rejected below.
41+ mtimes: '*',
42+ mrdivide: '/',
43+};
44+
45+/** Element-wise unary builtins -> WGSL prefix operator. */
46+const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };
47+
48+/** Element-wise builtin calls -> WGSL builtin of the same arity. */
49+const CALL_FNS: Record<string, string> = {
50+ abs: 'abs',
51+ acos: 'acos',
52+ asin: 'asin',
53+ atan: 'atan',
54+ atan2: 'atan2',
55+ ceil: 'ceil',
56+ cos: 'cos',
57+ cosh: 'cosh',
58+ exp: 'exp',
59+ floor: 'floor',
60+ log: 'log',
61+ log2: 'log2',
62+ max: 'max',
63+ min: 'min',
64+ round: 'round',
65+ sign: 'sign',
66+ sin: 'sin',
67+ sinh: 'sinh',
68+ sqrt: 'sqrt',
69+ tan: 'tan',
70+ tanh: 'tanh',
71+};
72+
73+/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
74+ * infers AbstractInt and rejects the mixed-type arithmetic. */
75+function f32Lit(v: number): string {
76+ if (!Number.isFinite(v)) {
77+ throw new UnsupportedOnGpu(`cannot emit non-finite literal ${v}`);
78+ }
79+ return Number.isInteger(v) && Math.abs(v) < 1e21
80+ ? `${v}.0`
81+ : String(v).includes('e')
82+ ? `${v}f`
83+ : String(v);
84+}
85+
86+/** How a scalar or tensor operand is read inside the kernel. */
87+export interface KernelInputs {
88+ /** cName -> storage binding index, for multi-element tensor operands. */
89+ tensors: Map<string, number>;
90+ /** cName -> slot in the params storage buffer, for runtime scalars. */
91+ params: Map<string, number>;
92+ /** cName -> defining expression, for scalars the .m computes from
93+ * parameters (`us = a + b`). These have no buffer and no param slot; they
94+ * become `let` bindings in the prologue of every kernel that reads them. */
95+ scalars: Map<string, { name: string; expr: IRExpr }>;
96+}
97+
98+/** Mutable state while emitting one kernel. */
99+interface Ctx {
100+ io: KernelInputs;
101+ /** `let` lines to emit before the body, in dependency order. */
102+ prologue: string[];
103+ /** cName -> WGSL identifier, for scalars already bound in the prologue. */
104+ bound: Map<string, string>;
105+}
106+
107+/** WGSL identifier for a derived scalar. Avoids a leading underscore, which
108+ * WGSL reserves. */
109+const scalarIdent = (cName: string): string =>
110+ `s_${cName.replace(/[^A-Za-z0-9_]/g, '_')}`;
111+
112+/**
113+ * Bind a .m-derived scalar in the prologue (once), after whatever it depends
114+ * on, and return its identifier.
115+ */
116+function bindScalar(cName: string, ctx: Ctx): string {
117+ const already = ctx.bound.get(cName);
118+ if (already) return already;
119+ const def = ctx.io.scalars.get(cName)!;
120+ const ident = scalarIdent(cName);
121+ // Claim the name before emitting the RHS so a (malformed) self-reference
122+ // cannot recurse forever.
123+ ctx.bound.set(cName, ident);
124+ const rhs = emitExpr(def.expr, ctx);
125+ ctx.prologue.push(` let ${ident} = ${rhs};`);
126+ return ident;
127+}
128+
129+/**
130+ * Emit the per-element WGSL expression for `e`. `i` is the element index
131+ * variable in scope.
132+ */
133+function emitExpr(e: IRExpr, ctx: Ctx): string {
134+ const io = ctx.io;
135+ switch (e.kind) {
136+ case 'NumLit':
137+ return f32Lit(e.value);
138+
139+ case 'Var': {
140+ if (isTensor(e.ty)) {
141+ const slot = io.tensors.get(e.cName);
142+ if (slot === undefined) {
143+ throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
144+ }
145+ return `in${slot}[i]`;
146+ }
147+ // Scalar: either an exact compile-time value or a runtime parameter.
148+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
149+ return f32Lit(e.ty.exact);
150+ }
151+ const slot = io.params.get(e.cName);
152+ if (slot !== undefined) return `prm[${slot}]`;
153+ if (io.scalars.has(e.cName)) return bindScalar(e.cName, ctx);
154+ throw new UnsupportedOnGpu(
155+ `scalar '${e.name}' is not a constant, a parameter, or computed in ` +
156+ `this model`,
157+ e.span,
158+ );
159+ }
160+
161+ case 'Binary': {
162+ if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
163+ isTensor(e.left.ty) && isTensor(e.right.ty)) {
164+ throw new UnsupportedOnGpu(
165+ `matrix '${e.builtin === 'mtimes' ? '*' : '/'}' is not supported; ` +
166+ `use the element-wise form ('.${e.builtin === 'mtimes' ? '*' : '/'}')`,
167+ e.span,
168+ );
169+ }
170+ if (e.builtin === 'power' || e.builtin === 'mpower') {
171+ return emitPower(e.left, e.right, ctx, e.span);
172+ }
173+ const op = BINARY_OPS[e.builtin];
174+ if (!op) {
175+ throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
176+ }
177+ return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
178+ }
179+
180+ case 'Unary': {
181+ const op = UNARY_OPS[e.builtin];
182+ if (!op) {
183+ throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
184+ }
185+ return `(${op}${emitExpr(e.operand, ctx)})`;
186+ }
187+
188+ case 'Call': {
189+ // A shape constructor used inside an element-wise expression
190+ // contributes the same constant at every slot, so it needs no buffer.
191+ // (The shape itself is validated against the target by checkShapes.)
192+ if (e.name === 'ones') return '1.0';
193+ if (e.name === 'zeros') return '0.0';
194+
195+ const fn = CALL_FNS[e.name];
196+ const b = getBuiltin(e.name);
197+ if (!fn || !b?.elementwise) {
198+ throw new UnsupportedOnGpu(
199+ `'${e.name}' cannot be evaluated element-wise on the GPU`,
200+ e.span,
201+ );
202+ }
203+ return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
204+ }
205+
206+ default:
207+ throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
208+ }
209+}
210+
211+/**
212+ * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
213+ * negative routinely, so expand small non-negative integer exponents into
214+ * repeated multiplication — which is also what makes `u.^2` free.
215+ */
216+function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
217+ const k =
218+ exponent.kind === 'NumLit'
219+ ? exponent.value
220+ : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
221+ ? exponent.ty.exact
222+ : undefined;
223+ const b = emitExpr(base, ctx);
224+ if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
225+ if (k === 0) return '1.0';
226+ // bind once so a compound base expression is not re-evaluated k times
227+ return `pow_i${k}(${b})`;
228+ }
229+ if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
230+ return `(1.0 / pow_i${-k}(${b}))`;
231+ }
232+ throw new UnsupportedOnGpu(
233+ `'.^' needs a literal integer exponent in [-8, 8] (got ` +
234+ `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
235+ `WGSL's pow() undefined`,
236+ span,
237+ );
238+}
239+
240+/** Fixed-exponent power helpers, emitted only when used. */
241+function powHelpers(used: Set<number>): string {
242+ const out: string[] = [];
243+ for (const k of [...used].sort((a, b) => a - b)) {
244+ const body =
245+ k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
246+ out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
247+ }
248+ return out.join('\n');
249+}
250+
251+/**
252+ * Reject implicit expansion (broadcasting).
253+ *
254+ * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
255+ * expansion semantics — but a kernel that walks one linear index across every
256+ * operand would quietly compute the wrong thing. So every multi-element
257+ * operand must have exactly the target's shape. Scalars are fine: they are
258+ * read from the params buffer or folded in as literals.
259+ */
260+function checkShapes(e: IRExpr, target: NumericType, name: string): void {
261+ const want = target.shape;
262+ const same = (t: NumericType): boolean => {
263+ const got = t.shape;
264+ return (
265+ !!want && !!got && want.length === got.length &&
266+ want.every((d, i) => d === got[i])
267+ );
268+ };
269+ const walk = (x: IRExpr): void => {
270+ if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
271+ const got = x.ty.shape?.join('x') ?? 'dynamic';
272+ throw new UnsupportedOnGpu(
273+ `'${name}' would need implicit expansion: an operand is ${got} but the ` +
274+ `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
275+ `(the GPU kernel walks one index across every operand).`,
276+ x.span,
277+ );
278+ }
279+ switch (x.kind) {
280+ case 'Binary':
281+ walk(x.left);
282+ walk(x.right);
283+ return;
284+ case 'Unary':
285+ walk(x.operand);
286+ return;
287+ case 'Call':
288+ // A shape constructor's own arguments are sizes, not data.
289+ if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
290+ return;
291+ default:
292+ return;
293+ }
294+ };
295+ walk(e);
296+}
297+
298+export const WORKGROUP_SIZE = 64;
299+
300+export interface Kernel {
301+ code: string;
302+ /** Number of output elements. */
303+ count: number;
304+ label: string;
305+}
306+
307+/**
308+ * Build the kernel for one element-wise `Assign`. `io` must already map every
309+ * tensor operand cName to a binding index and every runtime scalar to a
310+ * params slot; the output is binding 0 and the params buffer is the binding
311+ * after the last input.
312+ */
313+export function buildKernel(
314+ stmt: Assign,
315+ io: KernelInputs,
316+ count: number,
317+ label: string,
318+): Kernel {
319+ if (!isNumeric(stmt.ty)) {
320+ throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
321+ }
322+ if (stmt.ty.isComplex) {
323+ throw new UnsupportedOnGpu(
324+ `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
325+ `field is carried as a real 2 x nlm array)`,
326+ stmt.span,
327+ );
328+ }
329+
330+ checkShapes(stmt.expr, stmt.ty, stmt.name);
331+ const ctx: Ctx = { io, prologue: [], bound: new Map() };
332+ const body = emitExpr(stmt.expr, ctx);
333+
334+ // pow_iK helpers are discovered during emission; scan the result for them.
335+ const used = new Set<number>();
336+ const emitted = [...ctx.prologue, body].join('\n');
337+ for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
338+
339+ const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
340+ for (const [, slot] of io.tensors) {
341+ decls.push(
342+ `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
343+ );
344+ }
345+ // Params live in a read-only storage buffer rather than a uniform block:
346+ // uniform arrays would need 16-byte element stride.
347+ const prmBinding = io.tensors.size + 1;
348+ decls.push(
349+ `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
350+ );
351+
352+ const code = `${decls.join('\n')}
353+
354+${powHelpers(used)}
355+
356+@compute @workgroup_size(${WORKGROUP_SIZE})
357+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
358+ let i = gid.x;
359+ if (i >= ${count}u) { return; }
360+${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
361+}
362+`;
363+ return { code, count, label };
364+}
src/raw.d.tsadded+5−0View file
@@ -0,0 +1,5 @@
1+/** Vite's `?raw` suffix imports a file's text. Used to load .m model sources. */
2+declare module '*?raw' {
3+ const source: string;
4+ export default source;
5+}
src/sht/sht.tsmodified+65−13View file
@@ -20,6 +20,15 @@ import {
2020
2121 export type FourierMode = 'auto' | 'fft' | 'dft';
2222
23+/** The two bind groups (Legendre stage, Fourier stage) of one transform. */
24+export interface ShtBinding {
25+ readonly bgLeg: GPUBindGroup;
26+ readonly bgFour: GPUBindGroup;
27+}
28+
29+const bgEntries = (bufs: GPUBuffer[]) =>
30+ bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
31+
2332 export interface ShtOptions {
2433 /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
2534 fourier?: FourierMode;
@@ -180,8 +189,7 @@ export class ShtPlan {
180189 this.pipeFourSynth = pFourS;
181190 this.pipeFourAnalys = pFourA;
182191
183- const entries = (bufs: GPUBuffer[]) =>
184- bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
192+ const entries = bgEntries;
185193 this.bgLegSynth = dev.createBindGroup({
186194 layout: pLegS.getBindGroupLayout(0),
187195 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
@@ -200,37 +208,81 @@ export class ShtPlan {
200208 });
201209 }
202210
203- /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
204- encodeSynth(encoder: GPUCommandEncoder): void {
211+ /**
212+ * Bind groups for one transform against caller-supplied spectral/spatial
213+ * buffers, so a transform can read and write buffers it does not own (the
214+ * .m-driven executor keeps a buffer per IR variable). Build these once at
215+ * plan time, not per step. `fmBuf` stays internal scratch: passes and
216+ * dispatches within a submission execute in order, so sequential transforms
217+ * can share it.
218+ */
219+ createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
220+ return {
221+ bgLeg: this.device.createBindGroup({
222+ layout: this.pipeLegSynth.getBindGroupLayout(0),
223+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
224+ }),
225+ bgFour: this.device.createBindGroup({
226+ layout: this.pipeFourSynth.getBindGroupLayout(0),
227+ entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
228+ }),
229+ };
230+ }
231+
232+ createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
233+ return {
234+ bgFour: this.device.createBindGroup({
235+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
236+ entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
237+ }),
238+ bgLeg: this.device.createBindGroup({
239+ layout: this.pipeLegAnalys.getBindGroupLayout(0),
240+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
241+ }),
242+ };
243+ }
244+
245+ /** Record synthesis into an existing compute pass. */
246+ encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
205247 const { mmax, nlat, nphi } = this.cfg;
206- const pass = encoder.beginComputePass({ label: 'sht-synth' });
207248 pass.setPipeline(this.pipeLegSynth);
208- pass.setBindGroup(0, this.bgLegSynth);
249+ pass.setBindGroup(0, b.bgLeg);
209250 pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);
210251 pass.setPipeline(this.pipeFourSynth);
211- pass.setBindGroup(0, this.bgFourSynth);
252+ pass.setBindGroup(0, b.bgFour);
212253 if (this.fourierMode === 'fft') {
213254 pass.dispatchWorkgroups(nlat);
214255 } else {
215256 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
216257 }
217- pass.end();
218258 }
219259
220- /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
221- encodeAnalys(encoder: GPUCommandEncoder): void {
260+ /** Record analysis into an existing compute pass. */
261+ encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
222262 const { mmax, nlat } = this.cfg;
223- const pass = encoder.beginComputePass({ label: 'sht-analys' });
224263 pass.setPipeline(this.pipeFourAnalys);
225- pass.setBindGroup(0, this.bgFourAnalys);
264+ pass.setBindGroup(0, b.bgFour);
226265 if (this.fourierMode === 'fft') {
227266 pass.dispatchWorkgroups(nlat);
228267 } else {
229268 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
230269 }
231270 pass.setPipeline(this.pipeLegAnalys);
232- pass.setBindGroup(0, this.bgLegAnalys);
271+ pass.setBindGroup(0, b.bgLeg);
233272 pass.dispatchWorkgroups(mmax + 1);
273+ }
274+
275+ /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
276+ encodeSynth(encoder: GPUCommandEncoder): void {
277+ const pass = encoder.beginComputePass({ label: 'sht-synth' });
278+ this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
279+ pass.end();
280+ }
281+
282+ /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
283+ encodeAnalys(encoder: GPUCommandEncoder): void {
284+ const pass = encoder.beginComputePass({ label: 'sht-analys' });
285+ this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
234286 pass.end();
235287 }
236288
src/solver/models.tsmodified+8−0View file
@@ -1,4 +1,12 @@
11 /**
2+ * Model metadata and the reference reaction terms.
3+ *
4+ * The parameter metadata here (names, labels, defaults, slider ranges) is what
5+ * the app uses; `src/mgpu/registry.ts` attaches each model's .m source to it, so
6+ * the .m running on the GPU and the reference solver used by the tests cannot be
7+ * configured differently. The `reaction` / `init` closures below are the
8+ * reference implementation only — the app runs the .m instead.
9+ *
210 * Reaction-diffusion model presets, ported from websph's
311 * SphericalReactionDiffusionDriver.m. Each species k solves
412 *
src/solver/simulation.tsmodified+7−0View file
@@ -1,4 +1,11 @@
11 /**
2+ * Reference solver — NOT what the app runs.
3+ *
4+ * The app executes the .m models under `models/` on the GPU (see `src/mgpu/`).
5+ * This TypeScript port remains as an independent implementation of the same
6+ * scheme, which is what makes it usable as the test oracle: `test/mgpuChecks.ts`
7+ * runs both from the same seed and compares. Keep the two in step.
8+ *
29 * IMEX Euler reaction-diffusion timestepper on the sphere, ported from
310 * websph's SphericalReactionDiffusion.m. Diffusion is implicit and diagonal
411 * in spherical-harmonic space (Laplace-Beltrami eigenvalues -l(l+1));
test/mgpuChecks.tsadded+202−0View file
@@ -0,0 +1,202 @@
1+/**
2+ * Correctness of the .m -> WGSL path, against the TypeScript solver.
3+ *
4+ * src/solver/ is no longer what the app runs, but it is an independent
5+ * implementation of the same IMEX scheme, which makes it the oracle here: run
6+ * both from the same seeded perturbation, through the same fp32 transforms, and
7+ * compare the spectral state.
8+ *
9+ * The only difference between the two is where the reaction and the IMEX update
10+ * happen — f64 on the CPU for the reference, fp32 in generated WGSL for the .m.
11+ * The pattern-forming regime amplifies small differences, so this checks a
12+ * short run.
13+ */
14+import { GpuBackend } from '../src/solver/backend.ts';
15+import { Simulation, gridForLmax, makeRandn } from '../src/solver/simulation.ts';
16+import { models, defaultParams } from '../src/solver/models.ts';
17+import { ShtPlan } from '../src/sht/sht.ts';
18+import { GpuModel } from '../src/mgpu/model.ts';
19+import { mModels, type MModel } from '../src/mgpu/registry.ts';
20+
21+type Check = (name: string, ok: boolean, detail: string) => void;
22+type Log = (s: string) => void;
23+
24+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
25+ let num = 0;
26+ let den = 0;
27+ for (let i = 0; i < a.length; i++) {
28+ const d = a[i] - b[i];
29+ num += d * d;
30+ den += b[i] * b[i];
31+ }
32+ return Math.sqrt(num / Math.max(den, 1e-300));
33+}
34+
35+const LMAX = 31;
36+const STEPS = 10;
37+
38+/**
39+ * Kernels the step of each model should compile to — one per element-wise line
40+ * of MATLAB. This is a fusion guard: numbl's lowering emits one statement per
41+ * *operator*, and the inline pass folds those back into per-line expression
42+ * trees. If that stops happening the results stay correct but every operator
43+ * becomes its own dispatch, which is exactly the silent regression to catch.
44+ */
45+const EXPECTED_KERNELS: Record<string, number> = {
46+ schnakenberg: 5,
47+ brusselator: 5,
48+ allencahn: 2,
49+};
50+
51+/** One model: compile it, run it, and compare against the reference solver. */
52+async function checkModel(
53+ device: GPUDevice,
54+ m: MModel,
55+ check: Check,
56+ log: Log,
57+): Promise<{ mgpuMs: number; refMs: number } | null> {
58+ const spec = models.find((x) => x.key === m.key);
59+ if (!spec) {
60+ check(`${m.key}: reference model exists`, false, 'no matching ModelSpec');
61+ return null;
62+ }
63+ const params = defaultParams(spec);
64+ const { nlat, nphi } = gridForLmax(LMAX, m.pdeg);
65+ const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
66+ const npts = nlat * nphi;
67+
68+ const sht = await ShtPlan.create(device, cfg);
69+ const gpu = await GpuModel.create({
70+ device,
71+ sht,
72+ cfg,
73+ source: m.source,
74+ paramNames: m.params.map((p) => p.key),
75+ state: m.state,
76+ view: m.species,
77+ });
78+ gpu.setParams(params);
79+
80+ const plan = gpu.describe();
81+ const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
82+ const xforms = plan.step.filter(
83+ (l) => l.startsWith('synth') || l.startsWith('analys'),
84+ ).length;
85+ log(
86+ ` ${m.key}.m -> ${plan.step.length} ops/step ` +
87+ `(${kernels} generated kernels, ${xforms} transforms)`,
88+ );
89+ const expected = EXPECTED_KERNELS[m.key];
90+ check(
91+ `${m.key}: element-wise lines fused into one kernel each`,
92+ kernels === expected,
93+ `${kernels} kernels (expected ${expected})`,
94+ );
95+
96+ // One randn per grid point, in index order. Rounded to f32 once and fed to
97+ // BOTH sides, so the comparison is about the compute path, not the seed.
98+ const randn = makeRandn(1);
99+ const noise = new Float32Array(npts);
100+ for (let i = 0; i < npts; i++) noise[i] = m.seedAmp * randn();
101+
102+ gpu.init(noise);
103+
104+ // Reference, seeded from the same perturbation by handing the model's own
105+ // init the identical sequence.
106+ const backend = await GpuBackend.create(device, cfg);
107+ const ref = new Simulation(backend, spec, params);
108+ {
109+ let i = 0;
110+ const feed = (): number => noise[i++] / m.seedAmp;
111+ const grids = m.state.map(() => new Float64Array(npts));
112+ spec.init(params, ref.x, ref.y, ref.z, feed, grids);
113+ for (let k = 0; k < m.state.length; k++) {
114+ ref.U[k].set(await backend.analys(grids[k]));
115+ }
116+ }
117+
118+ let worstInit = 0;
119+ for (let k = 0; k < m.state.length; k++) {
120+ worstInit = Math.max(worstInit, relL2(await gpu.read(m.state[k]), ref.U[k]));
121+ }
122+ check(
123+ `${m.key}: init matches reference`,
124+ worstInit < 1e-5,
125+ `rel L2 ${worstInit.toExponential(2)}`,
126+ );
127+
128+ for (let s = 0; s < STEPS; s++) await ref.step();
129+ gpu.step(STEPS);
130+
131+ let worst = 0;
132+ let nan = false;
133+ for (let k = 0; k < m.state.length; k++) {
134+ const got = await gpu.read(m.state[k]);
135+ worst = Math.max(worst, relL2(got, ref.U[k]));
136+ for (const v of got) if (!Number.isFinite(v)) nan = true;
137+ }
138+ check(
139+ `${m.key}: .m vs reference after ${STEPS} steps`,
140+ worst < 2e-3 && !nan,
141+ `rel L2 ${worst.toExponential(2)}${nan ? ', NaN!' : ''}`,
142+ );
143+
144+ // Guard against "both sides computed nothing".
145+ const field = await gpu.read(m.species[0]);
146+ let peak = 0;
147+ for (const v of field) peak = Math.max(peak, Math.abs(v));
148+ check(
149+ `${m.key}: rendered field is non-trivial`,
150+ peak > 1e-4,
151+ `max |${m.species[0]}| ${peak.toExponential(2)}`,
152+ );
153+
154+ // Step rate. The reference maps a staging buffer on every transform, so it
155+ // pays four driver round-trips per step; the .m path keeps everything in GPU
156+ // buffers and submits once.
157+ const TIMED = 50;
158+ const t0 = performance.now();
159+ gpu.step(TIMED);
160+ await device.queue.onSubmittedWorkDone();
161+ const mgpuMs = (performance.now() - t0) / TIMED;
162+
163+ const t1 = performance.now();
164+ for (let s = 0; s < TIMED; s++) await ref.step();
165+ const refMs = (performance.now() - t1) / TIMED;
166+
167+ gpu.destroy();
168+ backend.destroy();
169+ sht.destroy();
170+ return { mgpuMs, refMs };
171+}
172+
173+export async function mgpuChecks(
174+ device: GPUDevice,
175+ check: Check,
176+ log: Log,
177+): Promise<void> {
178+ check(
179+ 'models: registry populated',
180+ mModels.length === models.length,
181+ `${mModels.length} .m models`,
182+ );
183+
184+ for (const m of mModels) {
185+ const timing = await checkModel(device, m, check, log);
186+ if (!timing) continue;
187+ const { mgpuMs, refMs } = timing;
188+ log(
189+ ` step rate: .m ${mgpuMs.toFixed(2)} ms vs reference ` +
190+ `${refMs.toFixed(2)} ms (${(refMs / mgpuMs).toFixed(1)}x)`,
191+ );
192+ // A soft bound, not a performance target: on a software rasterizer the
193+ // transforms dominate and the round-trips this path avoids are a small
194+ // share of the total, so the ratio understates what it is worth on real
195+ // hardware. The check is only that executing the .m did not make it worse.
196+ check(
197+ `${m.key}: step rate no worse than the readback path`,
198+ mgpuMs < refMs * 1.15,
199+ `${mgpuMs.toFixed(2)} vs ${refMs.toFixed(2)} ms/step`,
200+ );
201+ }
202+}
test/test-page.tsmodified+4−0View file
@@ -7,6 +7,7 @@ import { GpuBackend, CpuBackend, requestShtDevice } from '../src/solver/backend.
77 import { Simulation, gridForLmax } from '../src/solver/simulation.ts';
88 import { models, defaultParams } from '../src/solver/models.ts';
99 import { randomSpectrum } from '../src/sht/reference.ts';
10+import { mgpuChecks } from './mgpuChecks.ts';
1011
1112 declare global {
1213 interface Window {
@@ -164,6 +165,9 @@ async function main(): Promise<void> {
164165 gpu.destroy();
165166 }
166167
168+ // --- the .m model compiled to WGSL, against the reference solver ---
169+ await mgpuChecks(device, check, log);
170+
167171 window.__RESULTS__ = { ok: failures === 0, lines };
168172 log(failures === 0 ? 'ALL PASS' : `${failures} FAILURE(S)`);
169173 }
vite.config.tsmodified+16−0View file
@@ -1,8 +1,24 @@
11 import { defineConfig } from 'vite';
22 import { resolve } from 'node:path';
33
4+// numbl is a local `file:` dependency, so node_modules/numbl is a symlink to
5+// the sibling checkout. Its package `exports` map only publishes the runtime
6+// entry points, not the compiler internals we need (parser + JIT lowering), so
7+// we reach them through a path alias. (package.json's `imports` field cannot
8+// express this — Node rejects node_modules targets — and plain Node could not
9+// resolve numbl's internal `.js`->`.ts` imports anyway, which is why the GPU
10+// tests run in the browser harness rather than under `node`.)
11+const numblSrc = resolve(import.meta.dirname, 'node_modules/numbl/src');
12+
413 export default defineConfig({
514 base: './',
15+ resolve: {
16+ alias: { 'numbl-src': numblSrc },
17+ },
18+ server: {
19+ // the alias resolves outside the project root (through the symlink)
20+ fs: { allow: [import.meta.dirname, numblSrc] },
21+ },
622 build: {
723 target: 'es2022',
824 rollupOptions: {