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`. */
39function transformSource(
40 name: string,
41 inRows: number,
42 inCols: number,
43 outRows: number,
44 outCols: number,
45): string {
46 return `
47exports.name = ${JSON.stringify(name)};
49exports.transfer = function (argTypes, nargout) {
50 if (argTypes.length !== 1) {
51 throw new Error("${name} takes exactly one argument, got " + argTypes.length);
52 }
53 if (nargout > 1) {
54 throw new Error("${name} returns one value, but " + nargout + " were requested");
55 }
56 var a = argTypes[0];
57 if (!a || a.kind !== "Numeric" || a.isComplex) {
58 throw new Error("${name} requires a real numeric array");
59 }
60 var s = a.shape;
61 if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
62 throw new Error(
63 "${name} requires a ${inRows}x${inCols} array, got " +
64 (s ? s.join("x") : "unknown shape")
65 );
66 }
67 return [${numericType(outRows, outCols)}];
68};
70// Never called: this project executes the IR on WebGPU and emits no C.
71exports.emit = function () {
72 throw new Error("${name}: no C backend (this transform runs on WebGPU)");
73};
74exports.cBody = function () {
75 return "";
76};
77`;
78}
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. */
84function dotSource(): string {
85 return `
86exports.name = "dot";
88exports.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};
113// Never called: this project executes the IR on WebGPU and emits no C.
114exports.emit = function () {
115 throw new Error("dot: no C backend (this reduction runs on WebGPU)");
116};
117exports.cBody = function () {
118 return "";
119};
120`;
121}
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 */
141function indexOpFiles(nlm: number): { name: string; source: string }[] {
142 const shared = `
143function isRealNumeric(t) {
144 return t && t.kind === "Numeric" && !t.isComplex;
145}
146function 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. */
151function baseType(t) {
152 return {
153 kind: "Numeric", elem: t.elem, isComplex: false,
154 dims: t.dims, shape: t.shape, sign: "unknown",
155 };
156}
157function 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}
173exports.name = "getslab";
174exports.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};
181exports.emit = function () { throw new Error("getslab: no C backend"); };
182exports.cBody = function () { return ""; };
183`,
184 },
185 {
186 name: 'setslab.mtoc2.js',
187 source: `${shared}
188exports.name = "setslab";
189exports.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};
200exports.emit = function () { throw new Error("setslab: no C backend"); };
201exports.cBody = function () { return ""; };
202`,
203 },
204 {
205 name: 'getat.mtoc2.js',
206 source: `${shared}
207exports.name = "getat";
208exports.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};
219exports.emit = function () { throw new Error("getat: no C backend"); };
220exports.cBody = function () { return ""; };
221`,
222 },
223 {
224 name: 'setat.mtoc2.js',
225 source: `${shared}
226exports.name = "setat";
227exports.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};
239exports.emit = function () { throw new Error("setat: no C backend"); };
240exports.cBody = function () { return ""; };
241`,
242 },
243 ];
244}
246/**
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.
252 */
253export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
254 return [
255 {
256 name: 'synth.mtoc2.js',
257 source: transformSource('synth', 2, g.nlm, g.npts, 1),
258 },
259 {
260 name: 'analys.mtoc2.js',
261 source: transformSource('analys', g.npts, 1, 2, g.nlm),
262 },
263 {
264 name: 'dtheta.mtoc2.js',
265 source: transformSource('dtheta', 2, g.nlm, g.npts, 1),
266 },
267 {
268 name: 'dphi.mtoc2.js',
269 source: transformSource('dphi', 2, g.nlm, g.npts, 1),
270 },
271 {
272 name: 'dot.mtoc2.js',
273 source: dotSource(),
274 },
275 ...indexOpFiles(g.nlm),
276 ];
277}
279/** Names the WGSL backend must implement as GPU encodes rather than kernels,
280 * each with the argument counts it accepts. */
281export 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]);