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/** Workspace files that make `synth` / `analys` resolvable during lowering. */
81export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
82 return [
83 {
84 name: 'synth.mtoc2.js',
85 source: transformSource('synth', 2, g.nlm, g.npts, 1),
86 },
87 {
88 name: 'analys.mtoc2.js',
89 source: transformSource('analys', g.npts, 1, 2, g.nlm),
90 },
91 ];
92}
94/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
95export const EXTERNAL_OPS = new Set(['synth', 'analys']);