Factor the solver out of the models; add BiCGSTAB and GMRES
The models now form the reaction and hand the implicit diffusion solve to a
solver .m — richardson, bicgstab or gmres — all applying the same operator,
lib/dlap.m. Swapping solvers is editing the call line in the model.
- User-function calls (model-local or from the shared lib/ and solvers/
files) are expanded into the caller before the fusion pass
(src/mgpu/inlineCalls.ts), so a call compiles exactly like the same code
written inline. Recursion and runtime loop bounds are refused at compile
time; an inner loop bound may be an enclosing unrolled loop's variable.
- New GPU primitives: a dot reduction into a 1-element buffer
(src/mgpu/reduce.ts), GPU-resident scalars (1-element kernels) that
broadcast into vector expressions, half-spectrum inner-product weights
wlm, and indexed access (getslab/setslab on a bank of spectral fields,
getat/setat on small matrices) compiled to static-offset buffer copies.
- solvers/bicgstab.m and solvers/gmres.m keep every scalar on the GPU and
guard every ratio (a*b/(b*b + 1e-30)), so a converged iteration goes
stationary instead of dividing noise by noise. On the peanut at lmax 63,
where richardson diverges for niter >= 2, both converge monotonically.
- Richardson results are unchanged: schnakenberg and allencahn bit-identical
to the previous build; brusselator within 1 ulp/step (the old compile left
the second species' divisor unfused; the factored form fuses both).
- Tests: subroutine expansion, the reduction/indexing primitives against
exact expected values, and solver convergence comparisons. README and
docs/richardson-iteration.md updated; stale pre-operator UI text removed.
22 changed files+1912−334
README.mdmodified+154−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, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, niter)
112129 u = synth(U);
113130 v = synth(V);
114131 uuv = u .* u .* v;
@@ -116,44 +133,91 @@ 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 = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
137+ Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, 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+Trying a different solver against the same operator is a change to those two
142+call lines: every solver composes from `dlap` (the matvec is
143+`(1 + dtD.*lam).*x - dtD.*dlap(x)`, the preconditioner the elementwise
144+divide), and which one a model calls is part of what compiles — swapping
145+recompiles, like changing `niter` already does. The solver is written as a
146+full re-evaluation rather than an accumulated correction on purpose: where
147+`dlap` computes to zero there is no correction to mis-round, and the divide is
148+turing-sphere's arithmetic unchanged.
149+
150+Three solvers ship. [`solvers/bicgstab.m`](solvers/bicgstab.m) solves the
151+same system by preconditioned BiCGSTAB — same `dlap`, same preconditioner, a
152+Krylov recurrence instead of a stationary one, at two `dlap` evaluations per
153+iteration instead of one. Its scalars (`rho`, `alpha`, `omega`) never touch
154+the CPU: `dot` is a GPU reduction into a 1-element buffer, the recurrences on
155+its results compile to 1-element kernels, and a single-element value
156+broadcasts into the vector updates. Inner products carry the half-spectrum
157+weight `wlm` (m > 0 counts twice), making them the real L2 inner products on
158+the sphere. With no residual test, every ratio `a/b` is written in the
159+guarded form `a*b/(b*b + 1e-30)`, so a converged (or broken-down) iteration
160+goes stationary instead of dividing noise by noise. The difference is not
161+academic: at the app's default lmax, Schnakenberg on the peanut sits outside
162+the Richardson iteration's convergence radius for `niter ≥ 2` and diverges,
163+while BiCGSTAB on the identical operator converges monotonically — the tests
164+pin both behaviors, side by side.
165+
166+[`solvers/gmres.m`](solvers/gmres.m) is right-preconditioned GMRES(niter) —
167+one Arnoldi sweep, no restart — with the residual minimized over the whole
168+Krylov space. Its bookkeeping is what the other solvers never need: a basis
169+of niter+1 spectral fields, a Hessenberg matrix, Givens rotations, a
170+triangular back-substitution. The basis lives in a *bank* (`getslab` /
171+`setslab`: the k-th 2 × nlm field of a wider array), the small matrices are
172+element-addressed (`getat` / `setat`), and both are functional updates the
173+planner compiles to static-offset buffer copies — MATLAB's own `H(i,j) = h`
174+cannot lower, because numbl must prove an indexed write in bounds before the
175+loop unrolls, and a loop variable has no value yet at that point. Written as
176+calls, the index resolves at *planning*, where unrolling has made it a
177+literal. The same resolution lets an inner loop bound depend on the outer
178+loop's variable, which is what makes the `for i = 1:j` orthogonalization
179+sweep compile.
180+
181+### Where the geometry enters the operator
182+
183+[`lib/dlap.m`](lib/dlap.m) is Algorithm 3 of the evolving-surface notes: the
184+field's θ/φ derivatives (the `dtheta`/`dphi` transforms,
185+[`src/sht/deriv.ts`](src/sht/deriv.ts)) are contracted through the inverse
186+metric quantities into a tangential gradient; each Cartesian component is
187+re-analysed and differentiated again; the results recombine into the surface
188+divergence, and `lam .* F` adds back what the round-sphere part already
189+carries. The metric quantities `Vt*`/`Vp*`
190+([`src/geom/metric.ts`](src/geom/metric.ts)) are built once from the
191+embedding's derivatives when the geometry is (re)built — the geometry is
192+static, so per step they are just six more buffers the kernels read. `filt`
193+zeroes the top two spectral degrees wherever the operator re-differentiates,
194+because the derivative recurrences cannot exactly represent a derivative
195+there.
196+
197+On the sphere `dlap` computes to (numerical) zero, so any `niter` lands
198+within transform round-off of the exact round-sphere answer — asserted in the
199+tests. Off the sphere the correction genuinely moves the answer, and
200+convergence is a real constraint: the fixed-count loop has no residual check,
201+so the tests also pin which shape/niter combinations are known to sit outside
202+the convergence radius and diverge.
203+
204+### Subroutines
205+
206+A model file is not limited to `init` and `step`: it can define further
207+functions and call them, and every model compiles against the shared library
208+files — [`lib/`](lib/) for operators, [`solvers/`](solvers/) for solvers —
209+with MATLAB's visibility rules (a file's namesake function is public; a
210+model-local function of the same name shadows it). numbl specializes each
211+callee for the argument types at its call sites, and the host then splices
212+the lowered body into the caller, one clone per call site
213+([`src/mgpu/inlineCalls.ts`](src/mgpu/inlineCalls.ts)): arguments bind by
214+renaming rather than copying, and assignments to a callee output become
215+assignments to the caller's variable, which is what lets a solver iterate its
216+result in place. Expansion runs before the fusion pass, so a call fuses
217+exactly as the same code written inline would — the boundary costs nothing,
218+and `describe()`'s op listing names the expanded internals
219+(`richardson#1.X`). Recursion cannot unroll into a fixed op sequence and is
220+refused at compile time, like a runtime loop bound.
157221
158222 ### `for` loops, unrolled
159223
@@ -161,7 +225,8 @@ A plan is a fixed list of GPU operations with no branching, which is what makes
161225 a timestep pure command recording — one submit, no CPU in the loop. A counted
162226 loop still fits: the planner
163227 ([`src/mgpu/plan.ts`](src/mgpu/plan.ts)) unrolls it, planning the body once per
164-iteration.
228+iteration. The loop that matters is `solvers/richardson.m`'s `for k = 1:niter`,
229+expanded into each model's step at every solve call site.
165230
166231 Nothing else had to change for that, because numbl gives a variable one cName
167232 for every assignment to it: the buffer an iteration writes is the buffer the
@@ -173,8 +238,11 @@ Two consequences worth stating:
173238
174239 - **The bounds must be known when the model compiles.** `niter` is supplied as a
175240 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.
241+ parameter, which is a uniform. A bound may also be an enclosing unrolled
242+ loop's variable (`for i = 1:j` — each unrolled `j` plans its own inner trip
243+ count, which is how GMRES's triangular sweeps compile). A genuinely runtime
244+ bound is refused at compile time with a source position, not silently
245+ mis-compiled, and there is a test for that.
178246 - **Fusion survives.** numbl's inline pass recurses into loop bodies, so a line
179247 inside the loop is still one kernel. It runs there with no protected names,
180248 though, which means an assignment whose only visible use is later in the same
@@ -183,23 +251,32 @@ Two consequences worth stating:
183251 each loop body assigns before the pass and refuses the ones that escape, so
184252 that case is a compile error rather than a stale read.
185253
186-Unrolling is exactly linear in the trip count: 2 GPU ops per species per
187-iteration, asserted in the tests.
254+Unrolling is exactly linear in the trip count: 26 GPU ops per species per
255+iteration (the operator's twelve transforms and the solve's kernels), asserted
256+in the tests.
188257
189258 ## MATLAB, compiled to WebGPU
190259
191260 Unchanged from turing-sphere, and it now compiles the geometry files too. numbl
192261 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
262+grid; user-function calls are expanded into the caller, one clone per call
263+site; the inline pass folds single-use temps back into their consumer, so one
194264 line of MATLAB becomes one expression tree; and this repo emits one WGSL compute
195265 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
266+([`src/mgpu/wgsl.ts`](src/mgpu/wgsl.ts)). `synth` / `analys` (and the
267+derivative pair `dtheta` / `dphi`) are external operations whose type rules
268+numbl learns from a `.mtoc2.js` workspace file, and which the backend maps onto
269+the spherical-harmonic pipelines; `dot` is one more, mapped onto a
270+single-dispatch reduction ([`src/mgpu/reduce.ts`](src/mgpu/reduce.ts)) whose
271+1-element result stays on the GPU — scalars computed from it become 1-element
272+kernels, and reading one inside a vector expression broadcasts it. The
273+indexed-access ops (`getslab`/`setslab`, `getat`/`setat`) compile to
274+static-offset buffer copies, their indices evaluated at planning time where
275+the unrolled loop's variable is a literal. Anything it cannot express is
276+refused at compile time with a source position.
277+
278+The Schnakenberg step above compiles to 65 GPU operations at one solve
279+iteration: 28 transforms, 35 generated kernels, and 2 buffer copies feeding the
203280 new state back.
204281
205282 Two consequences carried over:
@@ -350,8 +427,16 @@ about the round sphere, so all three build on the sphere geometry:
350427 **the same coefficients give the same surface on a 2× grid** — the 2× Gauss
351428 latitudes share no point with the 1× ones, so agreeing there is agreeing
352429 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;
430+- unrolling is **exactly linear** in the trip count; on the sphere the
431+ geometric correction computes to (numerical) zero, so 0, 1 and 4 iterations
432+ agree to transform round-off; on the peanut it **measurably moves the
433+ answer**;
434+- a niter × geometry sweep stays finite except the combinations **known to sit
435+ outside the Richardson convergence radius**, which are pinned as diverging —
436+ and on exactly those combinations **bicgstab and gmres keep converging**,
437+ also pinned;
438+- at equal niter, **bicgstab and gmres land far closer to the converged
439+ answer** than richardson on the same operator;
355440 - a runtime loop bound is refused at compile time;
356441 - swapping the surface mid-run leaves the spectral state untouched.
357442
@@ -359,7 +444,13 @@ about the round sphere, so all three build on the sphere geometry:
359444 and asserts **how many kernels it compiles to**, split into the base step and
360445 what one solve iteration adds. That is a fusion guard: if numbl's inline pass
361446 stops folding, the results stay correct while every operator becomes its own
362-dispatch, which is invisible in the numbers.
447+dispatch, which is invisible in the numbers. It also compiles a model built of
448+**user-defined subroutines** — multi-output, scalar-returning, a solver-like
449+local with its own loop — asserts recursion is refused, and checks the
450+**reduction and indexing primitives** directly against exact expected values:
451+the `dot` sum and the `wlm`-weighted inner product against the CPU, scalar
452+arithmetic on a GPU-resident result bit for bit, the 1-element broadcast, and
453+element/slab round-trips through `setat`/`getat` and `setslab`/`getslab`.
363454
364455 [`test/transformChecks.ts`](test/transformChecks.ts) compares the WGSL transforms
365456 against shtns-webgpu's f64 CPU twin.
docs/richardson-iteration.mdmodified+37−7View file
@@ -64,13 +64,18 @@ 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 = richardson(Bu, dt * D1, ...)`), which is where the split pays:
72+a different solver for the same operator is a different call in the model,
73+with `lib/dlap.m` untouched. The solver is written as a full re-evaluation
74+rather than an accumulated correction `δ = Un^(k+1) - Un^(k)` on purpose:
75+where `dlap` evaluates to zero exactly, every `Un^(k)` is bit-for-bit
76+`Un^(0)`, with no cancellation to round differently. (In practice `dlap` is a
77+real computation through chained fp32 transforms, so on the round sphere it
78+lands near zero rather than at it — the tests bound how near.)
7479
7580 ## Convergence, and why it isn't GMRES
7681
@@ -102,6 +107,31 @@ algos.tex, one `dlap` evaluation per iteration, a fixed and
102107 recompile-on-change trip count, in exchange for linear rather than
103108 superlinear convergence.
104109
110+A Krylov method fits the shape too, as long as its scalars stay on the GPU —
111+which is what `solvers/bicgstab.m` does: `dot` is a reduction dispatch, the
112+alpha/omega/rho recurrences are 1-element kernels, and the iteration count is
113+still fixed and unrolled. What it cannot have is exactly what GMRES's `tol`
114+gives algos.tex: a stopping rule. It compensates in two ways — every ratio is
115+algebraically guarded (`a*b/(b*b + eps)`) so a converged iteration goes
116+stationary rather than dividing noise by noise, and the cost is fixed at two
117+`dlap` evaluations plus three reductions per iteration whether or not it has
118+already converged. In exchange it converges superlinearly, including on
119+shape/timestep combinations outside the Richardson iteration's spectral
120+radius — the tests pin Schnakenberg on the peanut at the app's default lmax
121+as exactly such a case.
122+
123+algos.tex's own method is here too, in the same fixed-count form:
124+`solvers/gmres.m` is right-preconditioned GMRES(niter) — one Arnoldi sweep,
125+Givens rotations, back-substitution — minus the restart loop and minus
126+`tol`/`maxiter`, since there is still no data-dependent stopping. Its basis
127+and Hessenberg bookkeeping run through the indexed-access ops
128+(`getslab`/`setslab`, `getat`/`setat`), which the planner compiles to
129+static-offset buffer copies once the unrolled loop's variable makes every
130+index a literal; the same guarded-ratio discipline covers the rotation and
131+back-substitution divides. Where algos.tex's GMRES stops at `tol`, this one
132+spends its fixed niter·(niter+3)/2 reductions and niter `dlap` evaluations
133+and keeps whatever residual that bought.
134+
105135 ## One more difference worth flagging
106136
107137 algos.tex maps the half-spectrum complex coefficients through a real vector
index.htmlmodified+1−23View 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,11 +164,6 @@
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
@@ -198,7 +176,7 @@
198176 <label>preset
199177 <select id="model"></select>
200178 </label>
201- <label title="The surface. Rendered, but not yet in the operator.">geometry
179+ <label title="The surface being solved on.">geometry
202180 <select id="geometry"></select>
203181 </label>
204182 <label title="Blend between the sphere (0) and the surface (1). Display only.">morph
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+3−31View file
@@ -2,7 +2,8 @@
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 solvers/richardson.m.
67
78 function [U, u] = init(noise)
89 U = analys(noise);
@@ -13,34 +14,5 @@ function [Un, u] = step(U, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
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 = richardson(Bu, dt * eps2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
4618 end
models/brusselator.mmodified+4−55View 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 solvers/richardson.m.
78
89 function [U, V, u, v] = init(noise, A, B)
910 U = analys(A + noise);
@@ -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 = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
25+ Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
7726 end
models/schnakenberg.mmodified+9−65View file
@@ -3,11 +3,12 @@
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 solvers/richardson.m — which applies the operator's geometric
9+% part through lib/dlap.m. Trying a different solver for the same operator is
10+% a change to these two call lines. Grid fields are npts x 1; spectral fields
11+% are real 2 x nlm. See docs/richardson-iteration.md.
1112
1213 function [U, V, u, v] = init(noise, a, b)
1314 us = a + b;
@@ -27,64 +28,7 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
2728 Bu = U + dt * analys(a - u + uuv);
2829 Bv = V + dt * analys(b - uuv);
2930
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
31+ % The species diffuse independently, so each gets its own solve.
32+ Un = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
33+ Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
9034 end
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/main.tsmodified+2−5View file
@@ -16,7 +16,6 @@ import {
1616 mGeometries,
1717 mGeometryByKey,
1818 defaultGeometryParams,
19- SPHERE_KEY,
2019 DEFAULT_GEOMETRY_KEY,
2120 type MGeometry,
2221 } from './geom/registry.ts';
@@ -103,7 +102,7 @@ elColormap.value = 'jet';
103102 const editor = new CodeEditor({
104103 textarea: elSource,
105104 overlay: elHighlight,
106- external: EXTERNAL_OPS,
105+ external: new Set(EXTERNAL_OPS.keys()),
107106 onInput: (value) => {
108107 if (editing === 'geometry') editedGeomSource = value;
109108 else editedSource = value;
@@ -588,11 +587,9 @@ function updateGeomNote(): void {
588587 return;
589588 }
590589 const { lo, hi } = session.geometry.radiusRange();
591- const isSphere = session.geometryModel.key === SPHERE_KEY;
592590 elGeomNote.innerHTML =
593591 `<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.');
592+ `Radius ${lo.toFixed(3)}–${hi.toFixed(3)}.`;
596593 }
597594
598595 /** Report a compile failure, and select the offending text in the editor. */
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+34−0View file
@@ -0,0 +1,34 @@
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];
src/mgpu/model.tsmodified+28−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 { modelLibs } from './libs.ts';
2627
2728 export interface ModelParams {
2829 [key: string]: number;
@@ -118,6 +119,27 @@ export function filterMask(cfg: ShtConfig, nlm: number): Float32Array {
118119 return filt;
119120 }
120121
122+/**
123+ * Inner-product weights for the half-spectrum layout: 1 at m = 0, 2 at
124+ * m > 0, duplicated across re/im like `lam`. The 2 x nlm state stores only
125+ * m >= 0 of a real field (the m < 0 coefficients are conjugates), so the
126+ * true L2 inner product on the sphere is sum(wlm .* a .* b) — which is what
127+ * a Krylov solver's `dot` calls should compute if its scalars are to mean
128+ * what they mean in the analysis.
129+ */
130+export function weightMask(cfg: ShtConfig, nlm: number): Float32Array {
131+ const wlm = new Float32Array(2 * nlm);
132+ for (let m = 0; m <= cfg.mmax; m++) {
133+ for (let l = m; l <= cfg.lmax; l++) {
134+ const i = lmIndex(cfg.lmax, l, m);
135+ const w = m === 0 ? 1 : 2;
136+ wlm[2 * i] = w;
137+ wlm[2 * i + 1] = w;
138+ }
139+ }
140+ return wlm;
141+}
142+
121143 export class GpuModel {
122144 readonly paramNames: string[];
123145 readonly state: string[];
@@ -175,6 +197,7 @@ export class GpuModel {
175197 const bindings: Record<string, Binding> = {
176198 lam: { kind: 'tensor', shape: [2, nlm] },
177199 filt: { kind: 'tensor', shape: [2, nlm] },
200+ wlm: { kind: 'tensor', shape: [2, nlm] },
178201 noise: { kind: 'tensor', shape: [npts, 1] },
179202 npts: { kind: 'const', value: npts },
180203 nlm: { kind: 'const', value: nlm },
@@ -189,7 +212,9 @@ export class GpuModel {
189212 for (const p of paramNames) bindings[p] = { kind: 'param' };
190213
191214 // Parsing belongs to the file, not to either function.
192- const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
215+ const compiled = inModel(
216+ () => new CompiledModel(source, bindings, { npts, nlm }, 'model.m', modelLibs),
217+ );
193218 // Both functions return the new state first, then the rendered grid fields.
194219 const nargout = state.length + view.length;
195220 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
@@ -207,6 +232,7 @@ export class GpuModel {
207232 for (const s of state) host.ensure(s, 2 * nlm);
208233 host.ensure('lam', 2 * nlm);
209234 host.ensure('filt', 2 * nlm);
235+ host.ensure('wlm', 2 * nlm);
210236 host.ensure('noise', npts);
211237 if (geometry) {
212238 for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
@@ -223,6 +249,7 @@ export class GpuModel {
223249
224250 host.upload('lam', eigenvalues(cfg, nlm));
225251 host.upload('filt', filterMask(cfg, nlm));
252+ host.upload('wlm', weightMask(cfg, nlm));
226253 if (geometry) {
227254 host.upload('gx', geometry.x);
228255 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/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+132−4View file
@@ -37,10 +37,32 @@ 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+
46+/**
47+ * The shipped Schnakenberg model with its solver calls switched from
48+ * richardson to another shipped solver — the same edit a user makes in the
49+ * page, which is the point: same operator, same model, different solver.
50+ * (gmres additionally takes `nlm`, for sizing its Krylov basis bank.)
51+ * Throws if the model text drifted from what this rewrites, so the test
52+ * fails loudly rather than silently comparing richardson with itself.
53+ */
54+function solverSchnak(source: string, solver: 'bicgstab' | 'gmres'): string {
55+ const out = source
56+ .replace('step(U, V, lam, filt, ', 'step(U, V, lam, filt, wlm, ')
57+ .replace('D2, dt, niter)', solver === 'gmres' ? 'D2, dt, nlm, niter)' : 'D2, dt, niter)')
58+ .replaceAll('richardson(Bu, dt * D1, lam, filt, ', `${solver}(Bu, dt * D1, lam, filt, wlm, `)
59+ .replaceAll('richardson(Bv, dt * D2, lam, filt, ', `${solver}(Bv, dt * D2, lam, filt, wlm, `)
60+ .replaceAll('Vpz, niter);', solver === 'gmres' ? 'Vpz, nlm, niter);' : 'Vpz, niter);');
61+ if (!out.includes('wlm,') || !out.includes(`${solver}(Bu`) || !out.includes(`${solver}(Bv`)) {
62+ throw new Error('solverSchnak: the model source no longer matches the rewrite');
63+ }
64+ return out;
65+}
4466 /** The app's actual default lmax (README: "at the default lmax 63 that is a
4567 * 128x256 grid"), used for the niter/geometry sweep below and the peanut
4668 * check next to it -- the divergence they're both about is a real, lmax-
@@ -312,12 +334,12 @@ export async function geometryChecks(
312334 // Unrolling has to be exactly linear in the trip count: the body planned
313335 // once per iteration, no more and no less. Per species per iteration: 8
314336 // 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
337+ // field and to each of its three Cartesian gradient components) plus 14
316338 // generated kernels -- see test/modelChecks.ts's KERNELS_PER_ITERATION,
317339 // which counts the kernels alone; this counts every op, transforms
318340 // included.
319341 const perIteration = ops[1] - ops[0];
320- const want = 54;
342+ const want = 52;
321343 check(
322344 'loop: unrolling is exactly linear in the trip count',
323345 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
@@ -377,6 +399,49 @@ export async function geometryChecks(
377399 );
378400 }
379401
402+ // ---- two solvers, one operator ------------------------------------------
403+ // solvers/bicgstab.m against solvers/richardson.m on the same implicit
404+ // system: a Krylov iteration converges superlinearly where the stationary
405+ // one converges linearly, so at equal niter it must land much closer to the
406+ // converged answer. The comparison is a ratio against the same reference,
407+ // which keeps it meaningful on SwiftShader's looser fp32 too.
408+ {
409+ const model = mModelByKey('schnakenberg')!;
410+ const params = defaultParams(model);
411+ const ellipsoid = mGeometryByKey('ellipsoid')!;
412+ const run = async (source: string | undefined, niter: number): Promise<Float32Array> => {
413+ const session = await ModelSession.create({
414+ device, model, params, lmax: LMAX,
415+ geometry: ellipsoid, geometryParams: defaultGeometryParams(ellipsoid),
416+ niter, ...(source ? { source } : {}),
417+ });
418+ session.seed(1);
419+ session.step(STEPS);
420+ const U = await session.read('U');
421+ session.destroy();
422+ return U;
423+ };
424+ const ref = await run(undefined, 8); // richardson, effectively converged
425+ const rich = await run(undefined, 2);
426+ const bicg = await run(solverSchnak(model.source, 'bicgstab'), 2);
427+ const gmres = await run(solverSchnak(model.source, 'gmres'), 2);
428+ const relRich = relL2(rich, ref);
429+ const relBicg = relL2(bicg, ref);
430+ const relGmres = relL2(gmres, ref);
431+ check(
432+ 'solvers: bicgstab(2) converges far past richardson(2) on the same operator',
433+ bicg.every((v) => Number.isFinite(v)) && relBicg < relRich / 5 && relBicg < 1e-4,
434+ `relL2 vs richardson(8): bicgstab ${relBicg.toExponential(2)}, ` +
435+ `richardson ${relRich.toExponential(2)}`,
436+ );
437+ check(
438+ 'solvers: gmres(2) converges far past richardson(2) on the same operator',
439+ gmres.every((v) => Number.isFinite(v)) && relGmres < relRich / 5 && relGmres < 1e-3,
440+ `relL2 vs richardson(8): gmres ${relGmres.toExponential(2)}, ` +
441+ `richardson ${relRich.toExponential(2)}`,
442+ );
443+ }
444+
380445 // ---- niter x geometry sweep: catch a "doesn't run" regression early -----
381446 // This is what actually turned up the two real issues found while building
382447 // the correction: peanut diverging at niter >= 4 with schnak-spots'
@@ -434,14 +499,77 @@ export async function geometryChecks(
434499 );
435500 }
436501 }
502+
503+ // The other side of KNOWN_DIVERGENT: on the very combinations where the
504+ // Richardson iteration leaves its convergence radius, the Krylov solvers
505+ // — same operator, same preconditioner — keep converging in niter. This
506+ // is what having the solver as its own .m is for.
507+ {
508+ const model = mModelByKey('schnakenberg')!;
509+ const params = defaultParams(model);
510+ const peanut = mGeometryByKey('peanut')!;
511+ const run = async (solver: 'bicgstab' | 'gmres', niter: number): Promise<Float32Array> => {
512+ const session = await ModelSession.create({
513+ device, model, params, lmax: SWEEP_LMAX,
514+ geometry: peanut, geometryParams: defaultGeometryParams(peanut),
515+ niter, source: solverSchnak(model.source, solver),
516+ });
517+ session.seed(1);
518+ session.step(STEPS);
519+ const U = await session.read('U');
520+ session.destroy();
521+ return U;
522+ };
523+ const finiteAll = (U: Float32Array): boolean => U.every((v) => Number.isFinite(v));
524+
525+ const bicg = new Map<number, Float32Array>();
526+ for (const niter of [1, 2, 4, 8]) bicg.set(niter, await run('bicgstab', niter));
527+ const bref = bicg.get(8)!;
528+ const brel = (n: number): number => relL2(bicg.get(n)!, bref);
529+ check(
530+ 'sweep: bicgstab converges on the peanut combinations richardson cannot',
531+ [...bicg.values()].every(finiteAll) && brel(4) < brel(2) && brel(2) < brel(1),
532+ `relL2 vs bicgstab(8): niter 1 -> ${brel(1).toExponential(2)}, ` +
533+ `2 -> ${brel(2).toExponential(2)}, 4 -> ${brel(4).toExponential(2)}`,
534+ );
535+
536+ // gmres exercises the whole indexed-access machinery (the basis bank,
537+ // the Hessenberg updates, the triangular inner loops) at the sweep's
538+ // full lmax, on the operator's hardest shipped case.
539+ const g1 = await run('gmres', 1);
540+ const g4 = await run('gmres', 4);
541+ const grel1 = relL2(g1, bref);
542+ const grel4 = relL2(g4, bref);
543+ check(
544+ 'sweep: gmres converges there too',
545+ finiteAll(g1) && finiteAll(g4) && grel4 < grel1,
546+ `relL2 vs bicgstab(8): niter 1 -> ${grel1.toExponential(2)}, ` +
547+ `4 -> ${grel4.toExponential(2)}`,
548+ );
549+ }
437550 }
438551
439552 // ---- a loop whose length is not known at compile time is refused --------
440553 {
441554 const model = mModelByKey('allencahn')!;
442555 // `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');
556+ // the plan cannot know how many iterations to emit. The model's own loop
557+ // lives in solvers/richardson.m now, so the bad loop is written out here.
558+ const bad = `
559+function [U, u] = init(noise)
560+ U = analys(noise);
561+ u = synth(U);
562+end
563+
564+function [Un, u] = step(U, lam, eps2, dt, niter)
565+ u = synth(U);
566+ Bu = U + dt * analys(u - u.^3);
567+ Un = Bu ./ (1 + (dt * eps2) * lam);
568+ for k = 1:dt
569+ Un = Un + 0 * Un;
570+ end
571+end
572+`;
445573 let message = '';
446574 try {
447575 const session = await ModelSession.create({
test/modelChecks.tsmodified+228−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
@@ -137,6 +140,219 @@ export async function modelChecks(
137140 session.destroy();
138141 }
139142
143+ // User-defined subroutines: a .m may define its own functions (and call the
144+ // shared solver/operator library), and each call is expanded into the caller
145+ // at compile time (src/mgpu/inlineCalls.ts). This model exercises the
146+ // shapes the shipped models do not: a multi-output function, a
147+ // scalar-returning function, a function reassigning its own parameter, and
148+ // a solver-like local whose loop bound arrives as the `niter` argument.
149+ {
150+ const model = mModels.find((m) => m.key === 'allencahn')!;
151+ const source = `
152+function [U, u] = init(noise)
153+ U = analys(noise);
154+ u = synth(U);
155+end
156+
157+function [Un, u] = step(U, lam, eps2, dt, niter)
158+ u = synth(U);
159+ [p, q] = react(u, dt);
160+ s = gain(eps2, dt);
161+ Bu = U + s * analys(p - q);
162+ Un = solveid(Bu, lam, dt, niter);
163+end
164+
165+function [p, q] = react(x, c)
166+ p = x + c * (x .* x);
167+ q = c * (x .* x);
168+end
169+
170+function y = gain(a, b)
171+ y = a + 2 * b;
172+end
173+
174+function X = solveid(B, lam, c, n)
175+ X = B ./ (1 + c * lam);
176+ for k = 1:n
177+ X = (B + c * (0 * X)) ./ (1 + c * lam);
178+ end
179+end
180+`;
181+ const session = await ModelSession.create({
182+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 2,
183+ });
184+ session.seed(1);
185+ session.step(STEPS);
186+ const values = await session.read('u');
187+ let finite = true;
188+ for (const v of values) if (!Number.isFinite(v)) finite = false;
189+ check(
190+ 'subroutines: a model composed of user functions compiles and runs',
191+ finite,
192+ `${session.describe().step.length} ops/step after expansion`,
193+ );
194+ session.destroy();
195+ }
196+
197+ // The reduction op and GPU-resident scalars: `dot` runs as a single
198+ // reduction dispatch into a 1-element buffer, scalars computed from its
199+ // result compile to 1-element kernels, and a single-element value
200+ // broadcasts into element-wise expressions as `in[0]`. These are the
201+ // primitives the Krylov solver is made of, checked directly against the
202+ // CPU here so a solver-level failure has somewhere smaller to point.
203+ {
204+ const model = mModels.find((m) => m.key === 'allencahn')!;
205+ const source = `
206+function [U, u] = init(noise)
207+ U = analys(noise);
208+ u = synth(U);
209+end
210+
211+function [Un, u] = step(U, lam, wlm, eps2, dt, niter)
212+ u = synth(U);
213+ s = dot(U, U);
214+ Uw = U .* wlm;
215+ sw = dot(Uw, lam);
216+ s2 = 2 * s;
217+ s3 = s2 - s;
218+ Un = (s * U) ./ s;
219+end
220+`;
221+ const session = await ModelSession.create({
222+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
223+ });
224+ session.seed(1);
225+ session.step(1);
226+ const U = await session.read('U');
227+ const nlm = U.length / 2;
228+ const cfg = session.cfg;
229+
230+ let cpuS = 0;
231+ for (let i = 0; i < U.length; i++) cpuS += U[i] * U[i];
232+ const gpuS = (await session.read('s'))[0];
233+ check(
234+ 'dot: matches the CPU sum',
235+ Math.abs(gpuS - cpuS) <= 1e-5 * Math.abs(cpuS),
236+ `gpu ${gpuS.toExponential(6)} vs cpu ${cpuS.toExponential(6)}`,
237+ );
238+
239+ const wlm = weightMask(cfg, nlm);
240+ const lam = eigenvalues(cfg, nlm);
241+ let cpuSw = 0;
242+ for (let i = 0; i < U.length; i++) cpuSw += U[i] * wlm[i] * lam[i];
243+ const gpuSw = (await session.read('sw'))[0];
244+ check(
245+ 'dot: the wlm-weighted inner product matches the CPU',
246+ Math.abs(gpuSw - cpuSw) <= 1e-5 * Math.abs(cpuSw),
247+ `gpu ${gpuSw.toExponential(6)} vs cpu ${cpuSw.toExponential(6)}`,
248+ );
249+
250+ // 2s - s is exact in any IEEE arithmetic, so the whole scalar chain
251+ // (reduction -> 1-element kernels -> readback) must return s's bits.
252+ const gpuS3 = (await session.read('s3'))[0];
253+ check('dot: scalar arithmetic on the result is exact', gpuS3 === gpuS,
254+ `s3 ${gpuS3.toExponential(6)} vs s ${gpuS.toExponential(6)}`);
255+
256+ const Un = await session.read('Un');
257+ let worst = 0;
258+ let scale = 0;
259+ for (let i = 0; i < U.length; i++) {
260+ worst = Math.max(worst, Math.abs(Un[i] - U[i]));
261+ scale = Math.max(scale, Math.abs(U[i]));
262+ }
263+ check(
264+ 'dot: a 1-element value broadcasts into an element-wise kernel',
265+ worst <= 1e-6 * scale,
266+ `(s*U)./s vs U: worst |d| = ${worst.toExponential(2)}`,
267+ );
268+ session.destroy();
269+ }
270+
271+ // The indexed-access ops (getslab/setslab on a bank of spectral fields,
272+ // getat/setat on a small matrix): functional updates the planner compiles
273+ // to static-offset buffer copies. Everything below has an exact expected
274+ // value, so the offsets themselves are what is being checked.
275+ {
276+ const model = mModels.find((m) => m.key === 'allencahn')!;
277+ const source = `
278+function [U, u] = init(noise)
279+ U = analys(noise);
280+ u = synth(U);
281+end
282+
283+function [Un, u] = step(U, lam, eps2, dt, nlm, niter)
284+ u = synth(U);
285+ A = zeros(2, 2);
286+ s1 = dot(U, U);
287+ s2 = 2 * s1;
288+ A = setat(A, s1, 1, 1);
289+ A = setat(A, s2, 2, 2);
290+ a11 = getat(A, 1, 1);
291+ a22 = getat(A, 2, 2);
292+ a21 = getat(A, 2, 1);
293+ chk = a22 - 2 * a11 + a21;
294+ VB = zeros(2, nlm * 2);
295+ VB = setslab(VB, U, 2);
296+ U2 = getslab(VB, 2);
297+ Z1 = getslab(VB, 1);
298+ Un = U2 + Z1;
299+end
300+`;
301+ const session = await ModelSession.create({
302+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
303+ });
304+ session.seed(1);
305+ session.step(1);
306+ // a22 - 2*a11 + a21 = 2*s - 2*s + 0, exactly, if every element landed
307+ // where its indices say.
308+ const chk = (await session.read('chk'))[0];
309+ check('indexing: matrix elements round-trip through setat/getat', chk === 0,
310+ `a22 - 2*a11 + a21 = ${chk}`);
311+ // The slab written at 2 must come back; the slab at 1 must still be zero.
312+ const U = await session.read('U');
313+ const Un = await session.read('Un');
314+ let same = U.length === Un.length;
315+ for (let i = 0; same && i < U.length; i++) if (Un[i] !== U[i]) same = false;
316+ check('indexing: a spectral field round-trips through setslab/getslab', same,
317+ same ? 'getslab(setslab(VB, U, 2), 2) + zeros = U, element for element' : 'mismatch');
318+ session.destroy();
319+ }
320+
321+ // A recursive function cannot unroll into a fixed op sequence, and must be
322+ // refused with a message that says so, not hang the compiler.
323+ {
324+ const model = mModels.find((m) => m.key === 'allencahn')!;
325+ const source = `
326+function [U, u] = init(noise)
327+ U = analys(noise);
328+ u = synth(U);
329+end
330+
331+function [Un, u] = step(U, lam, eps2, dt, niter)
332+ u = synth(U);
333+ Un = f(U);
334+end
335+
336+function y = f(x)
337+ y = f(x) + 1;
338+end
339+`;
340+ let message = '';
341+ try {
342+ const session = await ModelSession.create({
343+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
344+ });
345+ session.destroy();
346+ } catch (e) {
347+ message = e instanceof Error ? e.message : String(e);
348+ }
349+ check(
350+ 'subroutines: recursion is refused at compile time',
351+ message.includes('recursion'),
352+ message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
353+ );
354+ }
355+
140356 // The oversampled readback: readSpecies must be the state synthesized on the
141357 // display grid. Comparing against the display plan's own upload path
142358 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU