1/**
2 * The Laplacian stencils, as external operations the .m can call: `lap2` (the
3 * 5-point second-order stencil) and `lap4` (the 9-point fourth-order one).
4 *
5 * Everything else a model does is element-wise, so these are the only places
6 * where a grid point reads its neighbours — the one thing the element-wise
7 * kernel emitter cannot express, since it walks a single linear index across
8 * every operand. Keeping them as named operations rather than as array slicing
9 * (`p(2:end-1, :)` and friends) means the .m never has to know how the grid is
10 * laid out in the buffer, and the host is free to implement the stencil as one
11 * dispatch (src/mgpu/stencil.ts).
12 *
13 * numbl needs only their *type rule* in order to lower a call site. It gets
14 * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
15 * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
16 * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
17 * imports available, so `transfer` builds numbl `Type` objects as plain
18 * literals, and the grid size is baked in by the generator below (a grid change
19 * recompiles anyway).
20 *
21 * The `emit`/`cBody` exports exist only because the loader's contract requires
22 * them; we never emit C. The actual implementation is supplied by the WGSL
23 * backend.
24 */
26export interface GridSizes {
27 /** Grid points, nx*ny. Grid fields are npts x 1 column vectors. */
28 npts: number;
29}
31const numericType = (rows: number, cols: number): string =>
32 `{ kind: "Numeric", elem: "double", isComplex: false, ` +
33 `dims: [${dim(rows)}, ${dim(cols)}], shape: [${rows}, ${cols}], sign: "unknown" }`;
35// numbl's tensorDouble() canonicalizes an extent of 1 to its shared DIM_ONE
36// singleton; mirror that so types compare equal to host-built ones.
37const dim = (n: number): string =>
38 n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;
40/** Source for one stencil's `.mtoc2.js`: grid field in, grid field out. */
41function stencilSource(name: string, npts: number): string {
42 return `
43exports.name = ${JSON.stringify(name)};
45exports.transfer = function (argTypes, nargout) {
46 if (argTypes.length !== 1) {
47 throw new Error("${name} takes exactly one argument, got " + argTypes.length);
48 }
49 if (nargout > 1) {
50 throw new Error("${name} returns one value, but " + nargout + " were requested");
51 }
52 var a = argTypes[0];
53 if (!a || a.kind !== "Numeric" || a.isComplex) {
54 throw new Error("${name} requires a real numeric array");
55 }
56 var s = a.shape;
57 if (!s || s.length !== 2 || s[0] !== ${npts} || s[1] !== 1) {
58 throw new Error(
59 "${name} works on grid fields, so its argument must be ${npts}x1, not " +
60 (s ? s.join("x") : "unknown shape")
61 );
62 }
63 return [${numericType(npts, 1)}];
64};
66// Never called: this project executes the IR on WebGPU and emits no C.
67exports.emit = function () {
68 throw new Error("${name}: no C backend (this stencil runs on WebGPU)");
69};
70exports.cBody = function () {
71 return "";
72};
73`;
74}
76/** Workspace files that make `lap2` / `lap4` resolvable during lowering. */
77export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
78 return [
79 { name: 'lap2.mtoc2.js', source: stencilSource('lap2', g.npts) },
80 { name: 'lap4.mtoc2.js', source: stencilSource('lap4', g.npts) },
81 ];
82}
84/** Names the WGSL backend must implement as stencil dispatches rather than
85 * element-wise kernels. */
86export const EXTERNAL_OPS = new Set(['lap2', 'lap4']);