/ concept-collection / turing-surface
concept-collection / turing-surface
Batch independent transforms through one Legendre dispatch
The expensive part of every Legendre stage is generating the associated Legendre values on the fly by recurrence -- work that depends only on (m, theta), never on the field. Batched kernels walk the recurrence once for up to 4 fields, one accumulator lane each, with per-lane arithmetic textually identical to the scalar kernels': batched and scalar plans produce bit-identical states. synth/analys now take multiple fields -- [a, b] = synth(x, y) -- as explicit grouping: output k is the transform of input k, and the planner (materializeTransforms) chunks each group into whatever lane width the device supports. The lanes' fm scratch lives in one 256-aligned arena buffer, so the batched bind group is 3 tables + K fields + 1 arena = 8, exactly WebGPU's default storage-buffer limit at K = 4 -- full-width batching on every stack, SwiftShader included. Ungrouped transforms on consecutive independent lines batch the same way; SHT_BATCH=0 compiles scalar-only plans for A/B. The models group their solve-loop transforms across species (all 16 transforms of a schnakenberg step land in batches) and their init pairs. Misuse is refused at compile time: a group must bind every output, since each input costs a transform. Worth ~25% of the whole step: 0.88 vs 1.14 ms/step at lmax 127, niter 2, on bumpy -- 2.22 -> 0.88 ms cumulative with the flux-form reduction. Needs numbl >= 6035240 ("Route .mtoc2.js user functions through builtin multi-assign") for the multi-output external lowering; bump NUMBL_REF in .github/workflows when that lands.
Dan Fortunato <dan.fortunato@gmail.com> committed commit a4fee9c8e09f parent 591a4f5 Browse files
12 changed files+1098−106
README.mdmodified+25−1View file
@@ -206,6 +206,28 @@ The Schnakenberg step compiles to 51 GPU operations at one solve iteration:
206206 16 transforms, 8 coefficient-space shuffles, 25 generated kernels, and 2
207207 buffer copies feeding the new state back.
208208
209+**Transforms batch.** The expensive part of every Legendre stage is
210+generating the associated Legendre values on the fly by recurrence — work
211+that depends only on the grid, not on the field. `synth`/`analys` therefore
212+take multiple fields, and a grouped call runs as one batched dispatch: one
213+walk of the recurrence, one accumulator lane per field —
214+
215+```matlab
216+[Ftu, Fpu, Ftv, Fpv] = synth(vtu, vpu, vtv, vpv); % one Legendre dispatch
217+```
218+
219+The grouping is a promise of independence, never of a lane width: the
220+planner ([`src/mgpu/plan.ts`](src/mgpu/plan.ts), `materializeTransforms`)
221+chunks each group into whatever the device supports — one ×4 batch under the
222+default WebGPU limits, or scalar dispatches with `SHT_BATCH=0` for A/B — so
223+the same source runs anywhere. Ungrouped transforms that happen to sit on
224+consecutive independent lines are batched the same way. Per-lane arithmetic
225+is identical to the scalar kernels', so batched and scalar plans produce
226+bit-identical states, asserted in the tests along with compile-time refusal
227+of a group that drops one of its outputs. All 16 transforms of the step
228+above land in batches, worth ~25% of the whole step (0.88 vs 1.14 ms/step at
229+lmax 127, 2 iterations, on bumpy).
230+
209231 Two consequences carried over:
210232
211233 - **The step is synchronous.** WebGPU's encode path is synchronous and every
@@ -382,7 +404,9 @@ stops folding, the results stay correct while every operator becomes its own
382404 dispatch, which is invisible in the numbers.
383405
384406 [`test/transformChecks.ts`](test/transformChecks.ts) compares the WGSL transforms
385-against shtns-webgpu's f64 CPU twin.
407+against shtns-webgpu's f64 CPU twin, and holds every compiled batch width to
408+the scalar transforms lane by lane; a model run with `SHT_BATCH=0` must
409+reproduce the batched run's state exactly.
386410
387411 - `npm run test:node` — under Dawn on the desktop, via `vite-node`. Needs a GPU;
388412 `--skip-without-gpu` lets a machine without one say so and move on (which is
models/allencahn.mmodified+8−5View file
@@ -18,14 +18,17 @@ function [Un, u] = step(U, lam, filt, gx, gy, gz, p1, p2, q2, r, eps2, dt, niter
1818 for k = 1:niter
1919 % dlap = lap_g - lap_s, evaluated at the current iterate in flux form
2020 % (see models/schnakenberg.m, docs/richardson-iteration.md and
21- % docs/reduced-transforms.md for the derivation).
21+ % docs/reduced-transforms.md for the derivation; the grouped calls run
22+ % the gradient syntheses and the flux analyses as batched dispatches).
2223 Fu = Un .* filt;
23- Ftu = synth(dthetac(Fu));
24- Fpu = synth(dphic(Fu));
24+ vtu = dthetac(Fu);
25+ vpu = dphic(Fu);
26+ [Ftu, Fpu] = synth(vtu, vpu);
2527 Pu = p1 .* Ftu + p2 .* Fpu;
2628 Qu = p2 .* Ftu + q2 .* Fpu;
27- Pcu = analys(Pu) .* filt;
28- Qcu = analys(Qu) .* filt;
29+ [PAu, QAu] = analys(Pu, Qu);
30+ Pcu = PAu .* filt;
31+ Qcu = QAu .* filt;
2932 scu = dthetac(Pcu) + dphic(Qcu);
3033 lapu = r .* synth(scu);
3134 dLu = analys(lapu) + lam .* Un;
models/brusselator.mmodified+29−25View file
@@ -3,22 +3,23 @@
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, including the grouped transforms:
7+% [a, b] = synth(x, y) runs the group as batched Legendre dispatches.
78
89 function [U, V, u, v] = init(noise, A, B)
9- U = analys(A + noise);
10- V = analys((B / A) * ones(numel(noise), 1));
11- u = synth(U);
12- v = synth(V);
10+ [U, V] = analys(A + noise, (B / A) * ones(numel(noise), 1));
11+ [u, v] = synth(U, V);
1312 end
1413
1514 function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, p1, p2, q2, r, A, B, D1, D2, dt, niter)
16- u = synth(U);
17- v = synth(V);
15+ [u, v] = synth(U, V);
1816 uuv = u .* u .* v;
1917
20- Bu = U + dt * analys(A - (B + 1) * u + uuv);
21- Bv = V + dt * analys(B * u - uuv);
18+ ru = A - (B + 1) * u + uuv;
19+ rv = B * u - uuv;
20+ [Ru, Rv] = analys(ru, rv);
21+ Bu = U + dt * Ru;
22+ Bv = V + dt * Rv;
2223
2324 Un = Bu ./ (1 + (dt * D1) * lam);
2425 Vn = Bv ./ (1 + (dt * D2) * lam);
@@ -26,28 +27,31 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, p1, p2, q2, r, A, B,
2627 for k = 1:niter
2728 % dlap = lap_g - lap_s, evaluated at the current iterate in flux form
2829 % (see models/schnakenberg.m, docs/richardson-iteration.md and
29- % docs/reduced-transforms.md for the derivation).
30+ % docs/reduced-transforms.md for the derivation and the ordering).
3031 Fu = Un .* filt;
31- Ftu = synth(dthetac(Fu));
32- Fpu = synth(dphic(Fu));
32+ Fv = Vn .* filt;
33+ vtu = dthetac(Fu);
34+ vpu = dphic(Fu);
35+ vtv = dthetac(Fv);
36+ vpv = dphic(Fv);
37+ [Ftu, Fpu, Ftv, Fpv] = synth(vtu, vpu, vtv, vpv);
3338 Pu = p1 .* Ftu + p2 .* Fpu;
3439 Qu = p2 .* Ftu + q2 .* Fpu;
35- Pcu = analys(Pu) .* filt;
36- Qcu = analys(Qu) .* filt;
37- scu = dthetac(Pcu) + dphic(Qcu);
38- lapu = r .* synth(scu);
39- dLu = analys(lapu) + lam .* Un;
40-
41- Fv = Vn .* filt;
42- Ftv = synth(dthetac(Fv));
43- Fpv = synth(dphic(Fv));
4440 Pv = p1 .* Ftv + p2 .* Fpv;
4541 Qv = p2 .* Ftv + q2 .* Fpv;
46- Pcv = analys(Pv) .* filt;
47- Qcv = analys(Qv) .* filt;
42+ [PAu, QAu, PAv, QAv] = analys(Pu, Qu, Pv, Qv);
43+ Pcu = PAu .* filt;
44+ Qcu = QAu .* filt;
45+ Pcv = PAv .* filt;
46+ Qcv = QAv .* filt;
47+ scu = dthetac(Pcu) + dphic(Qcu);
4848 scv = dthetac(Pcv) + dphic(Qcv);
49- lapv = r .* synth(scv);
50- dLv = analys(lapv) + lam .* Vn;
49+ [Lu, Lv] = synth(scu, scv);
50+ lapu = r .* Lu;
51+ lapv = r .* Lv;
52+ [LAu, LAv] = analys(lapu, lapv);
53+ dLu = LAu + lam .* Un;
54+ dLv = LAv + lam .* Vn;
5155
5256 Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
5357 Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
models/schnakenberg.mmodified+46−34View file
@@ -18,20 +18,25 @@
1818 function [U, V, u, v] = init(noise, a, b)
1919 us = a + b;
2020 vs = b / (us * us);
21- U = analys(us + noise);
22- V = analys(vs * ones(numel(noise), 1));
23- u = synth(U);
24- v = synth(V);
21+ [U, V] = analys(us + noise, vs * ones(numel(noise), 1));
22+ [u, v] = synth(U, V);
2523 end
2624
2725 function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, p1, p2, q2, r, a, b, D1, D2, dt, niter)
28- u = synth(U);
29- v = synth(V);
26+ % Grouped transforms -- [a, b] = synth(x, y) -- are explicit batching:
27+ % output k is the transform of input k, and the whole group runs as one
28+ % batched Legendre dispatch, or as many as the device's lane width allows
29+ % (src/mgpu/plan.ts, materializeTransforms). The grouping is a promise of
30+ % independence, never of a lane width, so the same source runs anywhere.
31+ [u, v] = synth(U, V);
3032 uuv = u .* u .* v;
3133
3234 % Right-hand side of the implicit solve (I - dt*D*lap_g) Unew = B.
33- Bu = U + dt * analys(a - u + uuv);
34- Bv = V + dt * analys(b - uuv);
35+ ru = a - u + uuv;
36+ rv = b - uuv;
37+ [Ru, Rv] = analys(ru, rv);
38+ Bu = U + dt * Ru;
39+ Bv = V + dt * Rv;
3540
3641 % Round-sphere solve, then iterate the geometric correction.
3742 Un = Bu ./ (1 + (dt * D1) * lam);
@@ -39,38 +44,45 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, p1, p2, q2, r, a, b,
3944
4045 for k = 1:niter
4146 % dlap = lap_g - lap_s at the current iterate, in flux form
42- % (docs/reduced-transforms.md Sec 4). The sin-weighted
43- % derivatives sin(theta)*dtheta(u) and dphi(u) -- both smooth on the
44- % sphere, synthesized straight from the dthetac/dphic coefficient
45- % shuffles -- are combined pointwise through the precomputed weights
46- % p1,p2,q2 into two fluxes P,Q, also smooth. Their coefficients are then
47- % pushed through the *same* shuffles again and summed before the one
48- % synthesis of the divergence, which r scales into lap_g(u). The only
49- % division by sin(theta) anywhere is folded into p1,p2,q2,r at precompute
50- % time. lam.*Un adds back -lap_s(Un), since lam holds +l(l+1). filt
51- % zeroes the top two degrees, where the derivative recurrences cannot
52- % exactly represent a derivative.
47+ % (docs/reduced-transforms.md Sec 4). The sin-weighted derivatives
48+ % sin(theta)*dtheta(u) and dphi(u) -- both smooth on the sphere,
49+ % synthesized straight from the dthetac/dphic coefficient shuffles --
50+ % are combined pointwise through the precomputed weights p1,p2,q2 into
51+ % two fluxes P,Q, also smooth. Their coefficients are then pushed
52+ % through the *same* shuffles again and summed before the one synthesis
53+ % of the divergence, which r scales into lap_g(u). The only division by
54+ % sin(theta) anywhere is folded into p1,p2,q2,r at precompute time.
55+ % lam.*Un adds back -lap_s(Un), since lam holds +l(l+1). filt zeroes the
56+ % top two degrees, where the derivative recurrences cannot exactly
57+ % represent a derivative.
58+ %
59+ % The two species share each grouped call: the four gradient
60+ % syntheses, the four flux analyses, the two divergence syntheses and
61+ % the two final analyses each run as one batched dispatch.
5362 Fu = Un .* filt;
54- Ftu = synth(dthetac(Fu));
55- Fpu = synth(dphic(Fu));
63+ Fv = Vn .* filt;
64+ vtu = dthetac(Fu);
65+ vpu = dphic(Fu);
66+ vtv = dthetac(Fv);
67+ vpv = dphic(Fv);
68+ [Ftu, Fpu, Ftv, Fpv] = synth(vtu, vpu, vtv, vpv);
5669 Pu = p1 .* Ftu + p2 .* Fpu;
5770 Qu = p2 .* Ftu + q2 .* Fpu;
58- Pcu = analys(Pu) .* filt;
59- Qcu = analys(Qu) .* filt;
60- scu = dthetac(Pcu) + dphic(Qcu);
61- lapu = r .* synth(scu);
62- dLu = analys(lapu) + lam .* Un;
63-
64- Fv = Vn .* filt;
65- Ftv = synth(dthetac(Fv));
66- Fpv = synth(dphic(Fv));
6771 Pv = p1 .* Ftv + p2 .* Fpv;
6872 Qv = p2 .* Ftv + q2 .* Fpv;
69- Pcv = analys(Pv) .* filt;
70- Qcv = analys(Qv) .* filt;
73+ [PAu, QAu, PAv, QAv] = analys(Pu, Qu, Pv, Qv);
74+ Pcu = PAu .* filt;
75+ Qcu = QAu .* filt;
76+ Pcv = PAv .* filt;
77+ Qcv = QAv .* filt;
78+ scu = dthetac(Pcu) + dphic(Qcu);
7179 scv = dthetac(Pcv) + dphic(Qcv);
72- lapv = r .* synth(scv);
73- dLv = analys(lapv) + lam .* Vn;
80+ [Lu, Lv] = synth(scu, scv);
81+ lapu = r .* Lu;
82+ lapv = r .* Lv;
83+ [LAu, LAv] = analys(lapu, lapv);
84+ dLu = LAu + lam .* Un;
85+ dLv = LAv + lam .* Vn;
7486
7587 Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
7688 Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
src/mgpu/compile.tsmodified+8−3View file
@@ -215,13 +215,16 @@ function forLoops(stmts: IRStmt[]): For[] {
215215 return out;
216216 }
217217
218-/** cNames assigned anywhere in a statement list, including inside loops. */
218+/** cNames assigned anywhere in a statement list, including inside loops.
219+ * A MultiAssignCall assigns every bound output slot. */
219220 function assignedCNames(stmts: IRStmt[]): Set<string> {
220221 const out = new Set<string>();
221222 const walk = (list: IRStmt[]): void => {
222223 for (const s of list) {
223224 if (s.kind === 'Assign') out.add(s.cName);
224- else if (s.kind === 'For') walk(s.body);
225+ else if (s.kind === 'MultiAssignCall') {
226+ for (const o of s.outputs) if (o.binding) out.add(o.binding.cName);
227+ } else if (s.kind === 'For') walk(s.body);
225228 }
226229 };
227230 walk(stmts);
@@ -280,7 +283,9 @@ function checkLoopEscapes(fn: IRFunc, loop: For, assignedBefore: Set<string>): v
280283 for (const s of list) {
281284 if (s === (loop as IRStmt)) continue; // the loop's own body is not "outside"
282285 if (s.kind === 'Assign') forEachVarRead(s.expr, (c) => readOutside.add(c));
283- else if (s.kind === 'For') walk(s.body);
286+ else if (s.kind === 'MultiAssignCall') {
287+ for (const a of s.args) forEachVarRead(a, (c) => readOutside.add(c));
288+ } else if (s.kind === 'For') walk(s.body);
284289 }
285290 };
286291 walk(fn.body);
src/mgpu/externals.tsmodified+44−14View file
@@ -35,36 +35,66 @@ const numericType = (rows: number, cols: number): string =>
3535 const dim = (n: number): string =>
3636 n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;
3737
38-/** Source for one transform's `.mtoc2.js`. */
38+/** Source for one transform's `.mtoc2.js`. With `multi`, the op maps each of
39+ * N inputs to its own output — `[a, b] = synth(x, y)` — so the backend can
40+ * run the group as one batched Legendre dispatch (or split it to whatever
41+ * lane width the device supports; the syntax promises grouping intent, not
42+ * a width). */
3943 function transformSource(
4044 name: string,
4145 inRows: number,
4246 inCols: number,
4347 outRows: number,
4448 outCols: number,
49+ multi = false,
4550 ): string {
4651 return `
4752 exports.name = ${JSON.stringify(name)};
4853
4954 exports.transfer = function (argTypes, nargout) {
50- if (argTypes.length !== 1) {
55+ ${
56+ multi
57+ ? `if (argTypes.length < 1) {
58+ throw new Error("${name} takes at least one argument");
59+ }
60+ if (nargout > 1 && nargout !== argTypes.length) {
61+ throw new Error(
62+ "${name}: each input produces one output, so " + argTypes.length +
63+ " input(s) return " + argTypes.length + " output(s), but " +
64+ nargout + " were requested -- write [a, b] = ${name}(x, y)"
65+ );
66+ }
67+ if (nargout <= 1 && argTypes.length !== 1) {
68+ throw new Error(
69+ "${name}: " + argTypes.length + " inputs produce " + argTypes.length +
70+ " outputs -- bind each one: [a, b] = ${name}(x, y)"
71+ );
72+ }`
73+ : `if (argTypes.length !== 1) {
5174 throw new Error("${name} takes exactly one argument, got " + argTypes.length);
5275 }
5376 if (nargout > 1) {
5477 throw new Error("${name} returns one value, but " + nargout + " were requested");
78+ }`
5579 }
56- var a = argTypes[0];
57- if (!a || a.kind !== "Numeric" || a.isComplex) {
58- throw new Error("${name} requires a real numeric array");
80+ for (var i = 0; i < argTypes.length; i++) {
81+ var a = argTypes[i];
82+ if (!a || a.kind !== "Numeric" || a.isComplex) {
83+ throw new Error("${name} requires real numeric arrays (argument " + (i + 1) + ")");
84+ }
85+ var s = a.shape;
86+ if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
87+ throw new Error(
88+ "${name} requires ${inRows}x${inCols} arrays, argument " + (i + 1) +
89+ " is " + (s ? s.join("x") : "unknown shape")
90+ );
91+ }
5992 }
60- var s = a.shape;
61- if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
62- throw new Error(
63- "${name} requires a ${inRows}x${inCols} array, got " +
64- (s ? s.join("x") : "unknown shape")
65- );
93+ var out = [];
94+ for (var k = 0; k < Math.max(1, nargout); k++) {
95+ out.push(${numericType(outRows, outCols)});
6696 }
67- return [${numericType(outRows, outCols)}];
97+ return out;
6898 };
6999
70100 // Never called: this project executes the IR on WebGPU and emits no C.
@@ -93,11 +123,11 @@ export function externalOpFiles(g: GridSizes): { name: string; source: string }[
93123 return [
94124 {
95125 name: 'synth.mtoc2.js',
96- source: transformSource('synth', 2, g.nlm, g.npts, 1),
126+ source: transformSource('synth', 2, g.nlm, g.npts, 1, true),
97127 },
98128 {
99129 name: 'analys.mtoc2.js',
100- source: transformSource('analys', g.npts, 1, 2, g.nlm),
130+ source: transformSource('analys', g.npts, 1, 2, g.nlm, true),
101131 },
102132 {
103133 name: 'dtheta.mtoc2.js',
src/mgpu/numbl.d.tsmodified+22−2View file
@@ -137,16 +137,36 @@ declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
137137 body: IRStmt[];
138138 span: Span;
139139 }
140+ /**
141+ * Multi-output call statement: `[a, b] = f(x, y)`. For `isBuiltin: true`
142+ * the builtin's `transfer(argTypes, nargout)` typed the slots during
143+ * lowering; args arrive ANF'd. The planner accepts this only for the
144+ * batched transforms (`synth`/`analys`), where output k is the transform
145+ * of argument k.
146+ */
147+ export interface MultiAssignCall {
148+ kind: 'MultiAssignCall';
149+ cName: string;
150+ name: string;
151+ isBuiltin?: boolean;
152+ args: IRExpr[];
153+ outputs: ReadonlyArray<{
154+ ty: Type;
155+ binding: { name: string; cName: string } | null;
156+ }>;
157+ span: Span;
158+ }
159+
140160 /** Any other IR statement kind — rejected by the planner. */
141161 export interface OtherStmt {
142162 kind:
143163 | 'ExprStmt' | 'If' | 'While' | 'ReturnFromFunction' | 'Break'
144- | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
164+ | 'Continue' | 'TypeComment' | 'MemberStore'
145165 | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
146166 span: Span;
147167 }
148168
149- export type IRStmt = Assign | For | OtherStmt;
169+ export type IRStmt = Assign | For | MultiAssignCall | OtherStmt;
150170
151171 export interface IRFunc {
152172 name: string;
src/mgpu/plan.tsmodified+213−19View file
@@ -9,9 +9,15 @@
99 * encoded into one submit and keeps the CPU out of the loop.
1010 */
1111 import { isMultiElement, scalarDouble } from 'numbl-src/numbl-core/jit/lowering/types.ts';
12-import type { Assign, For, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
12+import type {
13+ Assign,
14+ For,
15+ IRExpr,
16+ IRStmt,
17+ MultiAssignCall,
18+} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
1319 import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
14-import { ShtPlan, type ShtBinding } from '../sht/sht.ts';
20+import { ShtPlan, type ShtBinding, type ShtBatchBinding } from '../sht/sht.ts';
1521 import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
1622 import type { CompiledFunction } from './compile.ts';
1723 import { EXTERNAL_OPS } from './externals.ts';
@@ -120,10 +126,106 @@ type Op =
120126 copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
121127 }
122128 | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
129+ | { kind: 'synth-batch' | 'analys-batch'; binding: ShtBatchBinding; labels: string[] }
123130 | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
124131 | { kind: 'dthetac' | 'dphic'; bindGroup: GPUBindGroup; label: string }
125132 | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
126133
134+/**
135+ * A transform op as planned, before bindings exist: `in`/`out` are the
136+ * caller-side buffers (spectral in / grid out for synth, the reverse for
137+ * analys). Kept unbound until every statement is planned so that adjacent
138+ * independent transforms of the same kind can be grouped into one batched
139+ * dispatch (ShtPlan.createSynthBatchBinding) — the Legendre recurrence is
140+ * the expensive shared part, and a batch walks it once for all lanes.
141+ */
142+interface PendingSht {
143+ pending: true;
144+ kind: 'synth' | 'analys';
145+ in: GPUBuffer;
146+ out: GPUBuffer;
147+ label: string;
148+}
149+
150+type Planned = Op | PendingSht;
151+
152+const isPending = (op: Planned): op is PendingSht => 'pending' in op;
153+
154+/**
155+ * Group maximal runs of adjacent same-kind transforms into batches of the
156+ * widest compiled lane count, and create all bindings. Only literal
157+ * adjacency in the op sequence is batched — no reordering — so the models
158+ * are written to keep batchable transforms consecutive (see the solve loops
159+ * in models/*.m). Batching changes dispatch shape only: per-lane arithmetic
160+ * is identical to the scalar kernels', so results do not depend on batchK.
161+ */
162+function materializeTransforms(planned: Planned[], sht: ShtPlan): Op[] {
163+ /** Lanes must not collide: distinct outputs, and no lane reading another's
164+ * output (repeated read-only inputs would be harmless, but WebGPU also
165+ * forbids aliasing a writable binding, so outputs are the hard rule). */
166+ const disjoint = (members: PendingSht[]): boolean => {
167+ const outs = new Set<GPUBuffer>();
168+ for (const m of members) {
169+ if (outs.has(m.out)) return false;
170+ outs.add(m.out);
171+ }
172+ return members.every((m) => !outs.has(m.in));
173+ };
174+ const bind = (m: PendingSht): Op =>
175+ m.kind === 'synth'
176+ ? { kind: 'synth', binding: sht.createSynthBinding(m.in, m.out), label: m.label }
177+ : { kind: 'analys', binding: sht.createAnalysBinding(m.in, m.out), label: m.label };
178+ const bindBatch = (members: PendingSht[]): Op =>
179+ members[0].kind === 'synth'
180+ ? {
181+ kind: 'synth-batch',
182+ binding: sht.createSynthBatchBinding(
183+ members.map((m) => ({ qlmIn: m.in, spatOut: m.out })),
184+ ),
185+ labels: members.map((m) => m.label),
186+ }
187+ : {
188+ kind: 'analys-batch',
189+ binding: sht.createAnalysBatchBinding(
190+ members.map((m) => ({ spatIn: m.in, qlmOut: m.out })),
191+ ),
192+ labels: members.map((m) => m.label),
193+ };
194+
195+ const out: Op[] = [];
196+ let i = 0;
197+ while (i < planned.length) {
198+ const op = planned[i];
199+ if (!isPending(op)) {
200+ out.push(op);
201+ i++;
202+ continue;
203+ }
204+ let j = i;
205+ while (j < planned.length) {
206+ const p = planned[j];
207+ if (!isPending(p) || p.kind !== op.kind) break;
208+ j++;
209+ }
210+ const run = planned.slice(i, j) as PendingSht[];
211+ let s = 0;
212+ while (s < run.length) {
213+ let take = 1;
214+ for (const K of [4, 2]) {
215+ if (K > sht.batchK || s + K > run.length) continue;
216+ if (disjoint(run.slice(s, s + K))) {
217+ take = K;
218+ break;
219+ }
220+ }
221+ out.push(take === 1 ? bind(run[s]) : bindBatch(run.slice(s, s + take)));
222+ s += take;
223+ }
224+ i = j;
225+ }
226+ return out;
227+}
228+
127229 export interface PlanSpec {
128230 /** The specialized function this plan executes. */
129231 fn: CompiledFunction;
@@ -272,7 +374,7 @@ export class ModelPlan {
272374 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
273375 });
274376
275- const ops: Op[] = [];
377+ const planned: Planned[] = [];
276378 for (const stmt of fn.body) {
277379 await planStatement(stmt);
278380 }
@@ -295,7 +397,7 @@ export class ModelPlan {
295397 `'${to}' (${dst.count})`,
296398 );
297399 }
298- ops.push({
400+ planned.push({
299401 kind: 'copy',
300402 from: src.buffer,
301403 to: dst.buffer,
@@ -304,6 +406,10 @@ export class ModelPlan {
304406 });
305407 });
306408
409+ // Group adjacent independent transforms into batched dispatches and
410+ // create every binding.
411+ const ops = materializeTransforms(planned, sht);
412+
307413 return new ModelPlan({
308414 device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
309415 });
@@ -311,6 +417,7 @@ export class ModelPlan {
311417 async function planStatement(stmt: IRStmt): Promise<void> {
312418 if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
313419 if (stmt.kind === 'For') return planFor(stmt);
420+ if (stmt.kind === 'MultiAssignCall') return planMultiTransform(stmt);
314421 if (stmt.kind !== 'Assign') {
315422 throw new UnsupportedOnGpu(
316423 `a model function body may only contain assignments ` +
@@ -357,16 +464,14 @@ export class ModelPlan {
357464 );
358465 }
359466 const label = `${stmt.name} = ${ext.name}(${ext.argName})`;
360- if (ext.name === 'synth') {
361- ops.push({
362- kind: 'synth',
363- binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
364- label,
365- });
366- } else if (ext.name === 'analys') {
367- ops.push({
368- kind: 'analys',
369- binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
467+ if (ext.name === 'synth' || ext.name === 'analys') {
468+ // Left unbound until materializeTransforms has grouped adjacent
469+ // independent transforms into batched dispatches.
470+ planned.push({
471+ pending: true,
472+ kind: ext.name,
473+ in: argSlot.buffer,
474+ out: dest.buffer,
370475 label,
371476 });
372477 } else if (
@@ -393,7 +498,7 @@ export class ModelPlan {
393498 stmt.span,
394499 );
395500 }
396- ops.push({
501+ planned.push({
397502 kind: ext.name,
398503 bindGroup:
399504 ext.name === 'dthetac'
@@ -403,7 +508,7 @@ export class ModelPlan {
403508 });
404509 return;
405510 }
406- ops.push(
511+ planned.push(
407512 ext.name === 'dtheta'
408513 ? { kind: 'dtheta', binding: deriv.createDthetaBinding(argSlot.buffer, dest.buffer), label }
409514 : { kind: 'dphi', binding: deriv.createDphiBinding(argSlot.buffer, dest.buffer), label },
@@ -458,7 +563,7 @@ export class ModelPlan {
458563 }
459564 entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
460565
461- ops.push({
566+ planned.push({
462567 kind: 'kernel',
463568 pipeline,
464569 bindGroup: device.createBindGroup({
@@ -473,6 +578,75 @@ export class ModelPlan {
473578 });
474579 }
475580
581+ /**
582+ * `[a, b] = synth(x, y)` / `[a, b] = analys(x, y)`: an explicitly grouped
583+ * transform — output k is the transform of argument k. The group is
584+ * planned as consecutive pending transforms, which materializeTransforms
585+ * then chunks into whatever batched dispatch widths the device supports
586+ * (one x4 batch, two x2, or scalars with SHT_BATCH=0) — the syntax
587+ * promises grouping intent, never a lane width, so the same source
588+ * compiles everywhere.
589+ */
590+ function planMultiTransform(stmt: MultiAssignCall): void {
591+ if (stmt.name !== 'synth' && stmt.name !== 'analys') {
592+ throw new UnsupportedOnGpu(
593+ `'${stmt.name}' does not return multiple values here — only the ` +
594+ `transforms ('synth', 'analys') support [a, b] = op(x, y) grouping`,
595+ stmt.span,
596+ );
597+ }
598+ const kind = stmt.name;
599+ for (let i = 0; i < stmt.outputs.length; i++) {
600+ const slot = stmt.outputs[i];
601+ const arg = stmt.args[i];
602+ if (!slot.binding) {
603+ throw new UnsupportedOnGpu(
604+ `every output of '${kind}' must be bound to a name — output ` +
605+ `${i + 1} is dropped, but each input costs a transform`,
606+ stmt.span,
607+ );
608+ }
609+ if (!arg || arg.kind !== 'Var') {
610+ throw new UnsupportedOnGpu(
611+ `'${kind}' must be applied to variables (argument ${i + 1})`,
612+ stmt.span,
613+ );
614+ }
615+ const argSlot = slots.get(arg.cName);
616+ if (!argSlot) {
617+ throw new UnsupportedOnGpu(
618+ `'${kind}' reads '${arg.name}', which has no buffer`,
619+ stmt.span,
620+ );
621+ }
622+ if (!isNumeric(slot.ty) || !isTensor(slot.ty)) {
623+ throw new UnsupportedOnGpu(
624+ `'${slot.binding.name}' is not a numeric array`,
625+ stmt.span,
626+ );
627+ }
628+ const count = numel(slot.ty);
629+ let dest = slots.get(slot.binding.cName);
630+ if (!dest) {
631+ dest = alloc(`mgpu-${slot.binding.name}`, count);
632+ slots.set(slot.binding.cName, dest);
633+ } else if (dest.count !== count) {
634+ throw new UnsupportedOnGpu(
635+ `'${slot.binding.name}' changes size between assignments`,
636+ stmt.span,
637+ );
638+ }
639+ byName.set(slot.binding.name, dest);
640+ planned.push({
641+ pending: true,
642+ kind,
643+ in: argSlot.buffer,
644+ out: dest.buffer,
645+ label: `${slot.binding.name} = ${kind}(${arg.name})`,
646+ });
647+ }
648+ }
649+
476650 /**
477651 * Unroll a counted loop into the op sequence.
478652 *
@@ -605,6 +779,12 @@ export class ModelPlan {
605779 case 'dphic':
606780 this.#deriv!.encodeDphicInto(inPass(), op.bindGroup);
607781 break;
782+ case 'synth-batch':
783+ this.#sht.encodeSynthBatchInto(inPass(), op.binding);
784+ break;
785+ case 'analys-batch':
786+ this.#sht.encodeAnalysBatchInto(inPass(), op.binding);
787+ break;
608788 case 'copy':
609789 endPass();
610790 encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
@@ -627,9 +807,23 @@ export class ModelPlan {
627807 else this.#deriv!.encodeDphiInto(pass, op.binding);
628808 }
629809
630- /** Human-readable op sequence — what the .m actually compiled to. */
810+ /**
811+ * Human-readable op sequence — what the .m actually compiled to. Batched
812+ * transforms list one line per lane, annotated: the line count equals the
813+ * logical op count regardless of the device's batch width, so op-count
814+ * assertions in the tests are batch-invariant.
815+ */
631816 describe(): string[] {
632- return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
817+ return this.#ops.flatMap((op) => {
818+ if ('labels' in op) {
819+ const kind = op.kind === 'synth-batch' ? 'synth' : 'analys';
820+ return op.labels.map(
821+ (label, i) =>
822+ `${kind.padEnd(7)} ${label} [batch lane ${i + 1}/${op.binding.size}]`,
823+ );
824+ }
825+ return [`${op.kind.padEnd(7)} ${op.label}`];
826+ });
633827 }
634828
635829 destroy(): void {
src/sht/sht.tsmodified+186−2View file
@@ -9,7 +9,12 @@
99 import { gaussNodesWeights } from './gauss.ts';
1010 import { legendreCoeffs } from './coeffs.ts';
1111 import { nlmCalc, validateConfig, isPowerOfTwo, type ShtConfig } from './layout.ts';
12-import { legSynthWGSL, legAnalysWGSL } from './wgsl/leg.ts';
12+import {
13+ legSynthWGSL,
14+ legAnalysWGSL,
15+ legSynthBatchWGSL,
16+ legAnalysBatchWGSL,
17+} from './wgsl/leg.ts';
1318 import {
1419 fftSynthWGSL,
1520 fftAnalysWGSL,
@@ -28,6 +33,20 @@ export interface ShtBinding {
2833 readonly bgFour: GPUBindGroup;
2934 }
3035
36+/**
37+ * One batched transform: K fields through a single Legendre dispatch (the
38+ * recurrence walked once, K accumulator lanes) plus K per-field Fourier
39+ * dispatches — the Fourier stage shares nothing across fields, so batching
40+ * it would save only bind-group switches.
41+ */
42+export interface ShtBatchBinding {
43+ /** Lanes in this batch — selects the pipeline compiled for that width. */
44+ readonly size: number;
45+ readonly bgLeg: GPUBindGroup;
46+ /** Per-lane Fourier bind group, lane k against fm arena k. */
47+ readonly bgFour: GPUBindGroup[];
48+}
49+
3150 const bgEntries = (bufs: GPUBuffer[]) =>
3251 bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
3352
@@ -119,6 +138,16 @@ export class ShtPlan {
119138 readonly fourierMode: 'fft' | 'dft';
120139 /** Latitudes leg_synth walks: nlat/2 when parity folding. */
121140 readonly legLat: number = 0;
141+ /**
142+ * Widest transform batch this plan supports: the largest even K <= 4 whose
143+ * Legendre bind group (3 tables + K caller fields + the shared fm arena)
144+ * fits the device's storage-buffer limit. K = 4 needs exactly the WebGPU
145+ * default of 8, so batching is fully available on every stack; 1 (no
146+ * batching) if SHT_BATCH is turned off. Batched and scalar transforms
147+ * compute identical per-lane arithmetic, so this only affects speed,
148+ * never results.
149+ */
150+ readonly batchK: number = 1;
122151 /** Colatitudes theta_i (f64, increasing: north to south). */
123152 readonly theta: Float64Array;
124153 readonly cosTheta: Float64Array;
@@ -146,6 +175,16 @@ export class ShtPlan {
146175 private pipeLegAnalys!: GPUComputePipeline;
147176 private pipeFourSynth!: GPUComputePipeline;
148177 private pipeFourAnalys!: GPUComputePipeline;
178+ /** Batched Legendre pipelines by lane count (even sizes up to batchK). */
179+ private pipeLegSynthB = new Map<number, GPUComputePipeline>();
180+ private pipeLegAnalysB = new Map<number, GPUComputePipeline>();
181+ /** One fm arena for all batch lanes (lane k at byte offset k * fmLaneBytes,
182+ * 256-aligned so the Fourier stage can bind a lane by buffer offset). A
183+ * single buffer keeps the batched Legendre bind group at 3 tables +
184+ * K fields + 1 arena — within WebGPU's default storage-buffer limit of 8
185+ * at K = 4, on every stack. */
186+ private fmArena: GPUBuffer | null = null;
187+ private fmLaneBytes = 0;
149188 private bgLegSynth!: GPUBindGroup;
150189 private bgLegAnalys!: GPUBindGroup;
151190 private bgFourSynth!: GPUBindGroup;
@@ -269,6 +308,41 @@ export class ShtPlan {
269308 this.pipeFourSynth = pFourS;
270309 this.pipeFourAnalys = pFourA;
271310
311+ // --- batched Legendre pipelines ---
312+ // The widest even K <= 4 whose bind group (3 tables + K fields + the fm
313+ // arena) fits the device's storage-buffer budget — K = 4 needs 8, the
314+ // WebGPU default, so batching is fully available everywhere unless
315+ // SHT_BATCH=0 disables it (SHT_BATCH=2 caps it, for A/B).
316+ const batchTuning = tuning('SHT_BATCH');
317+ const batchWant =
318+ batchTuning === false || batchTuning === 0
319+ ? 1
320+ : typeof batchTuning === 'number'
321+ ? batchTuning
322+ : 4;
323+ const batchFit = dev.limits.maxStorageBuffersPerShaderStage - 4;
324+ const batchK = Math.min(4, Math.max(1, batchWant), 2 * Math.floor(batchFit / 2));
325+ (this as { batchK: number }).batchK = batchK;
326+ if (batchK >= 2) {
327+ // Lane stride rounded to the 256-byte offset alignment buffer bindings
328+ // require; laneElems is that stride in vec2f units for the kernels.
329+ this.fmLaneBytes = Math.ceil((8 * (mmax + 1) * nlat) / 256) * 256;
330+ const laneElems = this.fmLaneBytes / 8;
331+ this.fmArena = mkBuf('sht-fm-arena', batchK * this.fmLaneBytes, GPUBufferUsage.STORAGE);
332+ const sizes = [];
333+ for (let k = 2; k <= batchK; k += 2) sizes.push(k);
334+ const pipes = await Promise.all(
335+ sizes.flatMap((k) => [
336+ makePipeline(dev, legSynthBatchWGSL(legP, k, laneElems), `leg_synth_batch`),
337+ makePipeline(dev, legAnalysBatchWGSL(legP, k, laneElems), `leg_analys_batch`),
338+ ]),
339+ );
340+ sizes.forEach((k, i) => {
341+ this.pipeLegSynthB.set(k, pipes[2 * i]);
342+ this.pipeLegAnalysB.set(k, pipes[2 * i + 1]);
343+ });
344+ }
345+
272346 const entries = bgEntries;
273347 this.bgLegSynth = dev.createBindGroup({
274348 layout: pLegS.getBindGroupLayout(0),
@@ -322,6 +396,116 @@ export class ShtPlan {
322396 };
323397 }
324398
399+ /**
400+ * Bind groups for one batched synthesis: members.length must be a compiled
401+ * lane count (an even size <= batchK). Member outputs must be distinct
402+ * buffers; each lane gets its own fm arena, so batches compose in a pass
403+ * exactly like sequential scalar transforms do.
404+ */
405+ /** The fm arena sliced at lane k, sized as one transform's fm. */
406+ #fmLane(k: number): GPUBufferBinding {
407+ const { mmax, nlat } = this.cfg;
408+ return {
409+ buffer: this.fmArena!,
410+ offset: k * this.fmLaneBytes,
411+ size: 8 * (mmax + 1) * nlat,
412+ };
413+ }
414+
415+ createSynthBatchBinding(
416+ members: { qlmIn: GPUBuffer; spatOut: GPUBuffer }[],
417+ ): ShtBatchBinding {
418+ const K = members.length;
419+ const pipe = this.pipeLegSynthB.get(K);
420+ if (!pipe) throw new Error(`no batched synthesis pipeline for ${K} lanes`);
421+ return {
422+ size: K,
423+ bgLeg: this.device.createBindGroup({
424+ layout: pipe.getBindGroupLayout(0),
425+ entries: [
426+ ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, ...members.map((m) => m.qlmIn)]),
427+ { binding: 3 + K, resource: { buffer: this.fmArena! } },
428+ ],
429+ }),
430+ bgFour: members.map((m, k) =>
431+ this.device.createBindGroup({
432+ layout: this.pipeFourSynth.getBindGroupLayout(0),
433+ entries: [
434+ { binding: 0, resource: this.#fmLane(k) },
435+ { binding: 1, resource: { buffer: m.spatOut } },
436+ { binding: 2, resource: { buffer: this.bufTrig } },
437+ ],
438+ }),
439+ ),
440+ };
441+ }
442+
443+ createAnalysBatchBinding(
444+ members: { spatIn: GPUBuffer; qlmOut: GPUBuffer }[],
445+ ): ShtBatchBinding {
446+ const K = members.length;
447+ const pipe = this.pipeLegAnalysB.get(K);
448+ if (!pipe) throw new Error(`no batched analysis pipeline for ${K} lanes`);
449+ return {
450+ size: K,
451+ bgFour: members.map((m, k) =>
452+ this.device.createBindGroup({
453+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
454+ entries: [
455+ { binding: 0, resource: { buffer: m.spatIn } },
456+ { binding: 1, resource: this.#fmLane(k) },
457+ { binding: 2, resource: { buffer: this.bufTrig } },
458+ ],
459+ }),
460+ ),
461+ bgLeg: this.device.createBindGroup({
462+ layout: pipe.getBindGroupLayout(0),
463+ entries: [
464+ ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw]),
465+ { binding: 3, resource: { buffer: this.fmArena! } },
466+ ...members.map((m, k) => ({
467+ binding: 4 + k,
468+ resource: { buffer: m.qlmOut },
469+ })),
470+ ],
471+ }),
472+ };
473+ }
474+
475+ /** Record a batched synthesis: one Legendre dispatch, K Fourier dispatches. */
476+ encodeSynthBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
477+ const { mmax, nlat, nphi } = this.cfg;
478+ pass.setPipeline(this.pipeLegSynthB.get(b.size)!);
479+ pass.setBindGroup(0, b.bgLeg);
480+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
481+ pass.setPipeline(this.pipeFourSynth);
482+ for (const bg of b.bgFour) {
483+ pass.setBindGroup(0, bg);
484+ if (this.fourierMode === 'fft') {
485+ pass.dispatchWorkgroups(nlat);
486+ } else {
487+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
488+ }
489+ }
490+ }
491+
492+ /** Record a batched analysis: K Fourier dispatches, one Legendre dispatch. */
493+ encodeAnalysBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
494+ const { mmax, nlat } = this.cfg;
495+ pass.setPipeline(this.pipeFourAnalys);
496+ for (const bg of b.bgFour) {
497+ pass.setBindGroup(0, bg);
498+ if (this.fourierMode === 'fft') {
499+ pass.dispatchWorkgroups(nlat);
500+ } else {
501+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
502+ }
503+ }
504+ pass.setPipeline(this.pipeLegAnalysB.get(b.size)!);
505+ pass.setBindGroup(0, b.bgLeg);
506+ pass.dispatchWorkgroups(mmax + 1);
507+ }
508+
325509 /** Record synthesis into an existing compute pass. */
326510 encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
327511 const { mmax, nlat, nphi } = this.cfg;
@@ -463,7 +647,7 @@ export class ShtPlan {
463647 destroy(): void {
464648 for (const b of [
465649 this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
466- this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
650+ this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ, this.fmArena,
467651 ]) b?.destroy();
468652 }
469653 }
src/sht/wgsl/leg.tsmodified+336−0View file
@@ -350,3 +350,339 @@ ${
350350 }
351351 `;
352352 }
353+
354+/**
355+ * Batched transforms: K independent fields through ONE walk of the Legendre
356+ * recurrence. The recurrence state (y0, y1, rescaling) depends only on
357+ * (m, theta), never on the field, so a batch shares it and pays only the
358+ * extra data fetches and accumulators per lane — the same amortization
359+ * SHTNS's GPU backend gets from batching fields. Per-lane arithmetic is
360+ * textually identical to the scalar kernels' (same operations, same order),
361+ * so a batched transform reproduces the scalar transform's results.
362+ *
363+ * K is a codegen parameter. The bind group needs 3 tables + K inputs +
364+ * K outputs storage buffers, so K = 2 (7 bindings) fits WebGPU's default
365+ * limit of 8 on every stack including SwiftShader, and K = 4 (11) needs the
366+ * raised limit requestShtDevice asks for where the adapter offers it.
367+ * The recurrence bodies are kept in exactly the shape the scalar kernels
368+ * use — see the driver-workaround comment in legSynthWGSL before
369+ * "simplifying" either copy.
370+ */
371+export function legSynthBatchWGSL(p: LegParams, K: number, laneElems: number): string {
372+ const half = p.parity === true;
373+ const lanes = Array.from({ length: K }, (_, k) => k);
374+ // K caller-owned inputs, ONE plan-owned fm arena: lane k writes at a fixed
375+ // 256-byte-aligned offset (laneElems vec2f), which is what keeps the bind
376+ // group at 3 + K + 1 storage buffers -- within WebGPU's default limit of 8
377+ // at K = 4. The Fourier stage binds the arena per lane with a buffer
378+ // offset, so it needs no changes.
379+ const bind =
380+ lanes
381+ .map((k) => `@group(0) @binding(${3 + k}) var<storage, read> qlm${k}: array<vec2f>;`)
382+ .join('\n') +
383+ `\n@group(0) @binding(${3 + K}) var<storage, read_write> fm: array<vec2f>;`;
384+ const decl = lanes
385+ .map((k) =>
386+ half
387+ ? ` var accE${k} = vec2f(0.0);\n var accO${k} = vec2f(0.0);`
388+ : ` var acc${k} = vec2f(0.0);`,
389+ )
390+ .join('\n');
391+ const accEven = lanes
392+ .map((k) => (half ? ` accE${k} += y0 * qlm${k}[i0];` : ` acc${k} += y0 * qlm${k}[i0];`))
393+ .join('\n');
394+ const accOdd = lanes
395+ .map((k) => (half ? ` accO${k} += y1 * qlm${k}[i1];` : ` acc${k} += y1 * qlm${k}[i1];`))
396+ .join('\n');
397+ const store = lanes
398+ .map((k) =>
399+ half
400+ ? ` fm[${k}u * LANE + m * NLAT + ilat] = accE${k} + accO${k};\n` +
401+ ` fm[${k}u * LANE + m * NLAT + (NLAT - 1u - ilat)] = accE${k} - accO${k};`
402+ : ` fm[${k}u * LANE + m * NLAT + ilat] = acc${k};`,
403+ )
404+ .join('\n');
405+ return /* wgsl */ `
406+${RESCALE_WGSL}
407+const LMAX: u32 = ${p.lmax}u;
408+const NLAT: u32 = ${p.nlat}u;
409+const NLAT_2: u32 = ${p.nlat / 2}u;
410+const LANE: u32 = ${laneElems}u;
411+${BINDINGS}
412+${bind}
413+
414+@compute @workgroup_size(${p.wgSynth})
415+fn leg_synth_batch(@builtin(global_invocation_id) gid: vec3u,
416+ @builtin(workgroup_id) wid: vec3u) {
417+ let ilat = gid.x;
418+ let m = wid.y;
419+ if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }
420+
421+ let ct = ctstw[ilat];
422+ let st = ctstw[NLAT + ilat];
423+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
424+
425+ var seed = sinpow_rescaled(st, m);
426+ var y0 = seed.y0 * amm[m];
427+ var ny = seed.ny;
428+ var y1: f32 = 0.0;
429+ if (m < LMAX) {
430+ y1 = ab[base + 1u].x * ct * y0;
431+ }
432+
433+${decl}
434+ var l = m;
435+ loop {
436+ if (ny == 0) {
437+ let i0 = base + (l - m);
438+${accEven}
439+ if (l + 1u <= LMAX) {
440+ let i1 = base + (l + 1u - m);
441+${accOdd}
442+ }
443+ } else if (abs(y0) > RESCALE_THR) {
444+ ny += 1;
445+ y0 *= INV_SCALE;
446+ y1 *= INV_SCALE;
447+ }
448+ if (l + 2u > LMAX) { break; }
449+ // Same two-coefficient, temporary-carried shape as leg_synth (see the
450+ // driver-workaround comment there).
451+ let a0 = ab[base + (l + 2u - m)];
452+ var a1 = vec2f(0.0);
453+ if (l + 3u <= LMAX) {
454+ a1 = ab[base + (l + 3u - m)];
455+ }
456+ let t0 = a0.x * ct * y1 + a0.y * y0;
457+ y1 = a1.x * ct * t0 + a1.y * y1;
458+ y0 = t0;
459+ l += 2u;
460+ }
461+${store}
462+}
463+`;
464+}
465+
466+/** Batched analysis: K spatial-Fourier fields reduced against one Legendre
467+ * recurrence walk. Structure follows legAnalysWGSL; see legSynthBatchWGSL
468+ * for the batching rationale and the binding budget. */
469+export function legAnalysBatchWGSL(p: LegParams, K: number, laneElems: number): string {
470+ const half = p.parity === true;
471+ const lanes = Array.from({ length: K }, (_, k) => k);
472+ const Kl = Math.ceil((half ? p.nlat / 2 : p.nlat) / p.wgAnalys);
473+ const sg = p.subgroups === true;
474+ const nsubMax = Math.max(1, p.wgAnalys / 4);
475+ // Same 8 KB workgroup-storage budget as the scalar kernel, now split
476+ // across K lanes, so spans shorten as K grows: barriers per unit of work
477+ // stay level.
478+ const pairs = sg
479+ ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16 * K))))
480+ : 1;
481+ const redLen = (sg ? nsubMax * pairs : p.wgAnalys) * K;
482+ // ONE fm arena in (lane offsets baked, as in legSynthBatchWGSL), K
483+ // caller-owned outputs: 3 + 1 + K storage buffers.
484+ const bind =
485+ `@group(0) @binding(3) var<storage, read> fm: array<vec2f>;\n` +
486+ lanes
487+ .map((k) => `@group(0) @binding(${4 + k}) var<storage, read_write> qout${k}: array<vec2f>;`)
488+ .join('\n');
489+ const laneState = lanes
490+ .map((k) =>
491+ half
492+ ? ` var wpv${k}: array<vec2f, ${Kl}>;\n var wmv${k}: array<vec2f, ${Kl}>;`
493+ : ` var wfv${k}: array<vec2f, ${Kl}>;`,
494+ )
495+ .join('\n');
496+ const laneLoad = lanes
497+ .map((k) =>
498+ half
499+ ? ` let gN${k} = fm[${k}u * LANE + m * NLAT + lat];
500+ let gS${k} = fm[${k}u * LANE + m * NLAT + (NLAT - 1u - lat)];
501+ wp${k} = (gN${k} + gS${k}) * w;
502+ wm${k} = (gN${k} - gS${k}) * w;`
503+ : ` wf${k} = fm[${k}u * LANE + m * NLAT + lat] * w;`,
504+ )
505+ .join('\n');
506+ const laneLoadDecl = lanes
507+ .map((k) =>
508+ half ? ` var wp${k} = vec2f(0.0);\n var wm${k} = vec2f(0.0);` : ` var wf${k} = vec2f(0.0);`,
509+ )
510+ .join('\n');
511+ const laneLoadStore = lanes
512+ .map((k) => (half ? ` wpv${k}[k] = wp${k};\n wmv${k}[k] = wm${k};` : ` wfv${k}[k] = wf${k};`))
513+ .join('\n');
514+ const cDecl = lanes.map((k) => ` var c0_${k} = vec2f(0.0);\n var c1_${k} = vec2f(0.0);`).join('\n');
515+ const cAcc = lanes
516+ .map((k) =>
517+ half
518+ ? ` c0_${k} += wpv${k}[k] * y0v[k];
519+ c1_${k} += wmv${k}[k] * y1v[k];`
520+ : ` c0_${k} += wfv${k}[k] * y0v[k];
521+ c1_${k} += wfv${k}[k] * y1v[k];`,
522+ )
523+ .join('\n');
524+ return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
525+${RESCALE_WGSL}
526+const LMAX: u32 = ${p.lmax}u;
527+const NLAT: u32 = ${p.nlat}u;
528+const WG: u32 = ${p.wgAnalys}u;
529+const K: u32 = ${Kl}u;
530+const NLAT_2: u32 = ${p.nlat / 2}u;
531+const PAIRS: u32 = ${pairs}u;
532+const NB: u32 = ${K}u;
533+const LANE: u32 = ${laneElems}u;
534+${BINDINGS}
535+${bind}
536+
537+var<workgroup> red: array<vec4f, ${redLen}>;
538+
539+@compute @workgroup_size(${p.wgAnalys})
540+fn leg_analys_batch(@builtin(local_invocation_id) lid3: vec3u,
541+ @builtin(workgroup_id) wid: vec3u${
542+ sg
543+ ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
544+ : ''
545+ }) {
546+ let lid = lid3.x;
547+ let m = wid.x;
548+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
549+
550+ var y0v: array<f32, ${Kl}>;
551+ var y1v: array<f32, ${Kl}>;
552+ var nyv: array<i32, ${Kl}>;
553+ var ctv: array<f32, ${Kl}>;
554+${laneState}
555+
556+ for (var k = 0u; k < K; k++) {
557+ let lat = lid + k * WG;
558+ var ct: f32 = 0.0;
559+ var st: f32 = 0.0;
560+ var w: f32 = 0.0;
561+${laneLoadDecl}
562+ if (lat < ${half ? 'NLAT_2' : 'NLAT'}) {
563+ ct = ctstw[lat];
564+ st = ctstw[NLAT + lat];
565+ w = ctstw[2u * NLAT + lat];
566+${laneLoad}
567+ }
568+ ctv[k] = ct;
569+ let seed = sinpow_rescaled(st, m);
570+ y0v[k] = seed.y0 * amm[m];
571+ nyv[k] = seed.ny;
572+ y1v[k] = 0.0;
573+ if (m < LMAX) {
574+ y1v[k] = ab[base + 1u].x * ct * y0v[k];
575+ }
576+${laneLoadStore}
577+ }
578+
579+ var l = m;
580+${
581+ sg
582+ ? ` loop {
583+ let lstart = l;
584+ var npairs = 0u;
585+ var last = false;
586+ let sub = lid / sgSize;
587+ for (var jj = 0u; jj < PAIRS; jj++) {
588+${cDecl}
589+ for (var k = 0u; k < K; k++) {
590+ if (nyv[k] == 0) {
591+${cAcc}
592+ } else if (abs(y0v[k]) > RESCALE_THR) {
593+ nyv[k] += 1;
594+ y0v[k] *= INV_SCALE;
595+ y1v[k] *= INV_SCALE;
596+ }
597+ }
598+${lanes
599+ .map(
600+ (k) => ` let part${k} = subgroupAdd(vec4f(c0_${k}, c1_${k}));
601+ if (sgLane == 0u) { red[(sub * PAIRS + jj) * NB + ${k}u] = part${k}; }`,
602+ )
603+ .join('\n')}
604+ npairs = jj + 1u;
605+ if (l + 2u > LMAX) { last = true; break; }
606+ let a0 = ab[base + (l + 2u - m)];
607+ var a1 = vec2f(0.0);
608+ if (l + 3u <= LMAX) {
609+ a1 = ab[base + (l + 3u - m)];
610+ }
611+ for (var k = 0u; k < K; k++) {
612+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
613+ y0v[k] = t0;
614+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
615+ }
616+ l += 2u;
617+ }
618+
619+ workgroupBarrier();
620+ if (lid == 0u) {
621+ let nsub = (WG + sgSize - 1u) / sgSize;
622+ for (var jj = 0u; jj < npairs; jj++) {
623+ let ll = lstart + 2u * jj;
624+${lanes
625+ .map(
626+ (k) => ` var tot${k} = vec4f(0.0);
627+ for (var i = 0u; i < nsub; i++) { tot${k} += red[(i * PAIRS + jj) * NB + ${k}u]; }
628+ qout${k}[base + (ll - m)] = tot${k}.xy;
629+ if (ll + 1u <= LMAX) {
630+ qout${k}[base + (ll + 1u - m)] = tot${k}.zw;
631+ }`,
632+ )
633+ .join('\n')}
634+ }
635+ }
636+ workgroupBarrier(); // red is reused by the next span
637+
638+ if (last) { break; }
639+ }`
640+ : ` loop {
641+${cDecl}
642+ for (var k = 0u; k < K; k++) {
643+ if (nyv[k] == 0) {
644+${cAcc.replace(/^ {10}/gm, ' ')}
645+ } else if (abs(y0v[k]) > RESCALE_THR) {
646+ nyv[k] += 1;
647+ y0v[k] *= INV_SCALE;
648+ y1v[k] *= INV_SCALE;
649+ }
650+ }
651+ // workgroup tree reduction, lane-strided
652+${lanes.map((k) => ` red[lid + ${k}u * WG] = vec4f(c0_${k}, c1_${k});`).join('\n')}
653+ workgroupBarrier();
654+ var s = WG / 2u;
655+ while (s > 0u) {
656+ if (lid < s) {
657+${lanes.map((k) => ` red[lid + ${k}u * WG] += red[lid + s + ${k}u * WG];`).join('\n')}
658+ }
659+ workgroupBarrier();
660+ s = s >> 1u;
661+ }
662+ if (lid == 0u) {
663+${lanes
664+ .map(
665+ (k) => ` qout${k}[base + (l - m)] = red[${k}u * WG].xy;
666+ if (l + 1u <= LMAX) {
667+ qout${k}[base + (l + 1u - m)] = red[${k}u * WG].zw;
668+ }`,
669+ )
670+ .join('\n')}
671+ }
672+ if (l + 2u > LMAX) { break; }
673+ let a0 = ab[base + (l + 2u - m)];
674+ var a1 = vec2f(0.0);
675+ if (l + 3u <= LMAX) {
676+ a1 = ab[base + (l + 3u - m)];
677+ }
678+ for (var k = 0u; k < K; k++) {
679+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
680+ y0v[k] = t0;
681+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
682+ }
683+ l += 2u;
684+ }`
685+ }
686+}
687+`;
688+}
test/modelChecks.tsmodified+94−1View file
@@ -10,7 +10,7 @@
1010 * but every operator becomes its own dispatch, which is invisible except here.
1111 */
1212 import { ModelSession } from '../src/mgpu/session.ts';
13-import { mModels, defaultParams } from '../src/mgpu/registry.ts';
13+import { mModels, mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
1414 import {
1515 formatCommand,
1616 parseArgs,
@@ -142,6 +142,99 @@ export async function modelChecks(
142142 session.destroy();
143143 }
144144
145+ // Batched transforms are an encoding of the same arithmetic, so a run with
146+ // batching disabled (SHT_BATCH=0 compiles scalar-only plans) must reproduce
147+ // the default run to shader-compiler latitude, and the default run must
148+ // actually be batching (the describe() lines say so). This is the guard
149+ // that the planner's adjacency grouping rewires buffers correctly — a lane
150+ // bound to the wrong field would miss by O(1), not O(1e-6).
151+ {
152+ const model = mModelByKey('schnakenberg')!;
153+ const params = defaultParams(model);
154+ const states: Float32Array[] = [];
155+ let batchedLanes = 0;
156+ for (const batch of [undefined, 0]) {
157+ const g = globalThis as Record<string, unknown>;
158+ if (batch !== undefined) g.SHT_BATCH = batch;
159+ try {
160+ const session = await ModelSession.create({
161+ device, model, params, lmax: LMAX, niter: NITER,
162+ });
163+ if (batch === undefined) {
164+ batchedLanes = session
165+ .describe()
166+ .step.filter((l) => l.includes('[batch lane')).length;
167+ }
168+ session.seed(1);
169+ session.step(STEPS);
170+ states.push(await session.read('U'));
171+ session.destroy();
172+ } finally {
173+ delete g.SHT_BATCH;
174+ }
175+ }
176+ // Every batchable run at one solve iteration: the u/v syntheses and the
177+ // reaction analyses outside the loop (2 + 2), the four gradient
178+ // syntheses, four flux analyses, two divergence syntheses and two final
179+ // analyses inside it (4 + 4 + 2 + 2). Lane counts are batch-width
180+ // invariant: a x4 run is one batch at K = 4 and two at K = 2, but the
181+ // lanes annotated are the same 16 either way.
182+ check(
183+ 'batch: the compiled step batches every adjacent transform pair',
184+ batchedLanes === 16,
185+ `${batchedLanes} batched transform lanes (expected 16)`,
186+ );
187+ let worst = 0;
188+ for (let i = 0; i < states[0].length; i++) {
189+ worst = Math.max(worst, Math.abs(states[0][i] - states[1][i]));
190+ }
191+ check(
192+ 'batch: batched and scalar plans agree through a real run',
193+ worst < 1e-4,
194+ `max |U_batched - U_scalar| = ${worst.toExponential(2)} after ${STEPS} steps`,
195+ );
196+ }
197+
198+ // Misusing the grouped-transform syntax is refused at compile time with a
199+ // message that says how to write it, not silently mis-planned: every input
200+ // must get an output (each one costs a transform), whether the mismatch is
201+ // an under-bound assignment or an ignored slot.
202+ {
203+ const model = mModelByKey('allencahn')!;
204+ const cases: [string, string, string][] = [
205+ [
206+ 'a single output bound to a grouped call',
207+ 'Ftu = synth(vtu, vpu);',
208+ 'bind each one',
209+ ],
210+ [
211+ 'an ignored output slot',
212+ // Fpu is reassigned so the only error left is the dropped slot
213+ // itself, which the planner refuses (numbl would otherwise catch
214+ // the undefined 'Fpu' first, masking the check under test).
215+ '[Ftu, ~] = synth(vtu, vpu);\n Fpu = Ftu;',
216+ 'must be bound',
217+ ],
218+ ];
219+ for (const [what, bad, expect] of cases) {
220+ const source = model.source.replace('[Ftu, Fpu] = synth(vtu, vpu);', bad);
221+ let message = '';
222+ try {
223+ const session = await ModelSession.create({
224+ device, model, params: defaultParams(model), lmax: LMAX, source, niter: 1,
225+ });
226+ session.destroy();
227+ } catch (e) {
228+ message = e instanceof Error ? e.message : String(e);
229+ }
230+ check(
231+ `batch: ${what} is refused at compile time`,
232+ message.includes(expect),
233+ message ? `refused: ${message.slice(0, 76)}…` : 'compiled anyway',
234+ );
235+ }
236+ }
237+
145238 // The oversampled readback: readSpecies must be the state synthesized on the
146239 // display grid. Comparing against the display plan's own upload path
147240 // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
test/transformChecks.tsmodified+87−0View file
@@ -126,5 +126,92 @@ export async function transformChecks(
126126 deriv.destroy();
127127 }
128128
129+ // ---- batched transforms reproduce the scalar transforms ------------------
130+ // A batch walks the Legendre recurrence once for K fields with per-lane
131+ // arithmetic textually identical to the scalar kernel's, so each lane must
132+ // agree with the scalar path to shader-compiler latitude (FMA contraction
133+ // may differ between the two modules; nothing else may).
134+ {
135+ const { nlat: gl, nphi: gp } = plan.cfg;
136+ const npts = gl * gp;
137+ const sizes = [];
138+ for (let k = 2; k <= plan.batchK; k += 2) sizes.push(k);
139+ check(
140+ 'batch: plan compiled batched pipelines',
141+ plan.batchK >= 2,
142+ `batchK = ${plan.batchK} (${sizes.map((s) => `x${s}`).join(', ') || 'none'})`,
143+ );
144+ for (const K of sizes) {
145+ const qs = Array.from({ length: K }, (_, k) => randomSpectrum(cfg, 1000 + k));
146+ const qBufs = qs.map((q, k) => {
147+ const b = device.createBuffer({
148+ label: `batch-test-q${k}`,
149+ size: 8 * plan.nlm,
150+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
151+ });
152+ device.queue.writeBuffer(b, 0, q as Float32Array<ArrayBuffer>);
153+ return b;
154+ });
155+ const spatBufs = qs.map((_, k) =>
156+ device.createBuffer({
157+ label: `batch-test-spat${k}`,
158+ size: 4 * npts,
159+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
160+ }),
161+ );
162+ const qOutBufs = qs.map((_, k) =>
163+ device.createBuffer({
164+ label: `batch-test-qout${k}`,
165+ size: 8 * plan.nlm,
166+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
167+ }),
168+ );
169+ const stage = device.createBuffer({
170+ label: 'batch-test-stage',
171+ size: K * (4 * npts + 8 * plan.nlm),
172+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
173+ });
174+
175+ // One pass: batched synthesis of all K, then batched analysis back.
176+ const synthB = plan.createSynthBatchBinding(
177+ qs.map((_, k) => ({ qlmIn: qBufs[k], spatOut: spatBufs[k] })),
178+ );
179+ const analysB = plan.createAnalysBatchBinding(
180+ qs.map((_, k) => ({ spatIn: spatBufs[k], qlmOut: qOutBufs[k] })),
181+ );
182+ const enc = device.createCommandEncoder({ label: 'batch-test' });
183+ const pass = enc.beginComputePass();
184+ plan.encodeSynthBatchInto(pass, synthB);
185+ plan.encodeAnalysBatchInto(pass, analysB);
186+ pass.end();
187+ for (let k = 0; k < K; k++) {
188+ enc.copyBufferToBuffer(spatBufs[k], 0, stage, k * 4 * npts, 4 * npts);
189+ enc.copyBufferToBuffer(qOutBufs[k], 0, stage, K * 4 * npts + k * 8 * plan.nlm, 8 * plan.nlm);
190+ }
191+ device.queue.submit([enc.finish()]);
192+ await stage.mapAsync(GPUMapMode.READ);
193+ const raw = new Float32Array(stage.getMappedRange().slice(0));
194+ stage.unmap();
195+
196+ let worstSynth = 0;
197+ let worstAnalys = 0;
198+ for (let k = 0; k < K; k++) {
199+ const spatLane = raw.subarray(k * npts, (k + 1) * npts);
200+ const qLane = raw.subarray(K * npts + k * 2 * plan.nlm, K * npts + (k + 1) * 2 * plan.nlm);
201+ const spatScalar = await plan.synth(qs[k]);
202+ const qScalar = await plan.analys(spatScalar);
203+ worstSynth = Math.max(worstSynth, relL2(spatLane, spatScalar));
204+ worstAnalys = Math.max(worstAnalys, relL2(qLane, qScalar));
205+ }
206+ check(
207+ `batch: x${K} lanes match the scalar transforms`,
208+ worstSynth < 1e-6 && worstAnalys < 1e-6,
209+ `synth ${worstSynth.toExponential(2)}, analys ${worstAnalys.toExponential(2)} ` +
210+ `across ${K} lanes`,
211+ );
212+ for (const b of [...qBufs, ...spatBufs, ...qOutBufs, stage]) b.destroy();
213+ }
214+ }
215+
129216 plan.destroy();
130217 }