1/**
2 * The external operations, as GPU dispatches.
3 *
4 * An air field is an npts x 1 column vector laid out x-fastest — the point
5 * (ix, iy, iz) is element `ix + nx*(iy + ny*iz)` — and a string field is ns
6 * nodes from x = 0 to x = Ls. Those layouts, and the placement of the string
7 * inside the air's coordinates, matter only here and to the renderer; a .m
8 * never sees them, which is the reason these are host-provided operations
9 * rather than something a model expresses with array slicing.
10 *
11 * `lapw` is the wall-masked Laplacian, and it is what makes the body's shell
12 * rigid. Each of the six face terms is scaled by the mask at the neighbour it
13 * reads, so a face into the wall contributes nothing — which is exactly the
14 * discrete Neumann (zero normal velocity, sound-hard) condition, in the
15 * divergence form div(w grad p). The alternative the flat siblings use, a
16 * wall as a fast material, is unusable here: wood at ~4000 m/s would cut the
17 * global timestep twelve-fold, where a mask costs nothing. Outside the domain
18 * the field is taken to be zero (a soft outer boundary); the scene's
19 * absorbing layer is meant to have swallowed the wave before it matters.
20 *
21 * `spread` and `bridge` are the two couplings from the string into the air.
22 * `spread` gives each air point the string value at its own x (linearly
23 * interpolated, zero beyond the ends): multiplied by a scene-built line
24 * profile, that is the string radiating directly. `bridge` broadcasts the
25 * string's slope at its bridge end to every air point: multiplied by a
26 * scene-built patch profile, that is the bridge force driving the top plate.
27 * Both read tiny buffers and write big ones, so each is one cheap dispatch.
28 */
29import { UnsupportedOnGpu, WORKGROUP_SIZE } from './wgsl.ts';
30import { EXTERNAL_OPS } from './externals.ts';
32export type OpKind = 'dxx' | 'dxxxx' | 'lapw' | 'spread' | 'bridge';
34export interface OpGeometry {
35 /** Air grid. */
36 nx: number;
37 ny: number;
38 nz: number;
39 h: number;
40 /** String grid. */
41 ns: number;
42 hs: number;
43 /** World x of the string's first node (the nut); the bridge is the last. */
44 xs0: number;
45 /** World x extent of the air domain, for mapping voxel index to metres. */
46 Lx: number;
47}
49function opWGSL(kind: OpKind, g: OpGeometry): { code: string; outCount: number; args: number } {
50 const { nx, ny, nz, h, ns, hs } = g;
51 const npts = nx * ny * nz;
52 const inv2 = 1 / (hs * hs);
53 const inv4 = 1 / (hs * hs * hs * hs);
55 const stringAt = `
56fn at(j: i32) -> f32 {
57 if (j < 0 || j >= ${ns}) { return 0.0; }
58 return src[u32(j)];
59}
60`;
62 switch (kind) {
63 case 'dxx':
64 return {
65 args: 1,
66 outCount: ns,
67 code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
68@group(0) @binding(1) var<storage, read> src: array<f32>;
69${stringAt}
70@compute @workgroup_size(${WORKGROUP_SIZE})
71fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
72 let i = gid.x;
73 if (i >= ${ns}u) { return; }
74 let j = i32(i);
75 dst[i] = ${inv2}f * (at(j - 1) - 2.0 * at(j) + at(j + 1));
76}
77`,
78 };
80 case 'dxxxx':
81 return {
82 args: 1,
83 outCount: ns,
84 code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
85@group(0) @binding(1) var<storage, read> src: array<f32>;
86${stringAt}
87@compute @workgroup_size(${WORKGROUP_SIZE})
88fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
89 let i = gid.x;
90 if (i >= ${ns}u) { return; }
91 let j = i32(i);
92 dst[i] = ${inv4}f * (at(j - 2) - 4.0 * at(j - 1) + 6.0 * at(j) - 4.0 * at(j + 1) + at(j + 2));
93}
94`,
95 };
97 case 'lapw':
98 return {
99 args: 2,
100 outCount: npts,
101 code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
102@group(0) @binding(1) var<storage, read> src: array<f32>;
103@group(0) @binding(2) var<storage, read> wall: array<f32>;
105fn idx(ix: i32, iy: i32, iz: i32) -> i32 {
106 return ix + ${nx} * (iy + ${ny} * iz);
107}
108fn inside(ix: i32, iy: i32, iz: i32) -> bool {
109 return ix >= 0 && ix < ${nx} && iy >= 0 && iy < ${ny} && iz >= 0 && iz < ${nz};
110}
111// One face's contribution: masked by the wall at the neighbour, so a face
112// into the wall carries no flux — the Neumann (rigid) condition.
113fn face(ix: i32, iy: i32, iz: i32, pc: f32) -> f32 {
114 if (!inside(ix, iy, iz)) { return -pc; } // outside the domain: p = 0, open
115 let k = u32(idx(ix, iy, iz));
116 return wall[k] * (src[k] - pc);
117}
119@compute @workgroup_size(${WORKGROUP_SIZE})
120fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
121 let i = gid.x;
122 if (i >= ${npts}u) { return; }
123 let ix = i32(i % ${nx}u);
124 let iy = i32((i / ${nx}u) % ${ny}u);
125 let iz = i32(i / ${nx * ny}u);
126 let pc = src[i];
127 let s = face(ix - 1, iy, iz, pc) + face(ix + 1, iy, iz, pc)
128 + face(ix, iy - 1, iz, pc) + face(ix, iy + 1, iz, pc)
129 + face(ix, iy, iz - 1, pc) + face(ix, iy, iz + 1, pc);
130 dst[i] = ${1 / (h * h)}f * wall[i] * s;
131}
132`,
133 };
135 case 'spread':
136 return {
137 args: 1,
138 outCount: npts,
139 code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
140@group(0) @binding(1) var<storage, read> src: array<f32>;
142@compute @workgroup_size(${WORKGROUP_SIZE})
143fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
144 let i = gid.x;
145 if (i >= ${npts}u) { return; }
146 let ix = i % ${nx}u;
147 let xw = ${-g.Lx / 2}f + (f32(ix) + 0.5) * ${h}f;
148 let u = (xw - ${g.xs0}f) / ${hs}f;
149 let j = i32(floor(u));
150 if (j < 0 || j >= ${ns - 1}) { dst[i] = 0.0; return; }
151 let f = u - f32(j);
152 dst[i] = mix(src[u32(j)], src[u32(j) + 1u], f);
153}
154`,
155 };
157 case 'bridge':
158 return {
159 args: 1,
160 outCount: npts,
161 code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
162@group(0) @binding(1) var<storage, read> src: array<f32>;
164@compute @workgroup_size(${WORKGROUP_SIZE})
165fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
166 let i = gid.x;
167 if (i >= ${npts}u) { return; }
168 // du/dx at the last node. The end is pinned, so this is the string's
169 // arriving slope — what the tension pulls the bridge with.
170 dst[i] = (src[${ns - 1}u] - src[${ns - 2}u]) * ${1 / hs}f;
171}
172`,
173 };
174 }
175}
177/** Compiled op pipelines for one pair of grids. Shared by every plan. */
178export class OpPlan {
179 readonly geometry: OpGeometry;
181 #device: GPUDevice;
182 #layouts = new Map<number, GPUBindGroupLayout>();
183 #pipelines = new Map<OpKind, { pipeline: GPUComputePipeline; outCount: number; args: number }>();
185 constructor(device: GPUDevice, geometry: OpGeometry) {
186 this.#device = device;
187 this.geometry = geometry;
188 }
190 /** The op's shape contract, for the planner's checks. */
191 spec(kind: OpKind): { argCounts: number[]; outCount: number } {
192 const s = EXTERNAL_OPS.get(kind);
193 if (!s) throw new UnsupportedOnGpu(`unknown external op '${kind}'`);
194 const { nx, ny, nz, ns } = this.geometry;
195 const size = (k: 'air' | 'string'): number => (k === 'air' ? nx * ny * nz : ns);
196 return { argCounts: s.args.map(size), outCount: size(s.out) };
197 }
199 #layout(args: number): GPUBindGroupLayout {
200 const existing = this.#layouts.get(args);
201 if (existing) return existing;
202 const entries: GPUBindGroupLayoutEntry[] = [
203 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
204 ];
205 for (let i = 0; i < args; i++) {
206 entries.push({
207 binding: i + 1,
208 visibility: GPUShaderStage.COMPUTE,
209 buffer: { type: 'read-only-storage' },
210 });
211 }
212 const layout = this.#device.createBindGroupLayout({ label: `op-${args}`, entries });
213 this.#layouts.set(args, layout);
214 return layout;
215 }
217 /** The pipeline for one op, compiled on first use. */
218 async pipeline(kind: OpKind): Promise<{ pipeline: GPUComputePipeline; outCount: number; args: number }> {
219 const existing = this.#pipelines.get(kind);
220 if (existing) return existing;
221 const { code, outCount, args } = opWGSL(kind, this.geometry);
222 const module = this.#device.createShaderModule({ code, label: kind });
223 let pipeline: GPUComputePipeline;
224 try {
225 pipeline = await this.#device.createComputePipelineAsync({
226 layout: this.#device.createPipelineLayout({ bindGroupLayouts: [this.#layout(args)] }),
227 compute: { module, entryPoint: 'main' },
228 label: kind,
229 });
230 } catch (e) {
231 // No error scope here either; see makePipeline in plan.ts.
232 throw new UnsupportedOnGpu(
233 `op '${kind}': ${e instanceof Error ? e.message : String(e)}`,
234 );
235 }
236 const built = { pipeline, outCount, args };
237 this.#pipelines.set(kind, built);
238 return built;
239 }
241 createBinding(kind: OpKind, srcs: GPUBuffer[], dst: GPUBuffer): GPUBindGroup {
242 const entries: GPUBindGroupEntry[] = [{ binding: 0, resource: { buffer: dst } }];
243 srcs.forEach((s, i) => entries.push({ binding: i + 1, resource: { buffer: s } }));
244 return this.#device.createBindGroup({ layout: this.#layout(srcs.length), entries });
245 }
247 workgroups(outCount: number): number {
248 return Math.ceil(outCount / WORKGROUP_SIZE);
249 }
250}