Comparing changes
solver-subroutines is 4 commits ahead of main and is 38 commits behind.
Create pull request
Use numbl main branch instead of a pinned commit in workflows
Jeremy Magland committed
abfe796Merge main into solver-subroutines
Jeremy Magland committed
cbe40a9Editable operator/solver files in the page, and a solver selector
Jeremy Magland committed
1015122Factor the solver out of the models; add BiCGSTAB and GMRES
Jeremy Magland committed
59f3e2228 changed files+2047−357
.github/workflows/ci.ymlmodified+1−1View file
@@ -24,7 +24,7 @@ jobs:
2424 # node_modules.
2525 - name: Check out numbl (sibling dependency)
2626 env:
27- NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
27+ NUMBL_REF: main
2828 run: |
2929 git clone --filter=blob:none --no-checkout \
3030 https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
.github/workflows/deploy.ymlmodified+1−1View file
@@ -36,7 +36,7 @@ jobs:
3636 # node_modules.
3737 - name: Check out numbl (sibling dependency)
3838 env:
39- NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
39+ NUMBL_REF: main
4040 run: |
4141 git clone --filter=blob:none --no-checkout \
4242 https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
.github/workflows/preview.ymlmodified+1−1View file
@@ -96,7 +96,7 @@ jobs:
9696 # pinned to the same ref as ci.yml and deploy.yml.
9797 - name: Check out numbl (sibling dependency)
9898 env:
99- NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
99+ NUMBL_REF: main
100100 run: |
101101 git clone --filter=blob:none --no-checkout \
102102 https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
README.mdmodified+159−63View file
@@ -9,15 +9,15 @@ This is the sibling of
99 solves the same systems on the round sphere. Everything there is here; what is
1010 added is a *surface*.
1111
12-> [!WARNING]
13-> **The geometry is rendered, not yet solved on.** The Laplace–Beltrami
14-> operator in the models is still the round sphere's — the term that carries
15-> the shape is a placeholder that is identically zero. On anything but the
16-> sphere you are looking at the sphere's pattern painted onto that surface, not
17-> the pattern that surface would grow. Everything the correction needs in order
18-> to be dropped in — the embedding, the split of the operator, the iterative
19-> solve, the unrolled loop — is built and tested. See
20-> [The geometry is not in the operator yet](#the-geometry-is-not-in-the-operator-yet).
12+> [!NOTE]
13+> **The geometry is in the operator.** The models solve with the surface's
14+> Laplace–Beltrami operator, iterated by a fixed-count preconditioned
15+> Richardson solve ([`solvers/richardson.m`](solvers/richardson.m), applying
16+> [`lib/dlap.m`](lib/dlap.m)). The iteration count is fixed at compile time
17+> with no residual check, so a shape/timestep/diffusivity combination outside
18+> its convergence radius diverges over many steps rather than being caught —
19+> the tests pin the known cases. See
20+> [Where the geometry enters the operator](#where-the-geometry-enters-the-operator).
2121
2222 ## What a surface is here
2323
@@ -103,12 +103,29 @@ Unew = (B + dt*D*dlap(Unew)) ./ (1 + dt*D*lam)
103103 and the loop iterates it from the round-sphere answer. That is preconditioned
104104 Richardson, with the operator we can invert exactly as the preconditioner; it
105105 converges while `dt*D*dlap` stays small against `(I - dt*D*lap_s)`, which is
106-what would keep the cost to a few transforms per step rather than a full
107-elliptic solve. Written out, the whole of
106+what keeps the cost to a few transforms per step rather than a full elliptic
107+solve.
108+
109+The pieces of that sentence are separate files, because they are separate
110+ideas. The **operator** — `dlap` applied to a spectral field — is
111+[`lib/dlap.m`](lib/dlap.m). The **solver** — the fixed point above, iterated
112+`niter` times — is [`solvers/richardson.m`](solvers/richardson.m):
113+
114+```matlab
115+function X = richardson(B, dtD, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)
116+ X = B ./ (1 + dtD * lam);
117+ for k = 1:niter
118+ dL = dlap(X, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
119+ X = (B + dtD * dL) ./ (1 + dtD * lam);
120+ end
121+end
122+```
123+
124+And a **model** is a reaction plus one solve per species — the whole of
108125 [`models/schnakenberg.m`](models/schnakenberg.m)'s step is:
109126
110127 ```matlab
111-function [Un, Vn, u, v] = step(U, V, lam, gx, gy, gz, a, b, D1, D2, dt, niter)
128+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, nlm, niter)
112129 u = synth(U);
113130 v = synth(V);
114131 uuv = u .* u .* v;
@@ -116,44 +133,95 @@ function [Un, Vn, u, v] = step(U, V, lam, gx, gy, gz, a, b, D1, D2, dt, niter)
116133 Bu = U + dt * analys(a - u + uuv);
117134 Bv = V + dt * analys(b - uuv);
118135
119- Un = Bu ./ (1 + (dt * D1) * lam);
120- Vn = Bv ./ (1 + (dt * D2) * lam);
121-
122- for k = 1:niter
123- dLu = 0 * Un; % <- the placeholder
124- dLv = 0 * Vn;
125- Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
126- Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
127- end
136+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
137+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
128138 end
129139 ```
130140
131-Written this way rather than as a residual correction on purpose: with `dlap`
132-zero, every iterate is *bit for bit* the first line, with no cancellation to
133-round differently. So the sphere case is not "close to" turing-sphere, it is
134-the same arithmetic, and the tests assert exactly that — the state after 20
135-steps is identical at 0, 1 and 4 iterations.
136-
137-### The geometry is not in the operator yet
138-
139-What belongs where `dLu` is now is `dlap = lap_g - lap_s` applied to the current
140-iterate. Getting it needs two things this repo does not have:
141-
142-1. **The induced metric**, `g_ij = ∂_i X · ∂_j X` for `X = (gx, gy, gz)`. The
143- geometry is static and low-degree, so this is a one-off precomputation, not
144- per-step work — but it needs θ- and φ-derivatives of the embedding.
145-2. **Surface derivatives of the field**, per iteration. In the round frame this
146- is the spheroidal transform pair — SHTNS's `SHsph_to_spat` and
147- `spat_to_SHsph`, i.e. `grad_s` and `div_s` — which lets the operator be
148- written as `div_s(A grad_s f)` with `A` built from the metric, with no
149- explicit `1/sin θ` to go singular at the poles.
150-
151-Both need Legendre *derivative* tables, which the vendored WGSL transforms under
152-[`src/sht/`](src/sht/) do not implement — they are scalar synthesis and analysis
153-only. That is the missing piece, and it is a substantial addition to the
154-transforms rather than a change to the models. Until it lands, the models take
155-`gx, gy, gz` (the surface on the grid) and `Gx, Gy, Gz` (the same surface as
156-coefficients) as arguments and do not use them, and the app says so.
141+`solve` is a one-line, host-generated shim (`src/mgpu/libs.ts`) forwarding to
142+whichever solver the app's **solver control** currently selects — every
143+solver composes from `dlap` (the matvec is
144+`(1 + dtD.*lam).*x - dtD.*dlap(x)`, the preconditioner the elementwise
145+divide), and the choice is part of what compiles, so switching recompiles,
146+like changing `niter` already does. A model may equally name a solver
147+directly in those call lines. The operator and every solver are files in the
148+page's editor, next to the model and the surface — edit one and recompile,
149+same as the model. The Richardson solver is written as a full re-evaluation
150+rather than an accumulated correction on purpose: where `dlap` computes to
151+zero there is no correction to mis-round, and the divide is turing-sphere's
152+arithmetic unchanged.
153+
154+Three solvers ship. [`solvers/bicgstab.m`](solvers/bicgstab.m) solves the
155+same system by preconditioned BiCGSTAB — same `dlap`, same preconditioner, a
156+Krylov recurrence instead of a stationary one, at two `dlap` evaluations per
157+iteration instead of one. Its scalars (`rho`, `alpha`, `omega`) never touch
158+the CPU: `dot` is a GPU reduction into a 1-element buffer, the recurrences on
159+its results compile to 1-element kernels, and a single-element value
160+broadcasts into the vector updates. Inner products carry the half-spectrum
161+weight `wlm` (m > 0 counts twice), making them the real L2 inner products on
162+the sphere. With no residual test, every ratio `a/b` is written in the
163+guarded form `a*b/(b*b + 1e-30)`, so a converged (or broken-down) iteration
164+goes stationary instead of dividing noise by noise. The difference is not
165+academic: at the app's default lmax, Schnakenberg on the peanut sits outside
166+the Richardson iteration's convergence radius for `niter ≥ 2` and diverges,
167+while BiCGSTAB on the identical operator converges monotonically — the tests
168+pin both behaviors, side by side.
169+
170+[`solvers/gmres.m`](solvers/gmres.m) is right-preconditioned GMRES(niter) —
171+one Arnoldi sweep, no restart — with the residual minimized over the whole
172+Krylov space. Its bookkeeping is what the other solvers never need: a basis
173+of niter+1 spectral fields, a Hessenberg matrix, Givens rotations, a
174+triangular back-substitution. The basis lives in a *bank* (`getslab` /
175+`setslab`: the k-th 2 × nlm field of a wider array), the small matrices are
176+element-addressed (`getat` / `setat`), and both are functional updates the
177+planner compiles to static-offset buffer copies — MATLAB's own `H(i,j) = h`
178+cannot lower, because numbl must prove an indexed write in bounds before the
179+loop unrolls, and a loop variable has no value yet at that point. Written as
180+calls, the index resolves at *planning*, where unrolling has made it a
181+literal. The same resolution lets an inner loop bound depend on the outer
182+loop's variable, which is what makes the `for i = 1:j` orthogonalization
183+sweep compile.
184+
185+### Where the geometry enters the operator
186+
187+[`lib/dlap.m`](lib/dlap.m) is Algorithm 3 of the evolving-surface notes: the
188+field's θ/φ derivatives (the `dtheta`/`dphi` transforms,
189+[`src/sht/deriv.ts`](src/sht/deriv.ts)) are contracted through the inverse
190+metric quantities into a tangential gradient; each Cartesian component is
191+re-analysed and differentiated again; the results recombine into the surface
192+divergence, and `lam .* F` adds back what the round-sphere part already
193+carries. The metric quantities `Vt*`/`Vp*`
194+([`src/geom/metric.ts`](src/geom/metric.ts)) are built once from the
195+embedding's derivatives when the geometry is (re)built — the geometry is
196+static, so per step they are just six more buffers the kernels read. `filt`
197+zeroes the top two spectral degrees wherever the operator re-differentiates,
198+because the derivative recurrences cannot exactly represent a derivative
199+there.
200+
201+On the sphere `dlap` computes to (numerical) zero, so any `niter` lands
202+within transform round-off of the exact round-sphere answer — asserted in the
203+tests. Off the sphere the correction genuinely moves the answer, and
204+convergence is a real constraint: the fixed-count loop has no residual check,
205+so the tests also pin which shape/niter combinations are known to sit outside
206+the convergence radius and diverge.
207+
208+### Subroutines
209+
210+A model file is not limited to `init` and `step`: it can define further
211+functions and call them, and every model compiles against the shared library
212+files — [`lib/`](lib/) for operators, [`solvers/`](solvers/) for solvers —
213+with MATLAB's visibility rules (a file's namesake function is public; a
214+model-local function of the same name shadows it). numbl specializes each
215+callee for the argument types at its call sites, and the host then splices
216+the lowered body into the caller, one clone per call site
217+([`src/mgpu/inlineCalls.ts`](src/mgpu/inlineCalls.ts)): arguments bind by
218+renaming rather than copying, and assignments to a callee output become
219+assignments to the caller's variable, which is what lets a solver iterate its
220+result in place. Expansion runs before the fusion pass, so a call fuses
221+exactly as the same code written inline would — the boundary costs nothing,
222+and `describe()`'s op listing names the expanded internals
223+(`richardson#1.X`). Recursion cannot unroll into a fixed op sequence and is
224+refused at compile time, like a runtime loop bound.
157225
158226 ### `for` loops, unrolled
159227
@@ -161,7 +229,8 @@ A plan is a fixed list of GPU operations with no branching, which is what makes
161229 a timestep pure command recording — one submit, no CPU in the loop. A counted
162230 loop still fits: the planner
163231 ([`src/mgpu/plan.ts`](src/mgpu/plan.ts)) unrolls it, planning the body once per
164-iteration.
232+iteration. The loop that matters is `solvers/richardson.m`'s `for k = 1:niter`,
233+expanded into each model's step at every solve call site.
165234
166235 Nothing else had to change for that, because numbl gives a variable one cName
167236 for every assignment to it: the buffer an iteration writes is the buffer the
@@ -173,8 +242,11 @@ Two consequences worth stating:
173242
174243 - **The bounds must be known when the model compiles.** `niter` is supplied as a
175244 fixed scalar rather than a tunable one, so changing it recompiles — unlike a
176- parameter, which is a uniform. A runtime bound is refused at compile time with
177- a source position, not silently mis-compiled, and there is a test for that.
245+ parameter, which is a uniform. A bound may also be an enclosing unrolled
246+ loop's variable (`for i = 1:j` — each unrolled `j` plans its own inner trip
247+ count, which is how GMRES's triangular sweeps compile). A genuinely runtime
248+ bound is refused at compile time with a source position, not silently
249+ mis-compiled, and there is a test for that.
178250 - **Fusion survives.** numbl's inline pass recurses into loop bodies, so a line
179251 inside the loop is still one kernel. It runs there with no protected names,
180252 though, which means an assignment whose only visible use is later in the same
@@ -183,23 +255,32 @@ Two consequences worth stating:
183255 each loop body assigns before the pass and refuses the ones that escape, so
184256 that case is a compile error rather than a stale read.
185257
186-Unrolling is exactly linear in the trip count: 2 GPU ops per species per
187-iteration, asserted in the tests.
258+Unrolling is exactly linear in the trip count: 26 GPU ops per species per
259+iteration (the operator's twelve transforms and the solve's kernels), asserted
260+in the tests.
188261
189262 ## MATLAB, compiled to WebGPU
190263
191264 Unchanged from turing-sphere, and it now compiles the geometry files too. numbl
192265 parses and lowers each function for the concrete argument types of the current
193-grid; its inline pass folds single-use temps back into their consumer, so one
266+grid; user-function calls are expanded into the caller, one clone per call
267+site; the inline pass folds single-use temps back into their consumer, so one
194268 line of MATLAB becomes one expression tree; and this repo emits one WGSL compute
195269 kernel per element-wise statement
196-([`src/mgpu/wgsl.ts`](src/mgpu/wgsl.ts)). `synth` / `analys` are external
197-operations whose type rules numbl learns from a `.mtoc2.js` workspace file, and
198-which the backend maps onto the spherical-harmonic pipelines. Anything it cannot
199-express is refused at compile time with a source position.
200-
201-The Schnakenberg step above compiles to 17 GPU operations at one solve
202-iteration: 4 transforms, 11 generated kernels, and 2 buffer copies feeding the
270+([`src/mgpu/wgsl.ts`](src/mgpu/wgsl.ts)). `synth` / `analys` (and the
271+derivative pair `dtheta` / `dphi`) are external operations whose type rules
272+numbl learns from a `.mtoc2.js` workspace file, and which the backend maps onto
273+the spherical-harmonic pipelines; `dot` is one more, mapped onto a
274+single-dispatch reduction ([`src/mgpu/reduce.ts`](src/mgpu/reduce.ts)) whose
275+1-element result stays on the GPU — scalars computed from it become 1-element
276+kernels, and reading one inside a vector expression broadcasts it. The
277+indexed-access ops (`getslab`/`setslab`, `getat`/`setat`) compile to
278+static-offset buffer copies, their indices evaluated at planning time where
279+the unrolled loop's variable is a literal. Anything it cannot express is
280+refused at compile time with a source position.
281+
282+The Schnakenberg step above compiles to 65 GPU operations at one solve
283+iteration: 28 transforms, 35 generated kernels, and 2 buffer copies feeding the
203284 new state back.
204285
205286 Two consequences carried over:
@@ -263,6 +344,7 @@ currently simulating:
263344
264345 ```
265346 npm run bench -- --preset schnak-spots --geometry ellipsoid --lmax 63 --niter 1 \
347+ --solver richardson \
266348 --steps 2000 --seed 1 --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05 \
267349 --gax 1.5 --gay 1 --gaz 0.6
268350 ```
@@ -350,8 +432,16 @@ about the round sphere, so all three build on the sphere geometry:
350432 **the same coefficients give the same surface on a 2× grid** — the 2× Gauss
351433 latitudes share no point with the 1× ones, so agreeing there is agreeing
352434 everywhere, which is what "rendered exactly, not subdivided" means;
353-- unrolling is **exactly linear** in the trip count, and the state after 20 steps
354- is **bit-identical** at 0, 1 and 4 iterations;
435+- unrolling is **exactly linear** in the trip count; on the sphere the
436+ geometric correction computes to (numerical) zero, so 0, 1 and 4 iterations
437+ agree to transform round-off; on the peanut it **measurably moves the
438+ answer**;
439+- a niter × geometry sweep stays finite except the combinations **known to sit
440+ outside the Richardson convergence radius**, which are pinned as diverging —
441+ and on exactly those combinations **bicgstab and gmres keep converging**,
442+ also pinned;
443+- at equal niter, **bicgstab and gmres land far closer to the converged
444+ answer** than richardson on the same operator;
355445 - a runtime loop bound is refused at compile time;
356446 - swapping the surface mid-run leaves the spectral state untouched.
357447
@@ -359,7 +449,13 @@ about the round sphere, so all three build on the sphere geometry:
359449 and asserts **how many kernels it compiles to**, split into the base step and
360450 what one solve iteration adds. That is a fusion guard: if numbl's inline pass
361451 stops folding, the results stay correct while every operator becomes its own
362-dispatch, which is invisible in the numbers.
452+dispatch, which is invisible in the numbers. It also compiles a model built of
453+**user-defined subroutines** — multi-output, scalar-returning, a solver-like
454+local with its own loop — asserts recursion is refused, and checks the
455+**reduction and indexing primitives** directly against exact expected values:
456+the `dot` sum and the `wlm`-weighted inner product against the CPU, scalar
457+arithmetic on a GPU-resident result bit for bit, the 1-element broadcast, and
458+element/slab round-trips through `setat`/`getat` and `setslab`/`getslab`.
363459
364460 [`test/transformChecks.ts`](test/transformChecks.ts) compares the WGSL transforms
365461 against shtns-webgpu's f64 CPU twin.
docs/richardson-iteration.mdmodified+38−7View file
@@ -64,13 +64,19 @@ Un^(0) = B ./ (1 + dt*D*lam) [dlap = 0]
6464 Un^(k+1) = (B + dt*D*dlap(Un^(k))) ./ (1 + dt*D*lam)
6565 ```
6666
67-for `k = 0 .. niter-1`. `models/schnakenberg.m`'s `for k = 1:niter` loop *is*
68-this: `Un^(0)` is the divide computed just before the loop, and each pass
69-computes `Un^(k+1)` from `Un^(k)`. It is written as a full re-evaluation
70-rather than an accumulated correction `δ = Un^(k+1) - Un^(k)` on purpose: at
71-`dlap ≡ 0` (the round sphere), every `Un^(k)` is then bit-for-bit `Un^(0)`,
72-with no cancellation to round differently — a stronger, and cheaper to
73-check, statement than "close to the round-sphere answer."
67+for `k = 0 .. niter-1`. `solvers/richardson.m`'s `for k = 1:niter` loop *is*
68+this: `Un^(0)` is the divide computed just before the loop, each pass
69+computes `Un^(k+1)` from `Un^(k)`, and `dlap` — evaluated once per iteration —
70+is its own function, `lib/dlap.m`. A model's step calls the solver once per
71+species (`Un = solve(Bu, dt * D1, ...)`, routed to the selected solvers/*.m
72+file by a host-generated shim), which is where the split pays: a different
73+solver for the same operator is a selector change — or a different call in
74+the model — with `lib/dlap.m` untouched. The solver is written as a full re-evaluation
75+rather than an accumulated correction `δ = Un^(k+1) - Un^(k)` on purpose:
76+where `dlap` evaluates to zero exactly, every `Un^(k)` is bit-for-bit
77+`Un^(0)`, with no cancellation to round differently. (In practice `dlap` is a
78+real computation through chained fp32 transforms, so on the round sphere it
79+lands near zero rather than at it — the tests bound how near.)
7480
7581 ## Convergence, and why it isn't GMRES
7682
@@ -102,6 +108,31 @@ algos.tex, one `dlap` evaluation per iteration, a fixed and
102108 recompile-on-change trip count, in exchange for linear rather than
103109 superlinear convergence.
104110
111+A Krylov method fits the shape too, as long as its scalars stay on the GPU —
112+which is what `solvers/bicgstab.m` does: `dot` is a reduction dispatch, the
113+alpha/omega/rho recurrences are 1-element kernels, and the iteration count is
114+still fixed and unrolled. What it cannot have is exactly what GMRES's `tol`
115+gives algos.tex: a stopping rule. It compensates in two ways — every ratio is
116+algebraically guarded (`a*b/(b*b + eps)`) so a converged iteration goes
117+stationary rather than dividing noise by noise, and the cost is fixed at two
118+`dlap` evaluations plus three reductions per iteration whether or not it has
119+already converged. In exchange it converges superlinearly, including on
120+shape/timestep combinations outside the Richardson iteration's spectral
121+radius — the tests pin Schnakenberg on the peanut at the app's default lmax
122+as exactly such a case.
123+
124+algos.tex's own method is here too, in the same fixed-count form:
125+`solvers/gmres.m` is right-preconditioned GMRES(niter) — one Arnoldi sweep,
126+Givens rotations, back-substitution — minus the restart loop and minus
127+`tol`/`maxiter`, since there is still no data-dependent stopping. Its basis
128+and Hessenberg bookkeeping run through the indexed-access ops
129+(`getslab`/`setslab`, `getat`/`setat`), which the planner compiles to
130+static-offset buffer copies once the unrolled loop's variable makes every
131+index a literal; the same guarded-ratio discipline covers the rotation and
132+back-substitution divides. Where algos.tex's GMRES stops at `tol`, this one
133+spends its fixed niter·(niter+3)/2 reductions and niter `dlap` evaluations
134+and keeps whatever residual that bought.
135+
105136 ## One more difference worth flagging
106137
107138 algos.tex maps the half-spectrum complex coefficients through a real vector
index.htmlmodified+12−26View file
@@ -18,9 +18,6 @@
1818 --tok-num: #0550ae;
1919 --tok-kw: #cf222e;
2020 --tok-ext: #8250df;
21- --warn-bg: #fff8e5;
22- --warn-line: #e3c37a;
23- --warn-edge: #bf8700;
2421 color-scheme: light dark;
2522 }
2623 @media (prefers-color-scheme: dark) {
@@ -36,9 +33,6 @@
3633 --tok-num: #79c0ff;
3734 --tok-kw: #ff7b72;
3835 --tok-ext: #d2a8ff;
39- --warn-bg: #2b2410;
40- --warn-line: #6b5518;
41- --warn-edge: #e3b341;
4236 }
4337 }
4438 body {
@@ -111,17 +105,6 @@
111105 color: var(--ink); user-select: all;
112106 }
113107 #blurb { margin-top: 4px; font-size: 13px; color: var(--ink-2); }
114- .warn {
115- margin: 0 0 14px;
116- padding: 10px 14px;
117- border: 1px solid var(--warn-line);
118- border-left: 5px solid var(--warn-edge);
119- border-radius: 6px;
120- background: var(--warn-bg);
121- color: var(--ink);
122- font-size: 13.5px; line-height: 1.5;
123- }
124- .warn b { color: var(--warn-edge); }
125108 #err { color: #b35900; white-space: pre-wrap; font-size: 13px; }
126109 .editor {
127110 margin-top: 12px; border: 1px solid var(--line); border-radius: 8px;
@@ -181,29 +164,32 @@
181164 <body>
182165 <main>
183166 <h1>turing-surface</h1>
184- <p class="warn">
185- <b>⚠ Work in progress: the geometry is drawn, not solved on.</b>
186- The solver still uses the round sphere's Laplace–Beltrami operator, so
187- on other shapes you see the sphere's pattern painted onto that surface.
188- </p>
189167 <p class="sub">
190168 Reaction-diffusion on closed surfaces, solved live with spherical
191169 harmonics on WebGPU via
192170 <a href="https://github.com/concept-collection/shtns-webgpu">shtns-webgpu</a>.
193- Both the solver and the shape are the MATLAB below, compiled in your
194- browser by <a href="https://numbl.org">numbl</a>. Edit either and watch
195- it change. Drag to rotate.
171+ The model, the shape, the diffusion operator and the solvers are all
172+ the MATLAB below, compiled in your browser by
173+ <a href="https://numbl.org">numbl</a>. Edit any of them and watch it
174+ change. Drag to rotate.
196175 </p>
197176 <div class="controls">
198177 <label>preset
199178 <select id="model"></select>
200179 </label>
201- <label title="The surface. Rendered, but not yet in the operator.">geometry
180+ <label title="The surface being solved on.">geometry
202181 <select id="geometry"></select>
203182 </label>
204183 <label title="Blend between the sphere (0) and the surface (1). Display only.">morph
205184 <input type="range" id="morph" min="0" max="1" step="0.01" value="1" />
206185 </label>
186+ <label title="Which solver answers the models' solve(...) call — see solvers/. Changing it recompiles.">solver
187+ <select id="solver">
188+ <option value="richardson" selected>richardson</option>
189+ <option value="bicgstab">bicgstab</option>
190+ <option value="gmres">gmres</option>
191+ </select>
192+ </label>
207193 <label title="Iterations of the implicit diffusion solve. Changing it recompiles.">solve iters
208194 <select id="niter">
209195 <option value="0">0</option>
lib/dlap.madded+44−0View file
@@ -0,0 +1,44 @@
1+% The geometric part of the surface Laplace-Beltrami operator:
2+%
3+% dlap = lap_g - lap_s
4+%
5+% applied to a spectral field F, where lap_g is the surface's operator and
6+% lap_s the round sphere's (diagonal, eigenvalues -l(l+1)). This is the piece
7+% of the implicit solve (I - dt*D*lap_g) X = B that a solver re-evaluates
8+% each iteration — the round-sphere part it inverts exactly. The operator is
9+% linear in F; the surface enters only through the inverse metric quantities
10+% Vt*/Vp* (src/geom/metric.ts), so swapping the geometry changes no code.
11+%
12+% Algorithm 3 of evolving_surface/notes/algos.tex: surface gradient of the
13+% field (theta/phi derivatives contracted through Vt*/Vp*); each Cartesian
14+% component re-analysed and differentiated again; recombined into the surface
15+% divergence. lam .* F adds back -lap_s(F), since lam holds +l(l+1). filt
16+% zeroes the top two degrees, where the theta/phi derivative recurrences
17+% cannot exactly represent a derivative. See docs/richardson-iteration.md.
18+%
19+% Spectral fields are real 2 x nlm; the intermediate d*/L fields live on the
20+% npts x 1 grid.
21+
22+function dL = dlap(F, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam)
23+ G = F .* filt;
24+ Ft = dtheta(G);
25+ Fp = dphi(G);
26+ dx = Ft .* Vtx + Fp .* Vpx;
27+ dy = Ft .* Vty + Fp .* Vpy;
28+ dz = Ft .* Vtz + Fp .* Vpz;
29+ cx = analys(dx) .* filt;
30+ cy = analys(dy) .* filt;
31+ cz = analys(dz) .* filt;
32+ Ftcx = dtheta(cx);
33+ Fpcx = dphi(cx);
34+ Ftcy = dtheta(cy);
35+ Fpcy = dphi(cy);
36+ Ftcz = dtheta(cz);
37+ Fpcz = dphi(cz);
38+ L = Ftcx .* Vtx + Fpcx .* Vpx;
39+ L = L + Ftcy .* Vty;
40+ L = L + Fpcy .* Vpy;
41+ L = L + Ftcz .* Vtz;
42+ L = L + Fpcz .* Vpz;
43+ dL = analys(L) + lam .* F;
44+end
models/allencahn.mmodified+4−32View file
@@ -2,45 +2,17 @@
22 %
33 % du/dt = eps2*lap_g(u) + u - u^3
44 %
5-% Same scheme as models/schnakenberg.m.
5+% Same scheme as models/schnakenberg.m: explicit reaction, then the implicit
6+% diffusion solve handed to solve(...) — the solver the app's selector picks.
67
78 function [U, u] = init(noise)
89 U = analys(noise);
910 u = synth(U);
1011 end
1112
12-function [Un, u] = step(U, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, eps2, dt, niter)
13+function [Un, u] = step(U, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, eps2, dt, nlm, niter)
1314 u = synth(U);
1415
1516 Bu = U + dt * analys(u - u.^3);
16- Un = Bu ./ (1 + (dt * eps2) * lam);
17-
18- for k = 1:niter
19- % dlap = lap_g - lap_s, evaluated at the current iterate (see
20- % models/schnakenberg.m and docs/richardson-iteration.md for the
21- % derivation).
22- Fu = Un .* filt;
23- Ftu = dtheta(Fu);
24- Fpu = dphi(Fu);
25- dux = Ftu .* Vtx + Fpu .* Vpx;
26- duy = Ftu .* Vty + Fpu .* Vpy;
27- duz = Ftu .* Vtz + Fpu .* Vpz;
28- cux = analys(dux) .* filt;
29- cuy = analys(duy) .* filt;
30- cuz = analys(duz) .* filt;
31- Ftcux = dtheta(cux);
32- Fpcux = dphi(cux);
33- Ftcuy = dtheta(cuy);
34- Fpcuy = dphi(cuy);
35- Ftcuz = dtheta(cuz);
36- Fpcuz = dphi(cuz);
37- lapu = Ftcux .* Vtx + Fpcux .* Vpx;
38- lapu = lapu + Ftcuy .* Vty;
39- lapu = lapu + Fpcuy .* Vpy;
40- lapu = lapu + Ftcuz .* Vtz;
41- lapu = lapu + Fpcuz .* Vpz;
42- dLu = analys(lapu) + lam .* Un;
43-
44- Un = (Bu + (dt * eps2) * dLu) ./ (1 + (dt * eps2) * lam);
45- end
17+ Un = solve(Bu, dt * eps2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
4618 end
models/brusselator.mmodified+5−56View file
@@ -3,7 +3,8 @@
33 % du/dt = D1*lap_g(u) + A - (B+1)*u + u^2*v
44 % dv/dt = D2*lap_g(v) + B*u - u^2*v
55 %
6-% Same scheme as models/schnakenberg.m.
6+% Same scheme as models/schnakenberg.m: explicit reaction, then the implicit
7+% diffusion solve handed to solve(...) — the solver the app's selector picks.
78
89 function [U, V, u, v] = init(noise, A, B)
910 U = analys(A + noise);
@@ -12,7 +13,7 @@ function [U, V, u, v] = init(noise, A, B)
1213 v = synth(V);
1314 end
1415
15-function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, A, B, D1, D2, dt, niter)
16+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, A, B, D1, D2, dt, nlm, niter)
1617 u = synth(U);
1718 v = synth(V);
1819 uuv = u .* u .* v;
@@ -20,58 +21,6 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
2021 Bu = U + dt * analys(A - (B + 1) * u + uuv);
2122 Bv = V + dt * analys(B * u - uuv);
2223
23- Un = Bu ./ (1 + (dt * D1) * lam);
24- Vn = Bv ./ (1 + (dt * D2) * lam);
25-
26- for k = 1:niter
27- % dlap = lap_g - lap_s, evaluated at the current iterate (see
28- % models/schnakenberg.m and docs/richardson-iteration.md for the
29- % derivation).
30- Fu = Un .* filt;
31- Ftu = dtheta(Fu);
32- Fpu = dphi(Fu);
33- dux = Ftu .* Vtx + Fpu .* Vpx;
34- duy = Ftu .* Vty + Fpu .* Vpy;
35- duz = Ftu .* Vtz + Fpu .* Vpz;
36- cux = analys(dux) .* filt;
37- cuy = analys(duy) .* filt;
38- cuz = analys(duz) .* filt;
39- Ftcux = dtheta(cux);
40- Fpcux = dphi(cux);
41- Ftcuy = dtheta(cuy);
42- Fpcuy = dphi(cuy);
43- Ftcuz = dtheta(cuz);
44- Fpcuz = dphi(cuz);
45- lapu = Ftcux .* Vtx + Fpcux .* Vpx;
46- lapu = lapu + Ftcuy .* Vty;
47- lapu = lapu + Fpcuy .* Vpy;
48- lapu = lapu + Ftcuz .* Vtz;
49- lapu = lapu + Fpcuz .* Vpz;
50- dLu = analys(lapu) + lam .* Un;
51-
52- Fv = Vn .* filt;
53- Ftv = dtheta(Fv);
54- Fpv = dphi(Fv);
55- dvx = Ftv .* Vtx + Fpv .* Vpx;
56- dvy = Ftv .* Vty + Fpv .* Vpy;
57- dvz = Ftv .* Vtz + Fpv .* Vpz;
58- cvx = analys(dvx) .* filt;
59- cvy = analys(dvy) .* filt;
60- cvz = analys(dvz) .* filt;
61- Ftcvx = dtheta(cvx);
62- Fpcvx = dphi(cvx);
63- Ftcvy = dtheta(cvy);
64- Fpcvy = dphi(cvy);
65- Ftcvz = dtheta(cvz);
66- Fpcvz = dphi(cvz);
67- lapv = Ftcvx .* Vtx + Fpcvx .* Vpx;
68- lapv = lapv + Ftcvy .* Vty;
69- lapv = lapv + Fpcvy .* Vpy;
70- lapv = lapv + Ftcvz .* Vtz;
71- lapv = lapv + Fpcvz .* Vpz;
72- dLv = analys(lapv) + lam .* Vn;
73-
74- Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
75- Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
76- end
24+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
25+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
7726 end
models/schnakenberg.mmodified+11−66View file
@@ -3,11 +3,13 @@
33 % du/dt = D1*lap_g(u) + a - u + u^2*v
44 % dv/dt = D2*lap_g(v) + b - u^2*v
55 %
6-% Explicit reaction, implicit diffusion (IMEX Euler). The implicit solve
7-% splits lap_g = lap_s + dlap: the round-sphere part lap_s is diagonal in
8-% spherical-harmonic space (eigenvalues -lam), and the loop iterates the
9-% geometric correction dlap from that exact solve. Grid fields are npts x 1;
10-% spectral fields are real 2 x nlm. See docs/richardson-iteration.md.
6+% Explicit reaction, implicit diffusion (IMEX Euler): the step forms the
7+% right-hand side B of the linear system (I - dt*D*lap_g) Unew = B, and hands
8+% the solve to solve(...) — the solver the app's selector picks (richardson,
9+% bicgstab or gmres, see solvers/), every one applying the operator's
10+% geometric part through lib/dlap.m. A model may also name a solver directly
11+% in these two call lines. Grid fields are npts x 1; spectral fields are
12+% real 2 x nlm. See docs/richardson-iteration.md.
1113
1214 function [U, V, u, v] = init(noise, a, b)
1315 us = a + b;
@@ -18,7 +20,7 @@ function [U, V, u, v] = init(noise, a, b)
1820 v = synth(V);
1921 end
2022
21-function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, niter)
23+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, nlm, niter)
2224 u = synth(U);
2325 v = synth(V);
2426 uuv = u .* u .* v;
@@ -27,64 +29,7 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
2729 Bu = U + dt * analys(a - u + uuv);
2830 Bv = V + dt * analys(b - uuv);
2931
30- % Round-sphere solve, then iterate the geometric correction.
31- Un = Bu ./ (1 + (dt * D1) * lam);
32- Vn = Bv ./ (1 + (dt * D2) * lam);
33-
34- for k = 1:niter
35- % dlap = lap_g - lap_s, evaluated at the current iterate (Algorithm 3 of
36- % evolving_surface/notes/algos.tex): surface gradient of the field,
37- % contracted through the inverse metric quantities Vt*/Vp*; each
38- % Cartesian component re-analysed and differentiated again; recombined
39- % into the surface divergence. lam.*Un adds back -lap_s(Un), since lam
40- % holds +l(l+1). filt zeroes the top two degrees, where the theta/phi
41- % derivative recurrences cannot exactly represent a derivative. See
42- % docs/richardson-iteration.md.
43- Fu = Un .* filt;
44- Ftu = dtheta(Fu);
45- Fpu = dphi(Fu);
46- dux = Ftu .* Vtx + Fpu .* Vpx;
47- duy = Ftu .* Vty + Fpu .* Vpy;
48- duz = Ftu .* Vtz + Fpu .* Vpz;
49- cux = analys(dux) .* filt;
50- cuy = analys(duy) .* filt;
51- cuz = analys(duz) .* filt;
52- Ftcux = dtheta(cux);
53- Fpcux = dphi(cux);
54- Ftcuy = dtheta(cuy);
55- Fpcuy = dphi(cuy);
56- Ftcuz = dtheta(cuz);
57- Fpcuz = dphi(cuz);
58- lapu = Ftcux .* Vtx + Fpcux .* Vpx;
59- lapu = lapu + Ftcuy .* Vty;
60- lapu = lapu + Fpcuy .* Vpy;
61- lapu = lapu + Ftcuz .* Vtz;
62- lapu = lapu + Fpcuz .* Vpz;
63- dLu = analys(lapu) + lam .* Un;
64-
65- Fv = Vn .* filt;
66- Ftv = dtheta(Fv);
67- Fpv = dphi(Fv);
68- dvx = Ftv .* Vtx + Fpv .* Vpx;
69- dvy = Ftv .* Vty + Fpv .* Vpy;
70- dvz = Ftv .* Vtz + Fpv .* Vpz;
71- cvx = analys(dvx) .* filt;
72- cvy = analys(dvy) .* filt;
73- cvz = analys(dvz) .* filt;
74- Ftcvx = dtheta(cvx);
75- Fpcvx = dphi(cvx);
76- Ftcvy = dtheta(cvy);
77- Fpcvy = dphi(cvy);
78- Ftcvz = dtheta(cvz);
79- Fpcvz = dphi(cvz);
80- lapv = Ftcvx .* Vtx + Fpcvx .* Vpx;
81- lapv = lapv + Ftcvy .* Vty;
82- lapv = lapv + Fpcvy .* Vpy;
83- lapv = lapv + Ftcvz .* Vtz;
84- lapv = lapv + Fpcvz .* Vpz;
85- dLv = analys(lapv) + lam .* Vn;
86-
87- Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
88- Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
89- end
32+ % The species diffuse independently, so each gets its own solve.
33+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
34+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
9035 end
scripts/bench.tsmodified+5−1View file
@@ -20,6 +20,7 @@ import { requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
2020 import { ModelSession } from '../src/mgpu/session.ts';
2121 import { presets } from '../src/mgpu/registry.ts';
2222 import { mGeometries, DEFAULT_GEOMETRY_KEY } from '../src/geom/registry.ts';
23+import { DEFAULT_SOLVER, solverKeys } from '../src/mgpu/libs.ts';
2324 import {
2425 parseArgs,
2526 modelForSpec,
@@ -47,6 +48,8 @@ const USAGE = `usage: ${BENCH_COMMAND} [options]
4748 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
4849 --niter <n> iterations of the implicit solve, unrolled into the compiled
4950 step (default ${DEFAULT_NITER})
51+ --solver <key> ${solverKeys.join(' | ')} — which solver answers the
52+ models' solve(...) call (default ${DEFAULT_SOLVER})
5053 --steps <n> timed steps (default ${DEFAULT_STEPS})
5154 --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
5255 --seed <n> initial-noise seed (default ${DEFAULT_SEED})
@@ -175,6 +178,7 @@ try {
175178 geometry,
176179 geometryParams: spec.geometryParams,
177180 niter: spec.niter,
181+ solver: spec.solver,
178182 });
179183 session.seed(spec.seed);
180184
@@ -198,7 +202,7 @@ try {
198202 );
199203 console.log(
200204 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
201- `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'}`,
205+ `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'} of ${spec.solver}`,
202206 );
203207 console.log(` compiled ${plan.step.length} GPU ops/step (${kernels} generated kernels)`);
204208 console.log(` fourier ${session.sht.fourierMode.toUpperCase()} stage`);
solvers/bicgstab.madded+63−0View file
@@ -0,0 +1,63 @@
1+% Solve the implicit diffusion system
2+%
3+% (I - dtD*lap_g) X = B, i.e. A*X = B with A*x = M.*x - dtD*dlap(x)
4+%
5+% by preconditioned BiCGSTAB (van der Vorst), with the round-sphere operator
6+% M = 1 + dtD*lam as the preconditioner — the same M solvers/richardson.m
7+% inverts, applied here as M\v inside the Krylov recurrence. Same operator,
8+% different solver: where Richardson converges only while the geometric
9+% correction stays small against M, BiCGSTAB builds a Krylov space and
10+% converges on configurations well outside that radius, at the price of two
11+% dlap evaluations per iteration instead of one, plus three dot products.
12+%
13+% The scalars (rho, alpha, omega) are GPU-resident 1-element values: `dot` is
14+% a reduction dispatch and the recurrences on its results compile to
15+% 1-element kernels, so the whole solve is still one command stream with no
16+% CPU in the loop. Inner products are taken in the half-spectrum weighting
17+% wlm (m > 0 counts twice), so they are the real L2 inner products on the
18+% sphere; the weight is folded into the fixed shadow residual once, and into
19+% t per iteration.
20+%
21+% niter is fixed at compile time: no residual test, and so no early exit —
22+% which for BiCGSTAB matters at the *converged* end, where the residual is
23+% fp32 noise and every textbook ratio is 0/0. Each ratio a/b is therefore
24+% written in the guarded form a*b/(b*b + 1e-30): identical to a/b (to a
25+% couple of ulp) whenever b is meaningfully sized, and 0 when it is not, so
26+% a converged or broken-down iteration degrades to a stationary one
27+% (coefficients 0, X carried unchanged) instead of poisoning the state with
28+% NaNs. X0 is the round-sphere answer, so niter = 0 computes richardson's
29+% starting divide.
30+
31+function X = bicgstab(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)
32+ M = 1 + dtD * lam;
33+ X = B ./ M;
34+ dL0 = dlap(X, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
35+ R = B - (M .* X - dtD * dL0);
36+ Rh = R .* wlm;
37+ P = 0 * R;
38+ V = 0 * R;
39+ rho = 1;
40+ alpha = 1;
41+ omega = 1;
42+ for k = 1:niter
43+ rho1 = dot(Rh, R);
44+ beta = (rho1 * rho) / (rho * rho + 1e-30) * ((alpha * omega) / (omega * omega + 1e-30));
45+ P = R + beta * (P - omega * V);
46+ Ph = P ./ M;
47+ dLp = dlap(Ph, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
48+ V = M .* Ph - dtD * dLp;
49+ rv = dot(Rh, V);
50+ alpha = (rho1 * rv) / (rv * rv + 1e-30);
51+ S = R - alpha * V;
52+ Sh = S ./ M;
53+ dLs = dlap(Sh, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
54+ T = M .* Sh - dtD * dLs;
55+ Tw = T .* wlm;
56+ ts = dot(Tw, S);
57+ tt = dot(Tw, T);
58+ omega = (ts * tt) / (tt * tt + 1e-30);
59+ X = X + alpha * Ph + omega * Sh;
60+ R = S - omega * T;
61+ rho = rho1;
62+ end
63+end
solvers/gmres.madded+124−0View file
@@ -0,0 +1,124 @@
1+% Solve the implicit diffusion system
2+%
3+% (I - dtD*lap_g) X = B, i.e. A*X = B with A*x = M.*x - dtD*dlap(x)
4+%
5+% by right-preconditioned GMRES(niter) — one Arnoldi sweep of niter
6+% iterations, no restart — with the round-sphere operator M = 1 + dtD*lam as
7+% the preconditioner. Same operator and preconditioner as the other two
8+% solvers; what GMRES adds over bicgstab is optimality (the residual is
9+% minimized over the whole Krylov space, monotonically nonincreasing) at the
10+% price of storing the basis and the O(niter^2) orthogonalization sweep.
11+%
12+% The bookkeeping no other solver needs — the basis, the Hessenberg matrix,
13+% the Givens rotations — lives in banks and small matrices accessed through
14+% getslab/setslab and getat/setat (src/mgpu/externals.ts): functional
15+% updates the planner turns into static-offset buffer copies once the
16+% unrolled loop's variable makes every index a literal. That is also why the
17+% triangular loops below (`for i = 1:j`) compile: each unrolled j plans its
18+% own inner trip count.
19+%
20+% Scalars are GPU-resident 1-element values throughout, as in bicgstab, and
21+% every ratio is guarded the same way (a*b/(b*b + 1e-30), and 1/sqrt(x) as
22+% 1/sqrt(x + 1e-30)), so a converged or broken-down iteration contributes
23+% zero coefficients instead of NaNs. sqrt takes abs() of its argument
24+% because the compiler cannot see that <w, w> is nonnegative.
25+%
26+% niter is fixed at compile time. X0 is the round-sphere answer, and the
27+% correction added at the end is X = X0 + M \ (V*y) accumulated slab by
28+% slab.
29+
30+function X = gmres(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter)
31+ M = 1 + dtD * lam;
32+ X = B ./ M;
33+ dL0 = dlap(X, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
34+ R = B - (M .* X - dtD * dL0);
35+
36+ % The Krylov basis (niter+1 spectral fields), the Hessenberg matrix, the
37+ % rotated right-hand side, the Givens coefficients, and the solve's y.
38+ VB = zeros(2, nlm * (niter + 1));
39+ H = zeros(niter + 1, niter);
40+ g = zeros(niter + 1, 1);
41+ c = zeros(niter, 1);
42+ s = zeros(niter, 1);
43+ y = zeros(niter, 1);
44+
45+ Rw = R .* wlm;
46+ r2 = dot(Rw, R);
47+ nr = sqrt(abs(r2));
48+ invnr = nr / (r2 + 1e-30);
49+ V1 = R * invnr;
50+ VB = setslab(VB, V1, 1);
51+ g = setat(g, nr, 1);
52+
53+ for j = 1:niter
54+ % w = A * M \ v_j
55+ Vj = getslab(VB, j);
56+ Zj = Vj ./ M;
57+ dLz = dlap(Zj, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
58+ W = M .* Zj - dtD * dLz;
59+
60+ % Modified Gram-Schmidt against every basis vector so far.
61+ for i = 1:j
62+ Vi = getslab(VB, i);
63+ Viw = Vi .* wlm;
64+ hij = dot(Viw, W);
65+ H = setat(H, hij, i, j);
66+ W = W - hij * Vi;
67+ end
68+ Ww = W .* wlm;
69+ w2 = dot(Ww, W);
70+ hn = sqrt(abs(w2));
71+ H = setat(H, hn, j + 1, j);
72+ invh = hn / (w2 + 1e-30);
73+ Vn = W * invh;
74+ VB = setslab(VB, Vn, j + 1);
75+
76+ % Apply the previous Givens rotations to column j, then form the new
77+ % one that zeroes H(j+1, j), and rotate g with it.
78+ for i = 1:j-1
79+ a = getat(H, i, j);
80+ b = getat(H, i + 1, j);
81+ ci = getat(c, i);
82+ si = getat(s, i);
83+ t1 = ci * a + si * b;
84+ t2 = ci * b - si * a;
85+ H = setat(H, t1, i, j);
86+ H = setat(H, t2, i + 1, j);
87+ end
88+ a = getat(H, j, j);
89+ b = getat(H, j + 1, j);
90+ rr = a * a + b * b;
91+ invr = 1 / sqrt(abs(rr) + 1e-30);
92+ cj = a * invr;
93+ sj = b * invr;
94+ c = setat(c, cj, j);
95+ s = setat(s, sj, j);
96+ t1 = cj * a + sj * b;
97+ H = setat(H, t1, j, j);
98+ gj = getat(g, j);
99+ t3 = cj * gj;
100+ t4 = -(sj * gj);
101+ g = setat(g, t3, j);
102+ g = setat(g, t4, j + 1);
103+ end
104+
105+ % Back-substitute the rotated (upper triangular) system H*y = g.
106+ for j = niter:-1:1
107+ acc = getat(g, j);
108+ for i = j+1:niter
109+ rji = getat(H, j, i);
110+ yi = getat(y, i);
111+ acc = acc - rji * yi;
112+ end
113+ rjj = getat(H, j, j);
114+ yj = (acc * rjj) / (rjj * rjj + 1e-30);
115+ y = setat(y, yj, j);
116+ end
117+
118+ % X = X0 + M \ (V * y), slab by slab.
119+ for j = 1:niter
120+ Vj = getslab(VB, j);
121+ yj = getat(y, j);
122+ X = X + yj * (Vj ./ M);
123+ end
124+end
solvers/richardson.madded+28−0View file
@@ -0,0 +1,28 @@
1+% Solve the implicit diffusion system
2+%
3+% (I - dtD*lap_g) X = B
4+%
5+% by preconditioned Richardson iteration, with the round-sphere operator as
6+% the preconditioner. Splitting lap_g = lap_s + dlap and moving the geometric
7+% part to the right-hand side gives the fixed point
8+%
9+% X = (B + dtD*dlap(X)) ./ (1 + dtD*lam)
10+%
11+% iterated from the round-sphere answer (lam holds +l(l+1), so the divide is
12+% the exact inverse of I - dtD*lap_s). It converges while dtD*dlap stays
13+% small against what the divide already inverts; niter is fixed at compile
14+% time — the loop is unrolled into the op sequence, so there is no residual
15+% check and no adaptive stopping. Full derivation and the map onto
16+% algos.tex's GMRES formulation: docs/richardson-iteration.md.
17+%
18+% Written as a full re-evaluation rather than an accumulated correction on
19+% purpose: on the round sphere dlap is identically zero, so every iterate is
20+% bit for bit the first divide, with no cancellation to round differently.
21+
22+function X = richardson(B, dtD, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)
23+ X = B ./ (1 + dtD * lam);
24+ for k = 1:niter
25+ dL = dlap(X, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, lam);
26+ X = (B + dtD * dL) ./ (1 + dtD * lam);
27+ end
28+end
src/bench/runSpec.tsmodified+12−0View file
@@ -22,6 +22,7 @@ import {
2222 type MGeometry,
2323 } from '../geom/registry.ts';
2424 import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
25+import { DEFAULT_SOLVER, solverKeys, type SolverKey } from '../mgpu/libs.ts';
2526
2627 export interface RunSpec {
2728 /** Preset key from the registry; fixes the model, params may still be edited. */
@@ -44,6 +45,9 @@ export interface RunSpec {
4445 /** Iterations of the .m's implicit solve. Structural: it is unrolled into
4546 * the compiled step, so it belongs to the spec rather than to the params. */
4647 niter: number;
48+ /** Which solver answers the models' solve(...) call. Structural like
49+ * niter: the choice compiles into the step. */
50+ solver: SolverKey;
4751 }
4852
4953 export const DEFAULT_NITER = 1;
@@ -113,6 +117,7 @@ export function formatCommand(spec: RunSpec): string {
113117 `--geometry ${spec.geometry}`,
114118 `--lmax ${spec.lmax}`,
115119 `--niter ${spec.niter}`,
120+ `--solver ${spec.solver}`,
116121 `--steps ${spec.steps}`,
117122 `--seed ${spec.seed}`,
118123 ...model.params.map((p) => `--${p.key} ${String(spec.params[p.key])}`),
@@ -160,6 +165,12 @@ export function parseArgs(argv: string[]): RunSpec {
160165 const { model, params } = resolvePreset(presetKey);
161166 const geometryKey = take('geometry') ?? DEFAULT_GEOMETRY_KEY;
162167 const { geometry, params: geometryParams } = resolveGeometry(geometryKey);
168+ const solverRaw = take('solver') ?? DEFAULT_SOLVER;
169+ if (!(solverKeys as string[]).includes(solverRaw)) {
170+ throw new Error(
171+ `--solver must be one of ${solverKeys.join(', ')} (got '${solverRaw}')`,
172+ );
173+ }
163174 const spec: RunSpec = {
164175 preset: presetKey,
165176 lmax: count('lmax', DEFAULT_LMAX, 1),
@@ -170,6 +181,7 @@ export function parseArgs(argv: string[]): RunSpec {
170181 geometry: geometryKey,
171182 geometryParams,
172183 niter: count('niter', DEFAULT_NITER, 0),
184+ solver: solverRaw as SolverKey,
173185 };
174186 const readInto = (into: Params, key: string, flag: string): void => {
175187 const raw = take(flag);
src/main.tsmodified+37−17View file
@@ -4,6 +4,7 @@ import { ModelSession } from './mgpu/session.ts';
44 import { mModelByKey, presets, type MModel, type Params } from './mgpu/registry.ts';
55 import { ModelCompileError, formatFailure } from './mgpu/errors.ts';
66 import { EXTERNAL_OPS } from './mgpu/externals.ts';
7+import { libPath, modelLibs, type SolverKey } from './mgpu/libs.ts';
78 import { CodeEditor } from './editor/codeEditor.ts';
89 import {
910 formatCommand,
@@ -16,7 +17,6 @@ import {
1617 mGeometries,
1718 mGeometryByKey,
1819 defaultGeometryParams,
19- SPHERE_KEY,
2020 DEFAULT_GEOMETRY_KEY,
2121 type MGeometry,
2222 } from './geom/registry.ts';
@@ -38,6 +38,7 @@ const $ = <T extends HTMLElement>(id: string): T =>
3838 const elModel = $<HTMLSelectElement>('model');
3939 const elGeometry = $<HTMLSelectElement>('geometry');
4040 const elMorph = $<HTMLInputElement>('morph');
41+const elSolver = $<HTMLSelectElement>('solver');
4142 const elNiter = $<HTMLSelectElement>('niter');
4243 const elLmax = $<HTMLSelectElement>('lmax');
4344 const elOversample = $<HTMLSelectElement>('oversample');
@@ -82,7 +83,11 @@ for (const g of mGeometries) {
8283 o.textContent = g.label;
8384 elGeometry.append(o);
8485 }
85-for (const [value, label] of [['model', 'the solver'], ['geometry', 'the surface']]) {
86+for (const [value, label] of [
87+ ['model', 'the model'],
88+ ['geometry', 'the surface'],
89+ ...modelLibs.map((f) => [`lib:${f.name}`, libPath(f.name)]),
90+]) {
8691 const o = document.createElement('option');
8792 o.value = value;
8893 o.textContent = label;
@@ -96,16 +101,17 @@ for (const name of colormapNames) {
96101 }
97102 elColormap.value = 'jet';
98103
99-/** Whichever .m is open: the solver or the surface. Both are MATLAB, compiled
100- * by the same backend, so one editor serves both. The host-provided operations
101- * are marked so the boundary between the file and what it is given is
102- * visible. */
104+/** Whichever .m is open: the model, the surface, or one of the shared
105+ * operator/solver files. All are MATLAB, compiled by the same backend, so
106+ * one editor serves them all. The host-provided operations are marked so
107+ * the boundary between the file and what it is given is visible. */
103108 const editor = new CodeEditor({
104109 textarea: elSource,
105110 overlay: elHighlight,
106- external: EXTERNAL_OPS,
111+ external: new Set(EXTERNAL_OPS.keys()),
107112 onInput: (value) => {
108113 if (editing === 'geometry') editedGeomSource = value;
114+ else if (editing.startsWith('lib:')) editedLibs.set(editing.slice(4), value);
109115 else editedSource = value;
110116 elRecompile.textContent = 'Recompile *';
111117 },
@@ -211,11 +217,14 @@ let model: MModel = mModelByKey(initial.model.key)!;
211217 let params: Params = initial.params;
212218 let geometry: MGeometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
213219 let geomParams: Params = defaultGeometryParams(geometry);
214-/** Which file the editor is showing. */
215-let editing: 'model' | 'geometry' = 'model';
220+/** Which file the editor is showing: the model, the surface, or a shared
221+ * file (`lib:<name>`). */
222+let editing: string = 'model';
216223 /** Each .m as edited in the page; `null` while it matches the file. */
217224 let editedSource: string | null = null;
218225 let editedGeomSource: string | null = null;
226+/** Shared-file working copies, by file name; absent while unedited. */
227+const editedLibs = new Map<string, string>();
219228 /** Sphere (0) to surface (1). Display only; does not touch the solver. */
220229 let morph = 1;
221230 let seed = 1;
@@ -236,6 +245,8 @@ let posBuf: Float32Array | null = null;
236245
237246 const source = (): string => editedSource ?? model.source;
238247 const geomSource = (): string => editedGeomSource ?? geometry.source;
248+const libSource = (name: string): string =>
249+ editedLibs.get(name) ?? modelLibs.find((f) => f.name === name)?.source ?? '';
239250
240251 // ---------------------------------------------------------------- UI wiring
241252 function buildParamInputs(): void {
@@ -324,10 +335,14 @@ function applyGeometryChoice(key: string): void {
324335
325336 /** Load the chosen file into the editor, keeping any unsaved edit to it. */
326337 function showEditorFile(): void {
327- editing = elEditorFile.value === 'geometry' ? 'geometry' : 'model';
338+ editing = elEditorFile.value;
328339 if (editing === 'geometry') {
329340 editor.value = geomSource();
330341 elEditorTitle.textContent = `geometries/${geometry.key}.m`;
342+ } else if (editing.startsWith('lib:')) {
343+ const name = editing.slice(4);
344+ editor.value = libSource(name);
345+ elEditorTitle.textContent = libPath(name);
331346 } else {
332347 editor.value = source();
333348 elEditorTitle.textContent = `models/${model.key}.m`;
@@ -346,6 +361,7 @@ function currentSpec(): RunSpec {
346361 geometry: geometry.key,
347362 geometryParams: geomParams,
348363 niter: Number(elNiter.value),
364+ solver: elSolver.value as SolverKey,
349365 };
350366 }
351367
@@ -359,8 +375,10 @@ elModel.addEventListener('change', () => {
359375 });
360376 elLmax.addEventListener('change', () => void rebuild());
361377 // The solve iteration count is unrolled into the compiled step, so unlike a
362-// parameter it cannot be changed without recompiling.
378+// parameter it cannot be changed without recompiling. The solver choice is a
379+// generated one-line shim compiled with the model, so it recompiles too.
363380 elNiter.addEventListener('change', () => void rebuild());
381+elSolver.addEventListener('change', () => void rebuild());
364382 // Oversampling and geometry are display-or-data changes, not code ones, so
365383 // they swap things in place rather than rebuilding the run. Serialized through
366384 // one chain: a rapid second change waits its turn.
@@ -407,11 +425,13 @@ elMovie.addEventListener('click', () => {
407425
408426 elRecompile.addEventListener('click', () => {
409427 if (editing === 'geometry') editedGeomSource = editor.value;
428+ else if (editing.startsWith('lib:')) editedLibs.set(editing.slice(4), editor.value);
410429 else editedSource = editor.value;
411430 void rebuild();
412431 });
413432 elRevert.addEventListener('click', () => {
414433 if (editing === 'geometry') editedGeomSource = null;
434+ else if (editing.startsWith('lib:')) editedLibs.delete(editing.slice(4));
415435 else editedSource = null;
416436 showEditorFile();
417437 void rebuild();
@@ -588,11 +608,9 @@ function updateGeomNote(): void {
588608 return;
589609 }
590610 const { lo, hi } = session.geometry.radiusRange();
591- const isSphere = session.geometryModel.key === SPHERE_KEY;
592611 elGeomNote.innerHTML =
593612 `<b>${session.geometryModel.label}</b> — ${session.geometryModel.blurb} ` +
594- `Radius ${lo.toFixed(3)}–${hi.toFixed(3)}.` +
595- (isSphere ? '' : ' <b>Rendered only</b> — not yet in the operator.');
613+ `Radius ${lo.toFixed(3)}–${hi.toFixed(3)}.`;
596614 }
597615
598616 /** Report a compile failure, and select the offending text in the editor. */
@@ -635,6 +653,8 @@ async function rebuild(): Promise<void> {
635653 geometryParams: geomParams,
636654 geometrySource: geomSource(),
637655 niter: Number(elNiter.value),
656+ solver: elSolver.value as SolverKey,
657+ libSources: Object.fromEntries(editedLibs),
638658 });
639659 } catch (e) {
640660 reportCompileError(e);
@@ -875,9 +895,9 @@ function submitSteps(n: number): void {
875895 * the Movie button itself becomes the cancel button. */
876896 function setMovieUi(on: boolean): void {
877897 const locked = [
878- elModel, elGeometry, elMorph, elNiter, elLmax, elOversample, elColormap,
879- elRunPause, elBenchmark, elReseed, elRecompile, elRevert, elEditorFile,
880- elMovieSpeed, elMovieRes, elMovieRotate, elMovieToggle,
898+ elModel, elGeometry, elMorph, elSolver, elNiter, elLmax, elOversample,
899+ elColormap, elRunPause, elBenchmark, elReseed, elRecompile, elRevert,
900+ elEditorFile, elMovieSpeed, elMovieRes, elMovieRotate, elMovieToggle,
881901 ];
882902 for (const el of locked) el.disabled = on;
883903 elParams.querySelectorAll('input').forEach((input) => (input.disabled = on));
src/mgpu/compile.tsmodified+30−4View file
@@ -25,7 +25,8 @@ import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
2525 import type { For, IRExpr, IRFunc, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
2626 import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
2727 import { externalOpFiles, type GridSizes } from './externals.ts';
28-import { ModelCompileError } from './errors.ts';
28+import { expandUserCalls } from './inlineCalls.ts';
29+import { inFunction, ModelCompileError } from './errors.ts';
2930
3031 /** What the host can supply for an argument the .m declares. */
3132 export type Binding =
@@ -84,16 +85,26 @@ export class CompiledModel {
8485 #lowerer: Lowerer;
8586 #decls: Map<string, FunctionDecl>;
8687 #bindings: Record<string, Binding>;
88+ /** Functions handed out by specialize(), in order — the ones whose bodies
89+ * the planner will execute, and so the ones finish() expands. */
90+ #entries: IRFunc[] = [];
8791
8892 constructor(
8993 source: string,
9094 bindings: Record<string, Binding>,
9195 grid: GridSizes,
9296 fileName = 'model.m',
97+ /** Shared .m files compiled alongside the model — the operator and solver
98+ * library. Only each file's namesake function is visible to the model,
99+ * as in MATLAB; a model function of the same name shadows it. */
100+ libs: { name: string; source: string }[] = [],
93101 ) {
94102 const ast = parseMFile(source, fileName);
95103 const ws = new Workspace(fileName, []);
96104 ws.addFile({ name: fileName, source, ast });
105+ for (const lib of libs) {
106+ ws.addFile({ name: lib.name, source: lib.source, ast: parseMFile(lib.source, lib.name) });
107+ }
97108 // synth / analys become resolvable, with their type rules.
98109 for (const f of externalOpFiles(grid)) ws.addFile(f);
99110 ws.finalize();
@@ -158,6 +169,7 @@ export class CompiledModel {
158169 nargout,
159170 undefined,
160171 );
172+ this.#entries.push(fn);
161173
162174 return {
163175 name,
@@ -180,13 +192,27 @@ export class CompiledModel {
180192 }
181193
182194 /**
183- * Run the inline pass over everything specialized so far. It rewrites the
195+ * Expand user-function calls, then run the inline pass. Both rewrite the
184196 * function bodies in place, so `CompiledFunction`s handed out earlier are
185197 * updated too.
198+ *
199+ * Expansion comes first: with every call spliced into its caller, each
200+ * entry function is one flat body, and the inline pass fuses it exactly as
201+ * it would the same code written out by hand — a call boundary neither
202+ * blocks fusion nor changes what compiles. Only the entry functions are
203+ * rewritten; the callee specializations they were cloned from are no longer
204+ * referenced.
186205 */
187206 finish(): void {
207+ for (const fn of this.#entries) {
208+ inFunction(fn.name, () =>
209+ expandUserCalls(fn, (cName) => this.#lowerer.specializations.get(cName)),
210+ );
211+ }
212+
188213 // Snapshot what each loop body assigns, before the pass can rewrite it.
189- const loops = [...this.#lowerer.specializations.values()].flatMap((fn) =>
214+ const entries = new Map(this.#entries.map((fn) => [fn.cName, fn]));
215+ const loops = [...entries.values()].flatMap((fn) =>
190216 forLoops(fn.body).map((loop) => ({
191217 fn,
192218 loop,
@@ -194,7 +220,7 @@ export class CompiledModel {
194220 })),
195221 );
196222
197- inlinePass({ topLevelStmts: [], functions: this.#lowerer.specializations });
223+ inlinePass({ topLevelStmts: [], functions: entries });
198224
199225 for (const { fn, loop, assignedBefore } of loops) checkLoopEscapes(fn, loop, assignedBefore);
200226 }
src/mgpu/externals.tsmodified+189−7View file
@@ -77,12 +77,178 @@ exports.cBody = function () {
7777 `;
7878 }
7979
80+/** Source for `dot`'s `.mtoc2.js`: two same-shape real arrays -> a 1x1
81+ * scalar. The result is GPU-resident (a 1-element buffer written by a
82+ * reduction dispatch, src/mgpu/reduce.ts), which is what lets a solver's
83+ * alpha/omega recurrences run without the CPU in the loop. */
84+function dotSource(): string {
85+ return `
86+exports.name = "dot";
87+
88+exports.transfer = function (argTypes, nargout) {
89+ if (argTypes.length !== 2) {
90+ throw new Error("dot takes exactly two arguments, got " + argTypes.length);
91+ }
92+ if (nargout > 1) {
93+ throw new Error("dot returns one value, but " + nargout + " were requested");
94+ }
95+ for (var i = 0; i < 2; i++) {
96+ var a = argTypes[i];
97+ if (!a || a.kind !== "Numeric" || a.isComplex) {
98+ throw new Error("dot requires real numeric arrays");
99+ }
100+ }
101+ var s0 = argTypes[0].shape;
102+ var s1 = argTypes[1].shape;
103+ if (!s0 || !s1 || s0.length !== s1.length ||
104+ !s0.every(function (d, i) { return d === s1[i]; })) {
105+ throw new Error(
106+ "dot requires two arrays of the same shape, got " +
107+ (s0 ? s0.join("x") : "unknown") + " and " + (s1 ? s1.join("x") : "unknown")
108+ );
109+ }
110+ return [${numericType(1, 1)}];
111+};
112+
113+// Never called: this project executes the IR on WebGPU and emits no C.
114+exports.emit = function () {
115+ throw new Error("dot: no C backend (this reduction runs on WebGPU)");
116+};
117+exports.cBody = function () {
118+ return "";
119+};
120+`;
121+}
122+
123+/**
124+ * Sources for the indexed-access ops a Krylov solver's bookkeeping needs:
125+ *
126+ * Vk = getslab(VB, k) k-th [2, nlm] field in a bank VB = [2, nlm*K]
127+ * VB = setslab(VB, W, k) the bank with field k replaced
128+ * h = getat(A, i) element of a small matrix (1- or 2-index,
129+ * h = getat(A, i, j) column-major)
130+ * A = setat(A, h, i) the matrix with that element replaced
131+ * A = setat(A, h, i, j)
132+ *
133+ * All four are functional updates at the MATLAB level — `A(i,j) = h` cannot
134+ * lower, because numbl's JIT must prove an indexed write in bounds at
135+ * lowering time and a loop variable has no value there. Written as calls,
136+ * lowering just types them; the *planner* resolves each index when the
137+ * unrolled loop makes it a literal, and compiles every one of these to a
138+ * static-offset buffer copy (src/mgpu/plan.ts) — an in-place write when the
139+ * result is assigned back over the base, as a solver's loop does.
140+ */
141+function indexOpFiles(nlm: number): { name: string; source: string }[] {
142+ const shared = `
143+function isRealNumeric(t) {
144+ return t && t.kind === "Numeric" && !t.isComplex;
145+}
146+function isScalarIndex(t) {
147+ return isRealNumeric(t) && (!t.shape || t.shape.every(function (d) { return d === 1; }));
148+}
149+/** The base's type, with any exact (constant-folded) value stripped — the
150+ * result differs from the base in one element, so it is not that constant. */
151+function baseType(t) {
152+ return {
153+ kind: "Numeric", elem: t.elem, isComplex: false,
154+ dims: t.dims, shape: t.shape, sign: "unknown",
155+ };
156+}
157+function checkBank(name, t) {
158+ if (!isRealNumeric(t) || !t.shape || t.shape.length !== 2 || t.shape[0] !== 2 ||
159+ t.shape[1] % ${nlm} !== 0) {
160+ throw new Error(
161+ name + " requires a 2 x (k*${nlm}) bank of spectral fields, got " +
162+ (t && t.shape ? t.shape.join("x") : "unknown shape")
163+ );
164+ }
165+}
166+`;
167+ const scalar11 = numericType(1, 1);
168+ const slab = numericType(2, nlm);
169+ return [
170+ {
171+ name: 'getslab.mtoc2.js',
172+ source: `${shared}
173+exports.name = "getslab";
174+exports.transfer = function (argTypes, nargout) {
175+ if (argTypes.length !== 2) throw new Error("getslab takes (bank, k), got " + argTypes.length + " arguments");
176+ if (nargout > 1) throw new Error("getslab returns one value");
177+ checkBank("getslab", argTypes[0]);
178+ if (!isScalarIndex(argTypes[1])) throw new Error("getslab's index must be a real scalar");
179+ return [${slab}];
180+};
181+exports.emit = function () { throw new Error("getslab: no C backend"); };
182+exports.cBody = function () { return ""; };
183+`,
184+ },
185+ {
186+ name: 'setslab.mtoc2.js',
187+ source: `${shared}
188+exports.name = "setslab";
189+exports.transfer = function (argTypes, nargout) {
190+ if (argTypes.length !== 3) throw new Error("setslab takes (bank, field, k), got " + argTypes.length + " arguments");
191+ if (nargout > 1) throw new Error("setslab returns one value");
192+ checkBank("setslab", argTypes[0]);
193+ var f = argTypes[1];
194+ if (!isRealNumeric(f) || !f.shape || f.shape.length !== 2 || f.shape[0] !== 2 || f.shape[1] !== ${nlm}) {
195+ throw new Error("setslab's field must be 2 x ${nlm}, got " + (f && f.shape ? f.shape.join("x") : "unknown shape"));
196+ }
197+ if (!isScalarIndex(argTypes[2])) throw new Error("setslab's index must be a real scalar");
198+ return [baseType(argTypes[0])];
199+};
200+exports.emit = function () { throw new Error("setslab: no C backend"); };
201+exports.cBody = function () { return ""; };
202+`,
203+ },
204+ {
205+ name: 'getat.mtoc2.js',
206+ source: `${shared}
207+exports.name = "getat";
208+exports.transfer = function (argTypes, nargout) {
209+ if (argTypes.length !== 2 && argTypes.length !== 3) {
210+ throw new Error("getat takes (A, i) or (A, i, j), got " + argTypes.length + " arguments");
211+ }
212+ if (nargout > 1) throw new Error("getat returns one value");
213+ if (!isRealNumeric(argTypes[0]) || !argTypes[0].shape) throw new Error("getat's base must be a real array of known shape");
214+ for (var k = 1; k < argTypes.length; k++) {
215+ if (!isScalarIndex(argTypes[k])) throw new Error("getat's indices must be real scalars");
216+ }
217+ return [${scalar11}];
218+};
219+exports.emit = function () { throw new Error("getat: no C backend"); };
220+exports.cBody = function () { return ""; };
221+`,
222+ },
223+ {
224+ name: 'setat.mtoc2.js',
225+ source: `${shared}
226+exports.name = "setat";
227+exports.transfer = function (argTypes, nargout) {
228+ if (argTypes.length !== 3 && argTypes.length !== 4) {
229+ throw new Error("setat takes (A, v, i) or (A, v, i, j), got " + argTypes.length + " arguments");
230+ }
231+ if (nargout > 1) throw new Error("setat returns one value");
232+ if (!isRealNumeric(argTypes[0]) || !argTypes[0].shape) throw new Error("setat's base must be a real array of known shape");
233+ if (!isScalarIndex(argTypes[1])) throw new Error("setat's value must be a real scalar");
234+ for (var k = 2; k < argTypes.length; k++) {
235+ if (!isScalarIndex(argTypes[k])) throw new Error("setat's indices must be real scalars");
236+ }
237+ return [baseType(argTypes[0])];
238+};
239+exports.emit = function () { throw new Error("setat: no C backend"); };
240+exports.cBody = function () { return ""; };
241+`,
242+ },
243+ ];
244+}
245+
80246 /**
81- * Workspace files that make `synth` / `analys` / `dtheta` / `dphi` resolvable
82- * during lowering. `dtheta` and `dphi` (the surface's first partial
83- * derivatives, coefficients -> grid — see src/sht/deriv.ts) have exactly
84- * `synth`'s shape rule: both take spectral coefficients and produce a grid
85- * field.
247+ * Workspace files that make `synth` / `analys` / `dtheta` / `dphi` / `dot`
248+ * (and the indexed-access ops above) resolvable during lowering. `dtheta`
249+ * and `dphi` (the surface's first partial derivatives, coefficients -> grid
250+ * — see src/sht/deriv.ts) have exactly `synth`'s shape rule: both take
251+ * spectral coefficients and produce a grid field.
86252 */
87253 export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
88254 return [
@@ -102,8 +268,24 @@ export function externalOpFiles(g: GridSizes): { name: string; source: string }[
102268 name: 'dphi.mtoc2.js',
103269 source: transformSource('dphi', 2, g.nlm, g.npts, 1),
104270 },
271+ {
272+ name: 'dot.mtoc2.js',
273+ source: dotSource(),
274+ },
275+ ...indexOpFiles(g.nlm),
105276 ];
106277 }
107278
108-/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
109-export const EXTERNAL_OPS = new Set(['synth', 'analys', 'dtheta', 'dphi']);
279+/** Names the WGSL backend must implement as GPU encodes rather than kernels,
280+ * each with the argument counts it accepts. */
281+export const EXTERNAL_OPS = new Map<string, { minArgs: number; maxArgs: number }>([
282+ ['synth', { minArgs: 1, maxArgs: 1 }],
283+ ['analys', { minArgs: 1, maxArgs: 1 }],
284+ ['dtheta', { minArgs: 1, maxArgs: 1 }],
285+ ['dphi', { minArgs: 1, maxArgs: 1 }],
286+ ['dot', { minArgs: 2, maxArgs: 2 }],
287+ ['getslab', { minArgs: 2, maxArgs: 2 }],
288+ ['setslab', { minArgs: 3, maxArgs: 3 }],
289+ ['getat', { minArgs: 2, maxArgs: 3 }],
290+ ['setat', { minArgs: 3, maxArgs: 4 }],
291+]);
src/mgpu/inlineCalls.tsadded+323−0View file
@@ -0,0 +1,323 @@
1+/**
2+ * Expand calls to user-defined MATLAB functions into the caller's body.
3+ *
4+ * numbl lowers `dL = dlap(X, ...)` to a single `Call` statement whose cName
5+ * names the callee's specialization — one IRFunc per distinct argument-type
6+ * signature, shared across call sites. The WGSL planner, though, executes a
7+ * flat statement list: every value is a buffer, every statement a dispatch,
8+ * and a call boundary has no runtime meaning on the GPU. So this pass gives a
9+ * call the only meaning it can have there: the callee's lowered body, spliced
10+ * in at the call site.
11+ *
12+ * Each call site gets its own clone of the callee body, with every
13+ * callee-local cName made unique to the site — so two calls to the same
14+ * function get separate buffers, while the *same* site re-planned per
15+ * unrolled loop iteration keeps reusing its buffers (the planner keys buffers
16+ * by cName, and the clone is made once, here, not per iteration).
17+ *
18+ * Arguments bind by renaming, not copying, wherever that is sound: a callee
19+ * parameter whose argument is a plain variable — after the inline pass they
20+ * all are, since it never folds an expression into a call argument — reads
21+ * the caller's buffer directly, unless the callee reassigns the parameter, in
22+ * which case a copy is materialized so the caller's value is not clobbered.
23+ * Outputs are the same in reverse: assignments to a callee output become
24+ * assignments to the caller's target variable, which is what makes a
25+ * loop-carried output (a solver iterating its result) work unchanged.
26+ *
27+ * The pass runs before numbl's inline (fusion) pass, which works per function
28+ * and treats a user call as an opaque producer. With every call already
29+ * spliced away, fusion sees one flat body and folds it exactly as it would
30+ * the same code written out by hand — a callee's final assignment can fuse
31+ * into its consumer, and the compiled plan is identical either way.
32+ */
33+import type {
34+ Assign,
35+ For,
36+ IRExpr,
37+ IRFunc,
38+ IRStmt,
39+ Span,
40+} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
41+import { ModelCompileError } from './errors.ts';
42+
43+/** A located compile failure. A span from a shared lib file still carries its
44+ * offsets; they are only ever mapped onto the model source for display, so a
45+ * failure inside a lib mislocates but still names the construct. */
46+const fail = (message: string, span: Span): ModelCompileError =>
47+ new ModelCompileError(message, { start: span.start, end: span.end });
48+
49+/** Looks up a callee's specialization by its mangled cName. */
50+export type ResolveFn = (cName: string) => IRFunc | undefined;
51+
52+/**
53+ * Expand every user-function call in `fn`, recursively, mutating `fn.body`.
54+ * Returns true if anything was expanded.
55+ */
56+export function expandUserCalls(fn: IRFunc, resolve: ResolveFn): boolean {
57+ const ctx = { resolve, site: 0, expanded: false };
58+ fn.body = expandBody(fn.body, ctx, [fn.cName]);
59+ return ctx.expanded;
60+}
61+
62+interface Ctx {
63+ resolve: ResolveFn;
64+ /** Call-site counter, for unique cNames. */
65+ site: number;
66+ expanded: boolean;
67+}
68+
69+function expandBody(stmts: IRStmt[], ctx: Ctx, stack: string[]): IRStmt[] {
70+ const out: IRStmt[] = [];
71+ for (const stmt of stmts) {
72+ if (stmt.kind === 'Assign' && stmt.expr.kind === 'Call') {
73+ const callee = ctx.resolve(stmt.expr.cName);
74+ if (callee) {
75+ out.push(
76+ ...expandCall(
77+ callee,
78+ stmt.expr.args,
79+ [{ name: stmt.name, cName: stmt.cName }],
80+ ctx,
81+ stack,
82+ stmt,
83+ ),
84+ );
85+ continue;
86+ }
87+ }
88+ if (stmt.kind === 'MultiAssignCall' && !stmt.isBuiltin) {
89+ const callee = ctx.resolve(stmt.cName);
90+ if (!callee) {
91+ throw fail(
92+ `call to '${stmt.name}' does not resolve to a function in this model`,
93+ stmt.span,
94+ );
95+ }
96+ out.push(
97+ ...expandCall(
98+ callee,
99+ stmt.args,
100+ stmt.outputs.map((o) => o.binding),
101+ ctx,
102+ stack,
103+ stmt,
104+ ),
105+ );
106+ continue;
107+ }
108+ if (stmt.kind === 'For') {
109+ stmt.body = expandBody(stmt.body, ctx, stack);
110+ out.push(stmt);
111+ continue;
112+ }
113+ out.push(stmt);
114+ }
115+ return out;
116+}
117+
118+/**
119+ * One call site: the callee's body, cloned and renamed into the caller's
120+ * namespace, preceded by whatever argument bindings need materializing.
121+ */
122+function expandCall(
123+ callee: IRFunc,
124+ args: IRExpr[],
125+ outs: ({ name: string; cName: string } | null)[],
126+ ctx: Ctx,
127+ stack: string[],
128+ at: IRStmt,
129+): IRStmt[] {
130+ if (stack.includes(callee.cName)) {
131+ throw fail(
132+ `'${callee.name}' calls itself (perhaps through another function); ` +
133+ `recursion cannot be compiled to a fixed sequence of GPU operations`,
134+ at.span,
135+ );
136+ }
137+ ctx.expanded = true;
138+ const site = ++ctx.site;
139+ /** Site-unique cName for a callee-local. */
140+ const local = (cName: string): string => `${callee.name}$${site}$${cName}`;
141+ /** Display name for a callee-local — shows up in buffer labels and in
142+ * describe(), so a solver's internals read as `richardson#1.X`. */
143+ const display = (name: string): string => `${callee.name}#${site}.${name}`;
144+
145+ const assigned = assignedCNames(callee.body);
146+ const rename = new Map<string, string>();
147+ const names = new Map<string, string>();
148+ const prelude: IRStmt[] = [];
149+
150+ // Outputs first: assignments to a callee output become assignments to the
151+ // caller's target. An ignored output (`~`, or trailing outputs the caller
152+ // did not ask for) stays a site-local.
153+ callee.cOutputs.forEach((c, j) => {
154+ const target = j < outs.length ? outs[j] : null;
155+ if (target) {
156+ rename.set(c, target.cName);
157+ names.set(c, target.name);
158+ } else {
159+ rename.set(c, local(c));
160+ names.set(c, display(callee.outputs[j]));
161+ }
162+ });
163+
164+ // Parameters: rename onto the argument where sound, else materialize.
165+ callee.cParams.forEach((p, i) => {
166+ const arg = args[i];
167+ if (arg === undefined) {
168+ throw fail(
169+ `'${callee.name}' takes ${callee.cParams.length} arguments, ` +
170+ `but this call passes ${args.length}`,
171+ at.span,
172+ );
173+ }
174+ if (rename.has(p)) {
175+ // The parameter is also an output (`function X = f(X)`): seed the
176+ // caller's target with the argument, and let the body update it there.
177+ prelude.push(makeAssign(names.get(p)!, rename.get(p)!, arg, callee.paramTypes[i], at));
178+ } else if (arg.kind === 'Var' && !assigned.has(p)) {
179+ rename.set(p, arg.cName);
180+ names.set(p, arg.name);
181+ } else {
182+ // A non-variable argument, or a parameter the callee reassigns: bind it
183+ // to a site-local first. For a scalar this is a free derived-scalar
184+ // binding; for a tensor it is one copy kernel.
185+ const c = local(p);
186+ rename.set(p, c);
187+ names.set(p, display(callee.params[i]));
188+ prelude.push(makeAssign(names.get(p)!, c, arg, callee.paramTypes[i], at));
189+ }
190+ });
191+
192+ const body = cloneBody(callee.body, { rename, names, local, display, callee, at });
193+ return [...prelude, ...expandBody(body, ctx, [...stack, callee.cName])];
194+}
195+
196+const makeAssign = (
197+ name: string,
198+ cName: string,
199+ expr: IRExpr,
200+ ty: IRFunc['paramTypes'][number],
201+ at: IRStmt,
202+): Assign => ({ kind: 'Assign', name, cName, ty, expr, span: at.span });
203+
204+interface CloneCtx {
205+ /** Callee cName -> caller-namespace cName. Filled for outputs and params up
206+ * front; locals are added on first sight. */
207+ rename: Map<string, string>;
208+ names: Map<string, string>;
209+ local: (cName: string) => string;
210+ display: (name: string) => string;
211+ callee: IRFunc;
212+ at: IRStmt;
213+}
214+
215+function cloneBody(stmts: IRStmt[], c: CloneCtx): IRStmt[] {
216+ const out: IRStmt[] = [];
217+ for (const s of stmts) {
218+ switch (s.kind) {
219+ case 'Assign':
220+ out.push({
221+ ...s,
222+ name: mapName(s.name, s.cName, c),
223+ cName: mapCName(s.cName, s.name, c),
224+ expr: cloneExpr(s.expr, c),
225+ });
226+ break;
227+ case 'MultiAssignCall':
228+ out.push({
229+ ...s,
230+ args: s.args.map((a) => cloneExpr(a, c)),
231+ outputs: s.outputs.map((o) =>
232+ o.binding
233+ ? {
234+ ...o,
235+ binding: {
236+ name: mapName(o.binding.name, o.binding.cName, c),
237+ cName: mapCName(o.binding.cName, o.binding.name, c),
238+ },
239+ }
240+ : o,
241+ ),
242+ });
243+ break;
244+ case 'For': {
245+ const loop: For = {
246+ ...s,
247+ cVar: mapCName(s.cVar, s.varName, c),
248+ start: cloneExpr(s.start, c),
249+ end: cloneExpr(s.end, c),
250+ body: [],
251+ };
252+ loop.body = cloneBody(s.body, c);
253+ out.push(loop);
254+ break;
255+ }
256+ case 'ReturnFromFunction':
257+ // The callee's return has no meaning at the splice point; statements
258+ // never follow it (numbl refuses an early return at lowering).
259+ break;
260+ default:
261+ throw fail(
262+ `'${c.callee.name}' contains a '${s.kind}' statement, which cannot ` +
263+ `be compiled to a fixed sequence of GPU operations`,
264+ s.span,
265+ );
266+ }
267+ }
268+ return out;
269+}
270+
271+/** Caller-namespace cName for a callee-side cName, minting one for a local
272+ * seen for the first time. */
273+function mapCName(cName: string, name: string, c: CloneCtx): string {
274+ const existing = c.rename.get(cName);
275+ if (existing) return existing;
276+ const fresh = c.local(cName);
277+ c.rename.set(cName, fresh);
278+ c.names.set(cName, c.display(name));
279+ return fresh;
280+}
281+
282+function mapName(name: string, cName: string, c: CloneCtx): string {
283+ return c.names.get(cName) ?? c.display(name);
284+}
285+
286+function cloneExpr(e: IRExpr, c: CloneCtx): IRExpr {
287+ switch (e.kind) {
288+ case 'Var':
289+ return {
290+ ...e,
291+ name: mapName(e.name, e.cName, c),
292+ cName: mapCName(e.cName, e.name, c),
293+ };
294+ case 'Binary':
295+ return { ...e, left: cloneExpr(e.left, c), right: cloneExpr(e.right, c) };
296+ case 'Unary':
297+ return { ...e, operand: cloneExpr(e.operand, c) };
298+ case 'Call':
299+ return { ...e, args: e.args.map((a) => cloneExpr(a, c)) };
300+ default:
301+ // Literals and anything else without variable reads: share as-is (the
302+ // planner treats expressions as read-only).
303+ return e;
304+ }
305+}
306+
307+/** cNames assigned anywhere in a statement list, including loop variables. */
308+function assignedCNames(stmts: IRStmt[]): Set<string> {
309+ const out = new Set<string>();
310+ const walk = (list: IRStmt[]): void => {
311+ for (const s of list) {
312+ if (s.kind === 'Assign') out.add(s.cName);
313+ else if (s.kind === 'MultiAssignCall') {
314+ for (const o of s.outputs) if (o.binding) out.add(o.binding.cName);
315+ } else if (s.kind === 'For') {
316+ out.add(s.cVar);
317+ walk(s.body);
318+ }
319+ }
320+ };
321+ walk(stmts);
322+ return out;
323+}
src/mgpu/libs.tsadded+70−0View file
@@ -0,0 +1,70 @@
1+/**
2+ * Shared .m files every model compiles against: the operator library and the
3+ * solvers. A model calls these by name (`richardson(...)`, `dlap(...)`) the
4+ * way it calls `synth` — except these are ordinary MATLAB, compiled through
5+ * the same pipeline and expanded into the caller at compile time
6+ * (src/mgpu/inlineCalls.ts), so a call costs exactly what writing the body
7+ * inline would.
8+ *
9+ * MATLAB file-visibility rules apply: only a file's namesake function is
10+ * callable from other files, and a function defined in the model shadows a
11+ * lib of the same name.
12+ */
13+import dlapSource from '../../lib/dlap.m?raw';
14+import richardsonSource from '../../solvers/richardson.m?raw';
15+import bicgstabSource from '../../solvers/bicgstab.m?raw';
16+import gmresSource from '../../solvers/gmres.m?raw';
17+
18+export interface LibFile {
19+ name: string;
20+ source: string;
21+}
22+
23+/** The operator: dlap = lap_g - lap_s applied to a spectral field. */
24+export const operatorLibs: LibFile[] = [{ name: 'dlap.m', source: dlapSource }];
25+
26+/** The solvers for (I - dtD*lap_g) X = B. All share dlap's operator. */
27+export const solverLibs: LibFile[] = [
28+ { name: 'richardson.m', source: richardsonSource },
29+ { name: 'bicgstab.m', source: bicgstabSource },
30+ { name: 'gmres.m', source: gmresSource },
31+];
32+
33+/** Everything a model may call. */
34+export const modelLibs: LibFile[] = [...operatorLibs, ...solverLibs];
35+
36+/** Where each shared file lives on disk, for display. */
37+export const libPath = (name: string): string =>
38+ operatorLibs.some((f) => f.name === name) ? `lib/${name}` : `solvers/${name}`;
39+
40+export type SolverKey = 'richardson' | 'bicgstab' | 'gmres';
41+export const solverKeys: SolverKey[] = ['richardson', 'bicgstab', 'gmres'];
42+export const DEFAULT_SOLVER: SolverKey = 'richardson';
43+
44+/** What each solver actually takes: richardson needs no inner products, so
45+ * no weights; only gmres sizes a basis bank, so only it takes nlm. */
46+const SOLVE_FORWARD: Record<SolverKey, string> = {
47+ richardson: 'richardson(B, dtD, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)',
48+ bicgstab: 'bicgstab(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)',
49+ gmres: 'gmres(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter)',
50+};
51+
52+/**
53+ * The one-line shim behind the models' `solve(...)` call. The models pass
54+ * every argument any solver could want, and this host-generated file
55+ * forwards to the chosen one — so which solver runs is a compile-time choice
56+ * the app's solver control makes (swapping recompiles, like changing niter),
57+ * while the models stay identical across solvers. A model may also bypass
58+ * the shim and call a solver by name.
59+ */
60+export function solveShim(solver: SolverKey): LibFile {
61+ return {
62+ name: 'solve.m',
63+ source:
64+ `% Host-generated: forwards the models' solve(...) call to the solver\n` +
65+ `% selected in the app. See solvers/${solver}.m.\n` +
66+ `function X = solve(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter)\n` +
67+ ` X = ${SOLVE_FORWARD[solver]};\n` +
68+ `end\n`,
69+ };
70+}
src/mgpu/model.tsmodified+36−1View file
@@ -23,6 +23,7 @@ import { lmIndex, type ShtConfig } from '../sht/layout.ts';
2323 import { HostBuffers, ModelPlan } from './plan.ts';
2424 import { inFunction, inFunctionAsync, inModel } from './errors.ts';
2525 import { CompiledModel, type Binding } from './compile.ts';
26+import { DEFAULT_SOLVER, modelLibs, solveShim, type LibFile } from './libs.ts';
2627
2728 export interface ModelParams {
2829 [key: string]: number;
@@ -56,6 +57,13 @@ export interface GpuModelOptions {
5657 * the count is part of what compiles and changing it recompiles.
5758 */
5859 niter?: number;
60+ /**
61+ * The shared .m files compiled alongside the model. Defaults to the shipped
62+ * operator and solvers plus a `solve(...)` shim forwarding to richardson;
63+ * the session passes its own list to honor the app's solver selection and
64+ * any lib edits.
65+ */
66+ libs?: LibFile[];
5967 }
6068
6169 /** Host-supplied surface fields, in the layout the .m sees them. */
@@ -118,6 +126,27 @@ export function filterMask(cfg: ShtConfig, nlm: number): Float32Array {
118126 return filt;
119127 }
120128
129+/**
130+ * Inner-product weights for the half-spectrum layout: 1 at m = 0, 2 at
131+ * m > 0, duplicated across re/im like `lam`. The 2 x nlm state stores only
132+ * m >= 0 of a real field (the m < 0 coefficients are conjugates), so the
133+ * true L2 inner product on the sphere is sum(wlm .* a .* b) — which is what
134+ * a Krylov solver's `dot` calls should compute if its scalars are to mean
135+ * what they mean in the analysis.
136+ */
137+export function weightMask(cfg: ShtConfig, nlm: number): Float32Array {
138+ const wlm = new Float32Array(2 * nlm);
139+ for (let m = 0; m <= cfg.mmax; m++) {
140+ for (let l = m; l <= cfg.lmax; l++) {
141+ const i = lmIndex(cfg.lmax, l, m);
142+ const w = m === 0 ? 1 : 2;
143+ wlm[2 * i] = w;
144+ wlm[2 * i + 1] = w;
145+ }
146+ }
147+ return wlm;
148+}
149+
121150 export class GpuModel {
122151 readonly paramNames: string[];
123152 readonly state: string[];
@@ -175,6 +204,7 @@ export class GpuModel {
175204 const bindings: Record<string, Binding> = {
176205 lam: { kind: 'tensor', shape: [2, nlm] },
177206 filt: { kind: 'tensor', shape: [2, nlm] },
207+ wlm: { kind: 'tensor', shape: [2, nlm] },
178208 noise: { kind: 'tensor', shape: [npts, 1] },
179209 npts: { kind: 'const', value: npts },
180210 nlm: { kind: 'const', value: nlm },
@@ -189,7 +219,10 @@ export class GpuModel {
189219 for (const p of paramNames) bindings[p] = { kind: 'param' };
190220
191221 // Parsing belongs to the file, not to either function.
192- const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
222+ const libs = opts.libs ?? [...modelLibs, solveShim(DEFAULT_SOLVER)];
223+ const compiled = inModel(
224+ () => new CompiledModel(source, bindings, { npts, nlm }, 'model.m', libs),
225+ );
193226 // Both functions return the new state first, then the rendered grid fields.
194227 const nargout = state.length + view.length;
195228 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
@@ -207,6 +240,7 @@ export class GpuModel {
207240 for (const s of state) host.ensure(s, 2 * nlm);
208241 host.ensure('lam', 2 * nlm);
209242 host.ensure('filt', 2 * nlm);
243+ host.ensure('wlm', 2 * nlm);
210244 host.ensure('noise', npts);
211245 if (geometry) {
212246 for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
@@ -223,6 +257,7 @@ export class GpuModel {
223257
224258 host.upload('lam', eigenvalues(cfg, nlm));
225259 host.upload('filt', filterMask(cfg, nlm));
260+ host.upload('wlm', weightMask(cfg, nlm));
226261 if (geometry) {
227262 host.upload('gx', geometry.x);
228263 host.upload('gy', geometry.y);
src/mgpu/numbl.d.tsmodified+23−2View file
@@ -137,16 +137,37 @@ declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
137137 body: IRStmt[];
138138 span: Span;
139139 }
140+ /**
141+ * A multi-output call statement, `[a, b] = f(...)`. For a user function
142+ * (`isBuiltin` false or absent) `cName` is the callee's specialization and
143+ * the call is expanded into the caller (src/mgpu/inlineCalls.ts); a
144+ * multi-output builtin is rejected by the planner.
145+ */
146+ export interface MultiAssignCall {
147+ kind: 'MultiAssignCall';
148+ /** Mangled specialization cName (user function) or builtin name. */
149+ cName: string;
150+ /** Source-level callee name, for diagnostics. */
151+ name: string;
152+ isBuiltin?: boolean;
153+ args: IRExpr[];
154+ /** One entry per output slot; `binding` null for an ignored output. */
155+ outputs: ReadonlyArray<{
156+ ty: Type;
157+ binding: { name: string; cName: string } | null;
158+ }>;
159+ span: Span;
160+ }
140161 /** Any other IR statement kind — rejected by the planner. */
141162 export interface OtherStmt {
142163 kind:
143164 | 'ExprStmt' | 'If' | 'While' | 'ReturnFromFunction' | 'Break'
144- | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
165+ | 'Continue' | 'TypeComment' | 'MemberStore'
145166 | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
146167 span: Span;
147168 }
148169
149- export type IRStmt = Assign | For | OtherStmt;
170+ export type IRStmt = Assign | For | MultiAssignCall | OtherStmt;
150171
151172 export interface IRFunc {
152173 name: string;
src/mgpu/plan.tsmodified+305−43View file
@@ -13,6 +13,7 @@ import type { Assign, For, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lower
1313 import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
1414 import { ShtPlan, type ShtBinding } from '../sht/sht.ts';
1515 import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
16+import { ReducePlan, type DotBinding } from './reduce.ts';
1617 import type { CompiledFunction } from './compile.ts';
1718 import { EXTERNAL_OPS } from './externals.ts';
1819 import {
@@ -26,15 +27,16 @@ const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
2627 const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
2728 const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
2829
29-/**
30- * The compile-time value of a scalar expression, if it has one. A literal
31- * carries its own; a variable carries one when it was bound to a `const` (the
32- * host's fixed scalars) or computed from constants, because numbl propagates
33- * `exact` through the type lattice.
34- */
35-const exactValue = (e: IRExpr): number | undefined => {
36- if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
37- return e.kind === 'NumLit' ? e.value : undefined;
30+/** Scalar arithmetic a plan-time evaluator can fold. */
31+const PLAN_BINOPS: Record<string, (l: number, r: number) => number> = {
32+ plus: (l, r) => l + r,
33+ minus: (l, r) => l - r,
34+ times: (l, r) => l * r,
35+ mtimes: (l, r) => l * r,
36+ rdivide: (l, r) => l / r,
37+ mrdivide: (l, r) => l / r,
38+ power: (l, r) => Math.pow(l, r),
39+ mpower: (l, r) => Math.pow(l, r),
3840 };
3941
4042 /** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
@@ -121,7 +123,17 @@ type Op =
121123 }
122124 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
123125 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
124- | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
126+ | { kind: 'dot'; binding: DotBinding; label: string }
127+ | {
128+ kind: 'copy';
129+ from: GPUBuffer;
130+ to: GPUBuffer;
131+ bytes: number;
132+ label: string;
133+ /** Byte offsets, for the indexed-access ops. Absent means 0. */
134+ fromOffset?: number;
135+ toOffset?: number;
136+ };
125137
126138 export interface PlanSpec {
127139 /** The specialized function this plan executes. */
@@ -271,6 +283,74 @@ export class ModelPlan {
271283 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
272284 });
273285
286+ /** Built on first use — only a model that calls `dot` pays for it. */
287+ let reduce: ReducePlan | null = null;
288+
289+ /** Does this expression read any GPU-resident value? Decides whether a
290+ * scalar assignment can stay a compile-time derived scalar or needs a
291+ * 1-element kernel. Plan-order matters and is correct: a name is
292+ * buffer-backed from the statement that first computes it into one. */
293+ const readsBufferValue = (e: IRExpr): boolean => {
294+ let found = false;
295+ collectVars(e, (v) => {
296+ if (slots.has(v.cName)) found = true;
297+ });
298+ return found;
299+ };
300+
301+ /**
302+ * The value a scalar expression has *at this point in the plan*, if it
303+ * is decidable. A literal carries its own; a variable carries one via
304+ * numbl's `exact` lattice or — the case the lattice cannot see — via its
305+ * derived-scalar binding, which is how an unrolled loop's variable (and
306+ * anything computed from it, like an index or an inner loop bound)
307+ * resolves to that iteration's literal. A buffer-backed name is a
308+ * runtime value and never resolves.
309+ */
310+ const planTimeValue = (e: IRExpr): number | undefined => {
311+ if (e.kind === 'NumLit') return e.value;
312+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
313+ switch (e.kind) {
314+ case 'Var': {
315+ if (slots.has(e.cName)) return undefined;
316+ const d = derivedScalars.get(e.cName);
317+ return d ? planTimeValue(d.expr) : undefined;
318+ }
319+ case 'Binary': {
320+ const op = PLAN_BINOPS[e.builtin];
321+ if (!op) return undefined;
322+ const l = planTimeValue(e.left);
323+ const r = planTimeValue(e.right);
324+ return l === undefined || r === undefined ? undefined : op(l, r);
325+ }
326+ case 'Unary': {
327+ const v = planTimeValue(e.operand);
328+ if (v === undefined) return undefined;
329+ if (e.builtin === 'uminus') return -v;
330+ if (e.builtin === 'uplus') return v;
331+ return undefined;
332+ }
333+ default:
334+ return undefined;
335+ }
336+ };
337+
338+ /** A plan-time index: integral and 1-based. */
339+ const planTimeIndex = (e: IRExpr, what: string, span: unknown): number => {
340+ const v = planTimeValue(e);
341+ if (v === undefined) {
342+ throw new UnsupportedOnGpu(
343+ `${what} must be known when the model compiles — a literal, a fixed ` +
344+ `argument, or a value of the unrolled loop's variable`,
345+ span,
346+ );
347+ }
348+ if (!Number.isInteger(v) || v < 1) {
349+ throw new UnsupportedOnGpu(`${what} must be a positive integer (got ${v})`, span);
350+ }
351+ return v;
352+ };
353+
274354 const ops: Op[] = [];
275355 for (const stmt of fn.body) {
276356 await planStatement(stmt);
@@ -323,10 +403,14 @@ export class ModelPlan {
323403 stmt.span,
324404 );
325405 }
326- if (!isTensor(stmt.ty)) {
406+ const ext = externalCall(stmt);
407+ if (!isTensor(stmt.ty) && !ext && !readsBufferValue(stmt.expr)) {
327408 // A scalar the model derives from its parameters (`us = a + b`). It
328409 // gets no buffer and no dispatch: the kernels that read it bind it as
329- // a `let` in their prologue.
410+ // a `let` in their prologue. A scalar computed from GPU-resident
411+ // values (a `dot` result, or anything downstream of one) instead
412+ // falls through to a 1-element kernel, because its inputs live in
413+ // buffers the CPU never sees.
330414 derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
331415 return;
332416 }
@@ -346,26 +430,175 @@ export class ModelPlan {
346430 }
347431 byName.set(stmt.name, dest);
348432
349- const ext = externalCall(stmt);
350433 if (ext) {
351- const argSlot = slots.get(ext.argCName);
352- if (!argSlot) {
353- throw new UnsupportedOnGpu(
354- `'${ext.name}' reads '${ext.argName}', which has no buffer`,
355- stmt.span,
434+ // Lazy per-argument resolution: buffer arguments must have slots,
435+ // while index arguments are plan-time scalars with no buffer at all.
436+ const argSlot = (i: number): Slot => {
437+ const a = ext.args[i];
438+ if (a.kind !== 'Var') {
439+ throw new UnsupportedOnGpu(
440+ `'${ext.name}' needs a plain variable here — assign the ` +
441+ `expression to a variable first`,
442+ stmt.span,
443+ );
444+ }
445+ const s = slots.get(a.cName);
446+ if (!s) {
447+ throw new UnsupportedOnGpu(
448+ `'${ext.name}' reads '${a.name}', which has no buffer`,
449+ stmt.span,
450+ );
451+ }
452+ return s;
453+ };
454+ const label = `${stmt.name} = ${ext.name}(${ext.args.map(extArgName).join(', ')})`;
455+ if (ext.name === 'dot') {
456+ const a = argSlot(0);
457+ const b = argSlot(1);
458+ if (a.count !== b.count) {
459+ throw new UnsupportedOnGpu(
460+ `'dot' needs equal-length arguments (${a.count} vs ${b.count})`,
461+ stmt.span,
462+ );
463+ }
464+ if (a.buffer === dest.buffer || b.buffer === dest.buffer) {
465+ throw new UnsupportedOnGpu(
466+ `'dot' cannot write over one of its own arguments`,
467+ stmt.span,
468+ );
469+ }
470+ reduce ??= new ReducePlan(device);
471+ ops.push({
472+ kind: 'dot',
473+ binding: await reduce.createDotBinding(a.buffer, b.buffer, dest.buffer, a.count),
474+ label,
475+ });
476+ return;
477+ }
478+ if (ext.name === 'getslab' || ext.name === 'setslab') {
479+ const slabElems = 2 * sht.nlm;
480+ const bank = argSlot(0);
481+ const nslabs = Math.floor(bank.count / slabElems);
482+ const kArg = ext.args[ext.name === 'getslab' ? 1 : 2];
483+ const k = planTimeIndex(kArg, `'${ext.name}'s index '${extArgName(kArg)}'`, stmt.span);
484+ if (bank.count % slabElems !== 0 || k > nslabs) {
485+ throw new UnsupportedOnGpu(
486+ `'${ext.name}': slab ${k} is out of range for a bank of ` +
487+ `${nslabs} spectral fields`,
488+ stmt.span,
489+ );
490+ }
491+ const slabBytes = 4 * slabElems;
492+ if (ext.name === 'getslab') {
493+ if (dest.count !== slabElems || bank.buffer === dest.buffer) {
494+ throw new UnsupportedOnGpu(`'getslab' cannot read into its own bank`, stmt.span);
495+ }
496+ ops.push({
497+ kind: 'copy', from: bank.buffer, fromOffset: (k - 1) * slabBytes,
498+ to: dest.buffer, bytes: slabBytes, label,
499+ });
500+ } else {
501+ const field = argSlot(1);
502+ if (field.count !== slabElems || field.buffer === dest.buffer) {
503+ throw new UnsupportedOnGpu(
504+ `'setslab' needs a distinct 2 x nlm field to write`,
505+ stmt.span,
506+ );
507+ }
508+ // Functional update: writing back over the base is the in-place
509+ // fast path; a fresh destination first takes a copy of the bank.
510+ if (dest.buffer !== bank.buffer) {
511+ ops.push({
512+ kind: 'copy', from: bank.buffer, to: dest.buffer,
513+ bytes: 4 * bank.count, label: `${label} (bank copy)`,
514+ });
515+ }
516+ ops.push({
517+ kind: 'copy', from: field.buffer,
518+ to: dest.buffer, toOffset: (k - 1) * slabBytes,
519+ bytes: slabBytes, label,
520+ });
521+ }
522+ return;
523+ }
524+ if (ext.name === 'getat' || ext.name === 'setat') {
525+ const base = argSlot(0);
526+ const baseTy = ext.args[0].ty;
527+ if (ext.args[0].kind !== 'Var') {
528+ throw new UnsupportedOnGpu(`'${ext.name}' needs a variable base`, stmt.span);
529+ }
530+ const shape = isNumeric(baseTy) ? baseTy.shape : undefined;
531+ if (!shape) {
532+ throw new UnsupportedOnGpu(`'${ext.name}' needs a base of known shape`, stmt.span);
533+ }
534+ const idxArgs = ext.args.slice(ext.name === 'getat' ? 1 : 2);
535+ const idx = idxArgs.map(
536+ (a) => planTimeIndex(a, `'${ext.name}'s index '${extArgName(a)}'`, stmt.span) - 1,
356537 );
538+ // Column-major, like everything else in the 2 x nlm layout: a
539+ // 2-index access is (i-1) + (j-1)*rows, a 1-index access is linear.
540+ let offset: number;
541+ if (idx.length === 2) {
542+ const [i, j] = idx;
543+ if (i >= shape[0] || j >= (shape[1] ?? 1)) {
544+ throw new UnsupportedOnGpu(
545+ `'${ext.name}': (${i + 1}, ${j + 1}) is outside ` +
546+ `${shape.join('x')} '${extArgName(ext.args[0])}'`,
547+ stmt.span,
548+ );
549+ }
550+ offset = i + j * shape[0];
551+ } else {
552+ offset = idx[0];
553+ if (offset >= base.count) {
554+ throw new UnsupportedOnGpu(
555+ `'${ext.name}': index ${offset + 1} is outside ` +
556+ `${base.count}-element '${extArgName(ext.args[0])}'`,
557+ stmt.span,
558+ );
559+ }
560+ }
561+ if (ext.name === 'getat') {
562+ if (dest.count !== 1 || base.buffer === dest.buffer) {
563+ throw new UnsupportedOnGpu(`'getat' cannot read into its own base`, stmt.span);
564+ }
565+ ops.push({
566+ kind: 'copy', from: base.buffer, fromOffset: 4 * offset,
567+ to: dest.buffer, bytes: 4, label,
568+ });
569+ } else {
570+ const value = argSlot(1);
571+ if (value.count !== 1 || value.buffer === dest.buffer) {
572+ throw new UnsupportedOnGpu(
573+ `'setat' needs a distinct 1-element value to write — compute ` +
574+ `it into a variable first`,
575+ stmt.span,
576+ );
577+ }
578+ if (dest.buffer !== base.buffer) {
579+ ops.push({
580+ kind: 'copy', from: base.buffer, to: dest.buffer,
581+ bytes: 4 * base.count, label: `${label} (base copy)`,
582+ });
583+ }
584+ ops.push({
585+ kind: 'copy', from: value.buffer,
586+ to: dest.buffer, toOffset: 4 * offset, bytes: 4, label,
587+ });
588+ }
589+ return;
357590 }
358- const label = `${stmt.name} = ${ext.name}(${ext.argName})`;
591+ const src = argSlot(0);
359592 if (ext.name === 'synth') {
360593 ops.push({
361594 kind: 'synth',
362- binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
595+ binding: sht.createSynthBinding(src.buffer, dest.buffer),
363596 label,
364597 });
365598 } else if (ext.name === 'analys') {
366599 ops.push({
367600 kind: 'analys',
368- binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
601+ binding: sht.createAnalysBinding(src.buffer, dest.buffer),
369602 label,
370603 });
371604 } else if (ext.name === 'dtheta' || ext.name === 'dphi') {
@@ -378,8 +611,8 @@ export class ModelPlan {
378611 }
379612 ops.push(
380613 ext.name === 'dtheta'
381- ? { kind: 'dtheta', binding: deriv.createDthetaBinding(argSlot.buffer, dest.buffer), label }
382- : { kind: 'dphi', binding: deriv.createDphiBinding(argSlot.buffer, dest.buffer), label },
614+ ? { kind: 'dtheta', binding: deriv.createDthetaBinding(src.buffer, dest.buffer), label }
615+ : { kind: 'dphi', binding: deriv.createDphiBinding(src.buffer, dest.buffer), label },
383616 );
384617 } else {
385618 throw new UnsupportedOnGpu(`unknown external op '${ext.name}'`, stmt.span);
@@ -387,11 +620,15 @@ export class ModelPlan {
387620 return;
388621 }
389622
390- // Element-wise kernel. Collect the distinct tensor operands and give
391- // them dense binding slots.
623+ // Element-wise kernel. Collect the distinct buffer-backed operands —
624+ // multi-element tensors, plus any single-element value living in a
625+ // buffer (a dot result or a scalar computed from one) — and give them
626+ // dense binding slots. The kernel reads a single-element operand as
627+ // `in<slot>[0]`, which is what broadcasts it across the output.
392628 const tensors = new Map<string, number>();
393- collectTensorVars(stmt.expr, (cName) => {
394- if (!tensors.has(cName)) tensors.set(cName, tensors.size);
629+ collectVars(stmt.expr, (v) => {
630+ if (!isTensor(v.ty) && !slots.has(v.cName)) return;
631+ if (!tensors.has(v.cName)) tensors.set(v.cName, tensors.size);
395632 });
396633
397634 const label = `${stmt.name} = <${count} elements, element-wise>`;
@@ -462,15 +699,16 @@ export class ModelPlan {
462699 * iteration's body is planned and its WGSL emitted.
463700 */
464701 async function planFor(stmt: For): Promise<void> {
465- const from = exactValue(stmt.start);
466- const to = exactValue(stmt.end);
702+ const from = planTimeValue(stmt.start);
703+ const to = planTimeValue(stmt.end);
467704 if (from === undefined || to === undefined) {
468705 throw new UnsupportedOnGpu(
469706 `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
470707 `be known when the model is compiled — ` +
471708 `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
472- `runtime value. Use a whole number, or a count the app supplies ` +
473- `as a fixed argument (changing it recompiles).`,
709+ `runtime value. Use a whole number, a count the app supplies ` +
710+ `as a fixed argument (changing it recompiles), or an enclosing ` +
711+ `unrolled loop's variable.`,
474712 stmt.span,
475713 );
476714 }
@@ -572,9 +810,18 @@ export class ModelPlan {
572810 case 'dphi':
573811 this.#derivInto(inPass(), op);
574812 break;
813+ case 'dot': {
814+ const p = inPass();
815+ p.setPipeline(op.binding.pipeline);
816+ p.setBindGroup(0, op.binding.bindGroup);
817+ p.dispatchWorkgroups(1);
818+ break;
819+ }
575820 case 'copy':
576821 endPass();
577- encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
822+ encoder.copyBufferToBuffer(
823+ op.from, op.fromOffset ?? 0, op.to, op.toOffset ?? 0, op.bytes,
824+ );
578825 break;
579826 }
580827 }
@@ -606,27 +853,42 @@ export class ModelPlan {
606853 }
607854 }
608855
609-/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
610-function externalCall(
611- stmt: Assign,
612-): { name: string; argCName: string; argName: string } | null {
856+/**
857+ * `x = synth(y)` / `x = dot(y, z)` -> the call's name and arguments. A
858+ * buffer argument must be a plain variable (an expression would need its own
859+ * buffer, which is exactly what writing it on its own line provides — the
860+ * per-argument check is in the planner); an index argument may be any
861+ * expression the plan can evaluate (`j + 1`).
862+ */
863+function externalCall(stmt: Assign): { name: string; args: IRExpr[] } | null {
613864 const e = stmt.expr;
614- if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
615- if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
865+ if (e.kind !== 'Call') return null;
866+ const arity = EXTERNAL_OPS.get(e.name);
867+ if (!arity) return null;
868+ if (e.args.length < arity.minArgs || e.args.length > arity.maxArgs) {
869+ const want =
870+ arity.minArgs === arity.maxArgs
871+ ? `${arity.minArgs}`
872+ : `${arity.minArgs} to ${arity.maxArgs}`;
616873 throw new UnsupportedOnGpu(
617- `'${e.name}' must be applied to a single variable`,
874+ `'${e.name}' takes ${want} argument${arity.maxArgs === 1 ? '' : 's'}`,
618875 stmt.span,
619876 );
620877 }
621- const arg = e.args[0];
622- return { name: e.name, argCName: arg.cName, argName: arg.name };
878+ return { name: e.name, args: e.args };
623879 }
624880
625-function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
881+const extArgName = (a: IRExpr): string =>
882+ a.kind === 'Var' ? a.name : a.kind === 'NumLit' ? String(a.value) : '<expression>';
883+
884+function collectVars(
885+ e: IRExpr,
886+ visit: (v: Extract<IRExpr, { kind: 'Var' }>) => void,
887+): void {
626888 const walk = (x: IRExpr): void => {
627889 switch (x.kind) {
628890 case 'Var':
629- if (isTensor(x.ty)) visit(x.cName);
891+ visit(x);
630892 return;
631893 case 'Binary':
632894 walk(x.left);
src/mgpu/reduce.tsadded+130−0View file
@@ -0,0 +1,130 @@
1+/**
2+ * `dot(x, y)` — the one reduction the solvers need, as a GPU operation.
3+ *
4+ * A Krylov solver's scalars (rho, alpha, omega, ...) are inner products of the
5+ * spectral state, and the whole point of the plan architecture is that no
6+ * value crosses back to the CPU mid-step — so the dot product must produce a
7+ * GPU-resident scalar: a 1-element buffer that later kernels read (the
8+ * planner binds any single-element value as `in<slot>[0]`, see
9+ * src/mgpu/wgsl.ts).
10+ *
11+ * One workgroup does the whole reduction: each thread accumulates a strided
12+ * partial sum, then a shared-memory tree combines them and thread 0 writes
13+ * the result. A single dispatch, no multi-pass bookkeeping, and — because
14+ * the striding is fixed — a bit-deterministic summation order on a given
15+ * device. n up to a few hundred thousand is a short loop per thread, far
16+ * below anything the solver grids produce.
17+ */
18+
19+const WG = 256;
20+
21+/** One dot call site: its bind group and the pipeline for its length. */
22+export interface DotBinding {
23+ readonly pipeline: GPUComputePipeline;
24+ readonly bindGroup: GPUBindGroup;
25+}
26+
27+const dotWGSL = (n: number): string => `
28+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
29+@group(0) @binding(1) var<storage, read> a: array<f32>;
30+@group(0) @binding(2) var<storage, read> b: array<f32>;
31+
32+var<workgroup> partials: array<f32, ${WG}>;
33+
34+@compute @workgroup_size(${WG})
35+fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
36+ var s = 0.0;
37+ var i = lid.x;
38+ loop {
39+ if (i >= ${n}u) { break; }
40+ s = s + a[i] * b[i];
41+ i = i + ${WG}u;
42+ }
43+ partials[lid.x] = s;
44+ workgroupBarrier();
45+ var stride = ${WG / 2}u;
46+ loop {
47+ if (stride == 0u) { break; }
48+ if (lid.x < stride) {
49+ partials[lid.x] = partials[lid.x] + partials[lid.x + stride];
50+ }
51+ workgroupBarrier();
52+ stride = stride / 2u;
53+ }
54+ if (lid.x == 0u) {
55+ out[0] = partials[0];
56+ }
57+}
58+`;
59+
60+export class ReducePlan {
61+ #device: GPUDevice;
62+ #layout: GPUBindGroupLayout;
63+ /** Element count is baked into the shader, so pipelines cache per length. */
64+ #byN = new Map<number, GPUComputePipeline>();
65+
66+ constructor(device: GPUDevice) {
67+ this.#device = device;
68+ const entry = (
69+ binding: number,
70+ type: GPUBufferBindingType,
71+ ): GPUBindGroupLayoutEntry => ({
72+ binding,
73+ visibility: GPUShaderStage.COMPUTE,
74+ buffer: { type },
75+ });
76+ this.#layout = device.createBindGroupLayout({
77+ entries: [entry(0, 'storage'), entry(1, 'read-only-storage'), entry(2, 'read-only-storage')],
78+ });
79+ }
80+
81+ async #pipeline(n: number): Promise<GPUComputePipeline> {
82+ const cached = this.#byN.get(n);
83+ if (cached) return cached;
84+ const device = this.#device;
85+ device.pushErrorScope('validation');
86+ const code = dotWGSL(n);
87+ const module = device.createShaderModule({ code, label: `dot-${n}` });
88+ const info = await module.getCompilationInfo();
89+ const errors = info.messages.filter((m) => m.type === 'error');
90+ if (errors.length) {
91+ throw new Error(
92+ `WGSL compile error in dot(${n}):\n` +
93+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
94+ );
95+ }
96+ const pipeline = await device.createComputePipelineAsync({
97+ layout: device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
98+ compute: { module, entryPoint: 'main' },
99+ label: `dot-${n}`,
100+ });
101+ const err = await device.popErrorScope();
102+ if (err) throw new Error(`pipeline dot(${n}): ${err.message}`);
103+ this.#byN.set(n, pipeline);
104+ return pipeline;
105+ }
106+
107+ async createDotBinding(
108+ a: GPUBuffer,
109+ b: GPUBuffer,
110+ out: GPUBuffer,
111+ n: number,
112+ ): Promise<DotBinding> {
113+ const pipeline = await this.#pipeline(n);
114+ const bindGroup = this.#device.createBindGroup({
115+ layout: this.#layout,
116+ entries: [
117+ { binding: 0, resource: { buffer: out } },
118+ { binding: 1, resource: { buffer: a } },
119+ { binding: 2, resource: { buffer: b } },
120+ ],
121+ });
122+ return { pipeline, bindGroup };
123+ }
124+
125+ encodeDotInto(pass: GPUComputePassEncoder, binding: DotBinding): void {
126+ pass.setPipeline(binding.pipeline);
127+ pass.setBindGroup(0, binding.bindGroup);
128+ pass.dispatchWorkgroups(1);
129+ }
130+}
src/mgpu/session.tsmodified+29−1View file
@@ -10,6 +10,7 @@ import { ShtPlan } from '../sht/sht.ts';
1010 import { DerivPlan } from '../sht/deriv.ts';
1111 import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
1212 import { GpuModel, type ModelParams } from './model.ts';
13+import { DEFAULT_SOLVER, modelLibs, solveShim, type SolverKey } from './libs.ts';
1314 import { seededNoise } from './noise.ts';
1415 import type { MModel } from './registry.ts';
1516 import { Geometry } from '../geom/geometry.ts';
@@ -36,6 +37,17 @@ export interface ModelSessionOptions {
3637 * is unrolled into the op sequence, so a change recompiles.
3738 */
3839 niter?: number;
40+ /**
41+ * Which solver answers the models' `solve(...)` call (default richardson).
42+ * Structural like niter: the choice is a one-line generated shim compiled
43+ * with the model, so changing it recompiles.
44+ */
45+ solver?: SolverKey;
46+ /**
47+ * Overrides for the shared .m files, by file name (`'dlap.m'`,
48+ * `'richardson.m'`, ...) — the editor's working copies.
49+ */
50+ libSources?: Record<string, string>;
3951 }
4052
4153 export class ModelSession {
@@ -47,6 +59,8 @@ export class ModelSession {
4759 readonly npts: number;
4860 /** Iterations of the implicit solve compiled into the step. */
4961 readonly niter: number;
62+ /** The solver compiled behind the models' solve(...) call. */
63+ readonly solver: SolverKey;
5064
5165 /** The surface being solved on, as spherical-harmonic coefficients. */
5266 #geometry: Geometry;
@@ -76,6 +90,7 @@ export class ModelSession {
7690 geometryModel: MGeometry;
7791 deriv: DerivPlan;
7892 niter: number;
93+ solver: SolverKey;
7994 }) {
8095 this.device = init.device;
8196 this.model = init.model;
@@ -90,6 +105,7 @@ export class ModelSession {
90105 this.#geometryModel = init.geometryModel;
91106 this.#deriv = init.deriv;
92107 this.niter = init.niter;
108+ this.solver = init.solver;
93109 }
94110
95111 get geometry(): Geometry {
@@ -109,6 +125,17 @@ export class ModelSession {
109125 const { device, model, params, lmax } = opts;
110126 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
111127 const niter = Math.max(0, Math.round(opts.niter ?? 1));
128+ const solver = opts.solver ?? DEFAULT_SOLVER;
129+ // The shared files, with the editor's working copies substituted, plus
130+ // the shim that routes solve(...) to the chosen solver.
131+ const libs = [
132+ ...modelLibs.map((f) =>
133+ opts.libSources?.[f.name] !== undefined
134+ ? { name: f.name, source: opts.libSources[f.name] }
135+ : f,
136+ ),
137+ solveShim(solver),
138+ ];
112139 const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
113140 const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
114141 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
@@ -156,11 +183,12 @@ export class ModelSession {
156183 geometry,
157184 deriv,
158185 niter,
186+ libs,
159187 });
160188 gpu.setParams(params);
161189 return new ModelSession({
162190 device, model, cfg, sht, displaySht, gpu, params, oversample,
163- geometry, geometryModel, deriv, niter,
191+ geometry, geometryModel, deriv, niter, solver,
164192 });
165193 } catch (e) {
166194 // The transform plans own GPU buffers; do not leak them on a compile error.
src/mgpu/wgsl.tsmodified+21−12View file
@@ -85,7 +85,10 @@ function f32Lit(v: number): string {
8585
8686 /** How a scalar or tensor operand is read inside the kernel. */
8787 export interface KernelInputs {
88- /** cName -> storage binding index, for multi-element tensor operands. */
88+ /** cName -> storage binding index, for buffer-backed operands. A
89+ * multi-element tensor is read at the output's index; a single-element
90+ * value (a `dot` result, or a scalar computed from one) is read at [0],
91+ * which broadcasts it across the output. */
8992 tensors: Map<string, number>;
9093 /** cName -> slot in the params storage buffer, for runtime scalars. */
9194 params: Map<string, number>;
@@ -137,12 +140,16 @@ function emitExpr(e: IRExpr, ctx: Ctx): string {
137140 return f32Lit(e.value);
138141
139142 case 'Var': {
143+ // Buffer-backed operands first: a single-element value in a buffer
144+ // shadows any compile-time definition the same name had earlier (a
145+ // scalar seeded `rho = 1` and then updated from a dot result inside
146+ // the solve loop reads as a buffer from the update on).
147+ const bound = io.tensors.get(e.cName);
148+ if (bound !== undefined) {
149+ return isTensor(e.ty) ? `in${bound}[i]` : `in${bound}[0]`;
150+ }
140151 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]`;
152+ throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
146153 }
147154 // Scalar: either an exact compile-time value or a runtime parameter.
148155 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
@@ -195,15 +202,17 @@ function emitExpr(e: IRExpr, ctx: Ctx): string {
195202 const fn = CALL_FNS[e.name];
196203 const b = getBuiltin(e.name);
197204 if (!fn || !b?.elementwise) {
198- // A call numbl resolved to another function in the file gets a mangled
199- // specialization name; a builtin keeps its source-level name. Only the
200- // model's entry points are compiled, so a helper is a distinct failure
201- // from an unsupported builtin and deserves to say so.
205+ // A call numbl resolved to another function in the workspace gets a
206+ // mangled specialization name; a builtin keeps its source-level name.
207+ // User-function calls are expanded into the caller before planning
208+ // (src/mgpu/inlineCalls.ts), so one surviving to kernel emission means
209+ // the expansion did not reach it — a distinct failure from an
210+ // unsupported builtin, and worth saying so.
202211 const isUserFunction = e.cName !== e.name;
203212 throw new UnsupportedOnGpu(
204213 isUserFunction
205- ? `'${e.name}' is a function defined in this model. Only init and ` +
206- `step are compiled — inline its body into the caller.`
214+ ? `the call to '${e.name}' was not expanded into the caller — ` +
215+ `assign its result to a variable on its own line`
207216 : `'${e.name}' cannot be evaluated element-wise on the GPU`,
208217 e.span,
209218 );
test/geometryChecks.tsmodified+117−4View file
@@ -37,10 +37,12 @@ import {
3737 SPHERE_KEY,
3838 } from '../src/geom/registry.ts';
3939 import { ModelCompileError } from '../src/mgpu/errors.ts';
40+import { relL2 } from '../src/mgpu/digest.ts';
4041 import type { Check, Log } from './analyticChecks.ts';
4142
4243 const LMAX = 31;
4344 const STEPS = 20;
45+
4446 /** The app's actual default lmax (README: "at the default lmax 63 that is a
4547 * 128x256 grid"), used for the niter/geometry sweep below and the peanut
4648 * check next to it -- the divergence they're both about is a real, lmax-
@@ -312,12 +314,12 @@ export async function geometryChecks(
312314 // Unrolling has to be exactly linear in the trip count: the body planned
313315 // once per iteration, no more and no less. Per species per iteration: 8
314316 // dtheta/dphi + 4 analys transforms (Algorithm 3's cost, applied to the
315- // field and to each of its three Cartesian gradient components) plus 15
317+ // field and to each of its three Cartesian gradient components) plus 14
316318 // generated kernels -- see test/modelChecks.ts's KERNELS_PER_ITERATION,
317319 // which counts the kernels alone; this counts every op, transforms
318320 // included.
319321 const perIteration = ops[1] - ops[0];
320- const want = 54;
322+ const want = 52;
321323 check(
322324 'loop: unrolling is exactly linear in the trip count',
323325 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
@@ -377,6 +379,54 @@ export async function geometryChecks(
377379 );
378380 }
379381
382+ // ---- three solvers, one operator ----------------------------------------
383+ // The Krylov solvers against richardson on the same implicit system: a
384+ // Krylov iteration converges superlinearly where the stationary one
385+ // converges linearly, so at equal niter it must land much closer to the
386+ // converged answer. Switched exactly the way the app's solver control does
387+ // — the session's `solver` option, which swaps the solve(...) shim. The
388+ // comparison is a ratio against the same reference, which keeps it
389+ // meaningful on SwiftShader's looser fp32 too.
390+ {
391+ const model = mModelByKey('schnakenberg')!;
392+ const params = defaultParams(model);
393+ const ellipsoid = mGeometryByKey('ellipsoid')!;
394+ const run = async (
395+ solver: 'richardson' | 'bicgstab' | 'gmres',
396+ niter: number,
397+ ): Promise<Float32Array> => {
398+ const session = await ModelSession.create({
399+ device, model, params, lmax: LMAX,
400+ geometry: ellipsoid, geometryParams: defaultGeometryParams(ellipsoid),
401+ niter, solver,
402+ });
403+ session.seed(1);
404+ session.step(STEPS);
405+ const U = await session.read('U');
406+ session.destroy();
407+ return U;
408+ };
409+ const ref = await run('richardson', 8); // effectively converged
410+ const rich = await run('richardson', 2);
411+ const bicg = await run('bicgstab', 2);
412+ const gmres = await run('gmres', 2);
413+ const relRich = relL2(rich, ref);
414+ const relBicg = relL2(bicg, ref);
415+ const relGmres = relL2(gmres, ref);
416+ check(
417+ 'solvers: bicgstab(2) converges far past richardson(2) on the same operator',
418+ bicg.every((v) => Number.isFinite(v)) && relBicg < relRich / 5 && relBicg < 1e-4,
419+ `relL2 vs richardson(8): bicgstab ${relBicg.toExponential(2)}, ` +
420+ `richardson ${relRich.toExponential(2)}`,
421+ );
422+ check(
423+ 'solvers: gmres(2) converges far past richardson(2) on the same operator',
424+ gmres.every((v) => Number.isFinite(v)) && relGmres < relRich / 5 && relGmres < 1e-3,
425+ `relL2 vs richardson(8): gmres ${relGmres.toExponential(2)}, ` +
426+ `richardson ${relRich.toExponential(2)}`,
427+ );
428+ }
429+
380430 // ---- niter x geometry sweep: catch a "doesn't run" regression early -----
381431 // This is what actually turned up the two real issues found while building
382432 // the correction: peanut diverging at niter >= 4 with schnak-spots'
@@ -434,14 +484,77 @@ export async function geometryChecks(
434484 );
435485 }
436486 }
487+
488+ // The other side of KNOWN_DIVERGENT: on the very combinations where the
489+ // Richardson iteration leaves its convergence radius, the Krylov solvers
490+ // — same operator, same preconditioner — keep converging in niter. This
491+ // is what having the solver as its own .m is for.
492+ {
493+ const model = mModelByKey('schnakenberg')!;
494+ const params = defaultParams(model);
495+ const peanut = mGeometryByKey('peanut')!;
496+ const run = async (solver: 'bicgstab' | 'gmres', niter: number): Promise<Float32Array> => {
497+ const session = await ModelSession.create({
498+ device, model, params, lmax: SWEEP_LMAX,
499+ geometry: peanut, geometryParams: defaultGeometryParams(peanut),
500+ niter, solver,
501+ });
502+ session.seed(1);
503+ session.step(STEPS);
504+ const U = await session.read('U');
505+ session.destroy();
506+ return U;
507+ };
508+ const finiteAll = (U: Float32Array): boolean => U.every((v) => Number.isFinite(v));
509+
510+ const bicg = new Map<number, Float32Array>();
511+ for (const niter of [1, 2, 4, 8]) bicg.set(niter, await run('bicgstab', niter));
512+ const bref = bicg.get(8)!;
513+ const brel = (n: number): number => relL2(bicg.get(n)!, bref);
514+ check(
515+ 'sweep: bicgstab converges on the peanut combinations richardson cannot',
516+ [...bicg.values()].every(finiteAll) && brel(4) < brel(2) && brel(2) < brel(1),
517+ `relL2 vs bicgstab(8): niter 1 -> ${brel(1).toExponential(2)}, ` +
518+ `2 -> ${brel(2).toExponential(2)}, 4 -> ${brel(4).toExponential(2)}`,
519+ );
520+
521+ // gmres exercises the whole indexed-access machinery (the basis bank,
522+ // the Hessenberg updates, the triangular inner loops) at the sweep's
523+ // full lmax, on the operator's hardest shipped case.
524+ const g1 = await run('gmres', 1);
525+ const g4 = await run('gmres', 4);
526+ const grel1 = relL2(g1, bref);
527+ const grel4 = relL2(g4, bref);
528+ check(
529+ 'sweep: gmres converges there too',
530+ finiteAll(g1) && finiteAll(g4) && grel4 < grel1,
531+ `relL2 vs bicgstab(8): niter 1 -> ${grel1.toExponential(2)}, ` +
532+ `4 -> ${grel4.toExponential(2)}`,
533+ );
534+ }
437535 }
438536
439537 // ---- a loop whose length is not known at compile time is refused --------
440538 {
441539 const model = mModelByKey('allencahn')!;
442540 // `dt` is a tunable parameter, so it reaches the compiler with no value:
443- // the plan cannot know how many iterations to emit.
444- const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
541+ // the plan cannot know how many iterations to emit. The model's own loop
542+ // lives in solvers/richardson.m now, so the bad loop is written out here.
543+ const bad = `
544+function [U, u] = init(noise)
545+ U = analys(noise);
546+ u = synth(U);
547+end
548+
549+function [Un, u] = step(U, lam, eps2, dt, niter)
550+ u = synth(U);
551+ Bu = U + dt * analys(u - u.^3);
552+ Un = Bu ./ (1 + (dt * eps2) * lam);
553+ for k = 1:dt
554+ Un = Un + 0 * Un;
555+ end
556+end
557+`;
445558 let message = '';
446559 try {
447560 const session = await ModelSession.create({
test/modelChecks.tsmodified+229−12View file
@@ -11,6 +11,7 @@
1111 */
1212 import { ModelSession } from '../src/mgpu/session.ts';
1313 import { mModels, defaultParams } from '../src/mgpu/registry.ts';
14+import { eigenvalues, weightMask } from '../src/mgpu/model.ts';
1415 import {
1516 formatCommand,
1617 parseArgs,
@@ -31,20 +32,22 @@ const EXPECTED_KERNELS: Record<string, number> = {
3132 };
3233
3334 /**
34- * What one unrolled iteration of the solve loop adds, total (not per
35- * species — the surface Laplace-Beltrami correction's per-species kernel
36- * count is a byproduct of exactly how its expression tree happens to fuse,
37- * not a clean per-species multiple, so this is measured per model rather
38- * than derived from `model.species.length`). Each species' correction is
39- * Algorithm 3 of evolving_surface/notes/algos.tex: a surface gradient
40- * (dtheta/dphi contracted through the metric), reanalysed per Cartesian
41- * component and differentiated again, recombined into the divergence, plus
42- * the round-sphere eigenvalue added back — see models/schnakenberg.m and
43- * docs/richardson-iteration.md.
35+ * What one unrolled iteration of the solve loop adds — 14 kernels per
36+ * species: 12 in lib/dlap.m's operator (the gradient contraction, the three
37+ * re-analysed components, the five-step divergence accumulation), plus
38+ * solvers/richardson.m's dtD*lam divisor temp and its update divide. Each
39+ * species' correction is Algorithm 3 of evolving_surface/notes/algos.tex: a
40+ * surface gradient (dtheta/dphi contracted through the metric), reanalysed
41+ * per Cartesian component and differentiated again, recombined into the
42+ * divergence, plus the round-sphere eigenvalue added back — see
43+ * solvers/richardson.m, lib/dlap.m and docs/richardson-iteration.md.
44+ * (Before the solver was factored out, the monolithic models compiled to 15
45+ * per species: the interleaved species order kept the second species'
46+ * divisor from fusing into its divide.)
4447 */
4548 const KERNELS_PER_ITERATION: Record<string, number> = {
46- schnakenberg: 30,
47- brusselator: 30,
49+ schnakenberg: 28,
50+ brusselator: 28,
4851 allencahn: 14,
4952 };
5053
@@ -75,6 +78,7 @@ export async function modelChecks(
7578 geometry: 'peanut',
7679 geometryParams: { waist: 0.45, stretch: 1.25 },
7780 niter: 3,
81+ solver: 'bicgstab',
7882 };
7983 const command = formatCommand(spec);
8084 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
@@ -137,6 +141,219 @@ export async function modelChecks(
137141 session.destroy();
138142 }
139143
144+ // User-defined subroutines: a .m may define its own functions (and call the
145+ // shared solver/operator library), and each call is expanded into the caller
146+ // at compile time (src/mgpu/inlineCalls.ts). This model exercises the
147+ // shapes the shipped models do not: a multi-output function, a
148+ // scalar-returning function, a function reassigning its own parameter, and
149+ // a solver-like local whose loop bound arrives as the `niter` argument.
150+ {
151+ const model = mModels.find((m) => m.key === 'allencahn')!;
152+ const source = `
153+function [U, u] = init(noise)
154+ U = analys(noise);
155+ u = synth(U);
156+end
157+
158+function [Un, u] = step(U, lam, eps2, dt, niter)
159+ u = synth(U);
160+ [p, q] = react(u, dt);
161+ s = gain(eps2, dt);
162+ Bu = U + s * analys(p - q);
163+ Un = solveid(Bu, lam, dt, niter);
164+end
165+
166+function [p, q] = react(x, c)
167+ p = x + c * (x .* x);
168+ q = c * (x .* x);
169+end
170+
171+function y = gain(a, b)
172+ y = a + 2 * b;
173+end
174+
175+function X = solveid(B, lam, c, n)
176+ X = B ./ (1 + c * lam);
177+ for k = 1:n
178+ X = (B + c * (0 * X)) ./ (1 + c * lam);
179+ end
180+end
181+`;
182+ const session = await ModelSession.create({
183+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 2,
184+ });
185+ session.seed(1);
186+ session.step(STEPS);
187+ const values = await session.read('u');
188+ let finite = true;
189+ for (const v of values) if (!Number.isFinite(v)) finite = false;
190+ check(
191+ 'subroutines: a model composed of user functions compiles and runs',
192+ finite,
193+ `${session.describe().step.length} ops/step after expansion`,
194+ );
195+ session.destroy();
196+ }
197+
198+ // The reduction op and GPU-resident scalars: `dot` runs as a single
199+ // reduction dispatch into a 1-element buffer, scalars computed from its
200+ // result compile to 1-element kernels, and a single-element value
201+ // broadcasts into element-wise expressions as `in[0]`. These are the
202+ // primitives the Krylov solver is made of, checked directly against the
203+ // CPU here so a solver-level failure has somewhere smaller to point.
204+ {
205+ const model = mModels.find((m) => m.key === 'allencahn')!;
206+ const source = `
207+function [U, u] = init(noise)
208+ U = analys(noise);
209+ u = synth(U);
210+end
211+
212+function [Un, u] = step(U, lam, wlm, eps2, dt, niter)
213+ u = synth(U);
214+ s = dot(U, U);
215+ Uw = U .* wlm;
216+ sw = dot(Uw, lam);
217+ s2 = 2 * s;
218+ s3 = s2 - s;
219+ Un = (s * U) ./ s;
220+end
221+`;
222+ const session = await ModelSession.create({
223+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
224+ });
225+ session.seed(1);
226+ session.step(1);
227+ const U = await session.read('U');
228+ const nlm = U.length / 2;
229+ const cfg = session.cfg;
230+
231+ let cpuS = 0;
232+ for (let i = 0; i < U.length; i++) cpuS += U[i] * U[i];
233+ const gpuS = (await session.read('s'))[0];
234+ check(
235+ 'dot: matches the CPU sum',
236+ Math.abs(gpuS - cpuS) <= 1e-5 * Math.abs(cpuS),
237+ `gpu ${gpuS.toExponential(6)} vs cpu ${cpuS.toExponential(6)}`,
238+ );
239+
240+ const wlm = weightMask(cfg, nlm);
241+ const lam = eigenvalues(cfg, nlm);
242+ let cpuSw = 0;
243+ for (let i = 0; i < U.length; i++) cpuSw += U[i] * wlm[i] * lam[i];
244+ const gpuSw = (await session.read('sw'))[0];
245+ check(
246+ 'dot: the wlm-weighted inner product matches the CPU',
247+ Math.abs(gpuSw - cpuSw) <= 1e-5 * Math.abs(cpuSw),
248+ `gpu ${gpuSw.toExponential(6)} vs cpu ${cpuSw.toExponential(6)}`,
249+ );
250+
251+ // 2s - s is exact in any IEEE arithmetic, so the whole scalar chain
252+ // (reduction -> 1-element kernels -> readback) must return s's bits.
253+ const gpuS3 = (await session.read('s3'))[0];
254+ check('dot: scalar arithmetic on the result is exact', gpuS3 === gpuS,
255+ `s3 ${gpuS3.toExponential(6)} vs s ${gpuS.toExponential(6)}`);
256+
257+ const Un = await session.read('Un');
258+ let worst = 0;
259+ let scale = 0;
260+ for (let i = 0; i < U.length; i++) {
261+ worst = Math.max(worst, Math.abs(Un[i] - U[i]));
262+ scale = Math.max(scale, Math.abs(U[i]));
263+ }
264+ check(
265+ 'dot: a 1-element value broadcasts into an element-wise kernel',
266+ worst <= 1e-6 * scale,
267+ `(s*U)./s vs U: worst |d| = ${worst.toExponential(2)}`,
268+ );
269+ session.destroy();
270+ }
271+
272+ // The indexed-access ops (getslab/setslab on a bank of spectral fields,
273+ // getat/setat on a small matrix): functional updates the planner compiles
274+ // to static-offset buffer copies. Everything below has an exact expected
275+ // value, so the offsets themselves are what is being checked.
276+ {
277+ const model = mModels.find((m) => m.key === 'allencahn')!;
278+ const source = `
279+function [U, u] = init(noise)
280+ U = analys(noise);
281+ u = synth(U);
282+end
283+
284+function [Un, u] = step(U, lam, eps2, dt, nlm, niter)
285+ u = synth(U);
286+ A = zeros(2, 2);
287+ s1 = dot(U, U);
288+ s2 = 2 * s1;
289+ A = setat(A, s1, 1, 1);
290+ A = setat(A, s2, 2, 2);
291+ a11 = getat(A, 1, 1);
292+ a22 = getat(A, 2, 2);
293+ a21 = getat(A, 2, 1);
294+ chk = a22 - 2 * a11 + a21;
295+ VB = zeros(2, nlm * 2);
296+ VB = setslab(VB, U, 2);
297+ U2 = getslab(VB, 2);
298+ Z1 = getslab(VB, 1);
299+ Un = U2 + Z1;
300+end
301+`;
302+ const session = await ModelSession.create({
303+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
304+ });
305+ session.seed(1);
306+ session.step(1);
307+ // a22 - 2*a11 + a21 = 2*s - 2*s + 0, exactly, if every element landed
308+ // where its indices say.
309+ const chk = (await session.read('chk'))[0];
310+ check('indexing: matrix elements round-trip through setat/getat', chk === 0,
311+ `a22 - 2*a11 + a21 = ${chk}`);
312+ // The slab written at 2 must come back; the slab at 1 must still be zero.
313+ const U = await session.read('U');
314+ const Un = await session.read('Un');
315+ let same = U.length === Un.length;
316+ for (let i = 0; same && i < U.length; i++) if (Un[i] !== U[i]) same = false;
317+ check('indexing: a spectral field round-trips through setslab/getslab', same,
318+ same ? 'getslab(setslab(VB, U, 2), 2) + zeros = U, element for element' : 'mismatch');
319+ session.destroy();
320+ }
321+
322+ // A recursive function cannot unroll into a fixed op sequence, and must be
323+ // refused with a message that says so, not hang the compiler.
324+ {
325+ const model = mModels.find((m) => m.key === 'allencahn')!;
326+ const source = `
327+function [U, u] = init(noise)
328+ U = analys(noise);
329+ u = synth(U);
330+end
331+
332+function [Un, u] = step(U, lam, eps2, dt, niter)
333+ u = synth(U);
334+ Un = f(U);
335+end
336+
337+function y = f(x)
338+ y = f(x) + 1;
339+end
340+`;
341+ let message = '';
342+ try {
343+ const session = await ModelSession.create({
344+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
345+ });
346+ session.destroy();
347+ } catch (e) {
348+ message = e instanceof Error ? e.message : String(e);
349+ }
350+ check(
351+ 'subroutines: recursion is refused at compile time',
352+ message.includes('recursion'),
353+ message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
354+ );
355+ }
356+
140357 // The oversampled readback: readSpecies must be the state synthesized on the
141358 // display grid. Comparing against the display plan's own upload path
142359 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU