4f822e1turing-surface-cache: reaction-diffusion solutions at a chosen end time, shared through a cloud cacheJeremy Magland 1/**
2 * The two spherical-harmonic transforms, as external operations the .m can
3 * call: `synth` (spectral -> grid) and `analys` (grid -> spectral).
4 *
5 * numbl needs only their *type rule* in order to lower a call site. It gets
6 * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
7 * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
8 * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
9 * imports available, so `transfer` builds numbl `Type` objects as plain
10 * literals, and the grid sizes are baked in by the generator below (a grid
11 * change recompiles anyway).
12 *
13 * The `emit`/`cBody` exports exist only because the loader's contract requires
14 * them; we never emit C. The actual implementation is supplied by the WGSL
15 * backend, which turns each of these calls into an ShtPlan encode.
16 *
17 * Spectral fields are carried as REAL 2 x nlm arrays (row 0 real part, row 1
18 * imaginary), matching the interleaved layout the GPU buffers already use.
19 * The IMEX update is real-linear, so no complex arithmetic is needed.
20 */
22export interface GridSizes {
23 /** Grid points, nlat*nphi. Grid fields are npts x 1 column vectors. */
24 npts: number;
25 /** Spectral coefficients. Spectral fields are 2 x nlm. */
26 nlm: number;
27}
29const numericType = (rows: number, cols: number): string =>
30 `{ kind: "Numeric", elem: "double", isComplex: false, ` +
31 `dims: [${dim(rows)}, ${dim(cols)}], shape: [${rows}, ${cols}], sign: "unknown" }`;
33// numbl's tensorDouble() canonicalizes an extent of 1 to its shared DIM_ONE
34// singleton; mirror that so types compare equal to host-built ones.
35const dim = (n: number): string =>
36 n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;
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). */
43function transformSource(
44 name: string,
45 inRows: number,
46 inCols: number,
47 outRows: number,
48 outCols: number,
49 multi = false,
50): string {
51 return `
52exports.name = ${JSON.stringify(name)};
54exports.transfer = function (argTypes, nargout) {
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) {
74 throw new Error("${name} takes exactly one argument, got " + argTypes.length);
75 }
76 if (nargout > 1) {
77 throw new Error("${name} returns one value, but " + nargout + " were requested");
78 }`
79 }
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 }
92 }
93 var out = [];
94 for (var k = 0; k < Math.max(1, nargout); k++) {
95 out.push(${numericType(outRows, outCols)});
96 }
97 return out;
98};
100// Never called: this project executes the IR on WebGPU and emits no C.
101exports.emit = function () {
102 throw new Error("${name}: no C backend (this transform runs on WebGPU)");
103};
104exports.cBody = function () {
105 return "";
106};
107`;
108}
110/**
111 * Workspace files that make `synth` / `analys` / `dtheta` / `dphi` /
112 * `dthetac` / `dphic` resolvable during lowering. `dtheta` and `dphi` (the
113 * surface's first partial derivatives, coefficients -> grid — see
114 * src/sht/deriv.ts) have exactly `synth`'s shape rule: both take spectral
115 * coefficients and produce a grid field. `dthetac` and `dphic` are their
116 * coefficient-space halves alone — the alpha^+/alpha^- shift and the i*m
117 * multiply, spectral -> spectral — which the six-transform Laplace-Beltrami
118 * scheme (docs/reduced-transforms.md) applies twice per
119 * matvec: to the field (gradient side) and to the analysed fluxes
120 * (divergence side, the same shift, not its transpose).
121 */
122export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
123 return [
124 {
125 name: 'synth.mtoc2.js',
126 source: transformSource('synth', 2, g.nlm, g.npts, 1, true),
127 },
128 {
129 name: 'analys.mtoc2.js',
130 source: transformSource('analys', g.npts, 1, 2, g.nlm, true),
131 },
132 {
133 name: 'dtheta.mtoc2.js',
134 source: transformSource('dtheta', 2, g.nlm, g.npts, 1),
135 },
136 {
137 name: 'dphi.mtoc2.js',
138 source: transformSource('dphi', 2, g.nlm, g.npts, 1),
139 },
140 {
141 name: 'dthetac.mtoc2.js',
142 source: transformSource('dthetac', 2, g.nlm, 2, g.nlm),
143 },
144 {
145 name: 'dphic.mtoc2.js',
146 source: transformSource('dphic', 2, g.nlm, 2, g.nlm),
147 },
148 {
149 // Grid-space phi-derivative: two Fourier stages and an i*m multiply,
150 // no Legendre work (d/dphi is diagonal in the Fourier index). What
151 // lets the flux-form divergence skip the Q-flux's spherical-harmonic
152 // analysis.
153 name: 'dphig.mtoc2.js',
154 source: transformSource('dphig', g.npts, 1, g.npts, 1),
155 },
156 {
157 // The seeded random field a model's `init` starts from
158 // (src/mgpu/randnfun3.ts): a wavelength and the three surface
159 // coordinates in, one value per grid point out.
160 name: 'randnfun3.mtoc2.js',
161 source: randnfun3Source(g),
162 },
163 ];
164}
166/** Source for `randnfun3`'s `.mtoc2.js`: `f = randnfun3(lambda, x, y, z)`. */
167function randnfun3Source(g: GridSizes): string {
168 return `
169exports.name = "randnfun3";
171exports.transfer = function (argTypes, nargout) {
172 if (argTypes.length !== 4) {
173 throw new Error(
174 "randnfun3 takes a wavelength and the three surface coordinates -- " +
175 "randnfun3(lambda, gx, gy, gz) -- got " + argTypes.length + " argument(s)"
176 );
177 }
178 if (nargout > 1) {
179 throw new Error("randnfun3 returns one value, but " + nargout + " were requested");
180 }
181 var lam = argTypes[0];
182 if (!lam || lam.kind !== "Numeric" || lam.isComplex) {
183 throw new Error("randnfun3's wavelength must be a real number");
184 }
185 var ls = lam.shape;
186 if (!ls || ls.length !== 2 || ls[0] !== 1 || ls[1] !== 1) {
187 throw new Error(
188 "randnfun3's wavelength must be a single number, not a " +
189 (ls ? ls.join("x") : "unknown shape") + " array"
190 );
191 }
192 var names = ["gx", "gy", "gz"];
193 for (var i = 1; i < 4; i++) {
194 var a = argTypes[i];
195 if (!a || a.kind !== "Numeric" || a.isComplex) {
196 throw new Error("randnfun3 requires real numeric arrays (" + names[i - 1] + ")");
197 }
198 var s = a.shape;
199 if (!s || s.length !== 2 || s[0] !== ${g.npts} || s[1] !== 1) {
200 throw new Error(
201 "randnfun3 evaluates on the grid, so " + names[i - 1] +
202 " must be ${g.npts}x1, not " + (s ? s.join("x") : "unknown shape")
203 );
204 }
205 }
206 return [${numericType(g.npts, 1)}];
207};
209// Never called: this project executes the IR on WebGPU and emits no C.
210exports.emit = function () {
211 throw new Error("randnfun3: no C backend (this runs on WebGPU)");
212};
213exports.cBody = function () {
214 return "";
215};
216`;
217}
219/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
220export const EXTERNAL_OPS = new Set([
221 'synth', 'analys', 'dtheta', 'dphi', 'dthetac', 'dphic', 'dphig', 'randnfun3',
222]);