1/**
2 * The host-provided operations a model .m can call — the only places where a
3 * point reads anything but itself.
4 *
5 * Everything else a model does is element-wise, so these five are where the
6 * physics that couples neighbours (and couples the two grids) lives:
7 *
8 * dxx(u) second difference along the string (string -> string)
9 * dxxxx(u) fourth difference along the string (string -> string)
10 * lapw(p, w) wall-masked 7-point Laplacian in the air (air, air -> air)
11 * spread(a) sample a string field at each air point's x (string -> air)
12 * bridge(u) du/dx at the string's bridge end, broadcast (string -> air)
13 *
14 * numbl needs only their *type rule* in order to lower a call site. It gets
15 * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
16 * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
17 * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
18 * imports available, so `transfer` builds numbl `Type` objects as plain
19 * literals, and the grid sizes are baked in by the generator below (a grid
20 * change recompiles anyway).
21 *
22 * The `emit`/`cBody` exports exist only because the loader's contract requires
23 * them; we never emit C. The actual implementation is supplied by the WGSL
24 * backend (src/mgpu/ops.ts).
25 */
27export interface GridSizes {
28 /** Air grid points, nx*ny*nz. Air fields are npts x 1 column vectors. */
29 npts: number;
30 /** String nodes. String fields are ns x 1 column vectors. */
31 ns: number;
32}
34/** One external operation's shape contract. */
35export interface ExternalOpSpec {
36 name: string;
37 /** Element counts of the arguments, in order. */
38 args: ('air' | 'string')[];
39 out: 'air' | 'string';
40}
42export const EXTERNAL_OP_SPECS: ExternalOpSpec[] = [
43 { name: 'dxx', args: ['string'], out: 'string' },
44 { name: 'dxxxx', args: ['string'], out: 'string' },
45 { name: 'lapw', args: ['air', 'air'], out: 'air' },
46 { name: 'spread', args: ['string'], out: 'air' },
47 { name: 'bridge', args: ['string'], out: 'air' },
48];
50/** Names the WGSL backend must implement as dispatches rather than
51 * element-wise kernels, with the argument count of each. */
52export const EXTERNAL_OPS = new Map(EXTERNAL_OP_SPECS.map((s) => [s.name, s]));
54const numericType = (rows: number): string =>
55 `{ kind: "Numeric", elem: "double", isComplex: false, ` +
56 `dims: [{ kind: "exact", value: ${rows} }, { kind: "exact", value: 1 }], ` +
57 `shape: [${rows}, 1], sign: "unknown" }`;
59/** Source for one op's `.mtoc2.js`: fields in, field out, shapes checked. */
60function opSource(spec: ExternalOpSpec, g: GridSizes): string {
61 const size = (k: 'air' | 'string'): number => (k === 'air' ? g.npts : g.ns);
62 const checks = spec.args
63 .map((kind, i) => {
64 const n = size(kind);
65 const what = kind === 'air' ? 'an air field' : 'a string field';
66 return `
67 var a${i} = argTypes[${i}];
68 if (!a${i} || a${i}.kind !== "Numeric" || a${i}.isComplex) {
69 throw new Error("${spec.name}: argument ${i + 1} must be a real numeric array");
70 }
71 var s${i} = a${i}.shape;
72 if (!s${i} || s${i}.length !== 2 || s${i}[0] !== ${n} || s${i}[1] !== 1) {
73 throw new Error(
74 "${spec.name}: argument ${i + 1} must be ${what} (${n}x1), not " +
75 (s${i} ? s${i}.join("x") : "unknown shape")
76 );
77 }`;
78 })
79 .join('\n');
80 return `
81exports.name = ${JSON.stringify(spec.name)};
83exports.transfer = function (argTypes, nargout) {
84 if (argTypes.length !== ${spec.args.length}) {
85 throw new Error("${spec.name} takes ${spec.args.length} argument(s), got " + argTypes.length);
86 }
87 if (nargout > 1) {
88 throw new Error("${spec.name} returns one value, but " + nargout + " were requested");
89 }
90${checks}
91 return [${numericType(size(spec.out))}];
92};
94// Never called: this project executes the IR on WebGPU and emits no C.
95exports.emit = function () {
96 throw new Error("${spec.name}: no C backend (this op runs on WebGPU)");
97};
98exports.cBody = function () {
99 return "";
100};
101`;
102}
104/** Workspace files that make the ops resolvable during lowering. */
105export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
106 return EXTERNAL_OP_SPECS.map((spec) => ({
107 name: `${spec.name}.mtoc2.js`,
108 source: opSource(spec, g),
109 }));
110}