2 * The Laplacian stencils, as GPU dispatches.
3 *
4 * A grid field is an npts x 1 column vector in the buffer, laid out x-fastest:
5 * the point (ix, iy) is element `ix + nx*iy`. That is the only place in the
6 * project where the layout matters to anything but the renderer, which is the
7 * reason the stencils are host-provided operations rather than something a .m
8 * expresses with array slicing.
9 *
10 * Outside the grid the field is taken to be zero, which makes the outer
11 * boundary sound-soft (a pressure-release wall). Every scene puts an absorbing
12 * layer in front of it (tools/sponge.m), so in a well set-up run almost nothing
13 * reaches the wall to be reflected; what does is attenuated on the way out and
14 * again on the way back. See the README on how good an open boundary that is.
15 *
16 * Both stencils divide by h^2, so a .m reads `lap2(p)` as the Laplacian in the
17 * physical units the scene's coordinates are in. The grid spacing is compiled
18 * into the shader — the grid is fixed when a model is compiled, and changing it
19 * recompiles anyway.
20 */
21import { UnsupportedOnGpu, WORKGROUP_SIZE } from './wgsl.ts';
23export type StencilKind = 'lap2' | 'lap4';
25export interface StencilGrid {
26 nx: number;
27 ny: number;
28 /** Grid spacing, the same in x and y. */
29 h: number;
30}
32/**
33 * The 5-point second-order Laplacian, and the 9-point fourth-order one.
34 *
35 * The fourth-order stencil is the standard (-1/12, 4/3, -5/2, 4/3, -1/12)/h^2
36 * one-dimensional second derivative applied along each axis, so it is a
37 * five-point line in x plus a five-point line in y (nine points in all, not a
38 * 3x3 block). It costs twice the reads of `lap2` and buys two orders of
39 * accuracy, which for a wave problem shows up as much less grid dispersion:
40 * a pulse travelling many wavelengths stays a pulse instead of trailing
41 * numerical ripples. See models/leapfrog4.m.
42 */
43function stencilWGSL(kind: StencilKind, grid: StencilGrid, npts: number): string {
44 const { nx, ny, h } = grid;
45 const inv = 1 / (h * h);
46 // A read of a point outside the grid returns zero.
47 const at = `
48fn at(ix: i32, iy: i32) -> f32 {
49 if (ix < 0 || ix >= ${nx} || iy < 0 || iy >= ${ny}) { return 0.0; }
50 return src[u32(ix + ${nx} * iy)];
51}
52`;
53 const body =
54 kind === 'lap2'
55 ? ` let s = at(ix - 1, iy) + at(ix + 1, iy) + at(ix, iy - 1) + at(ix, iy + 1)
56 - 4.0 * at(ix, iy);`
57 : ` let s = (-1.0 / 12.0) * (at(ix - 2, iy) + at(ix + 2, iy) + at(ix, iy - 2) + at(ix, iy + 2))
58 + (4.0 / 3.0) * (at(ix - 1, iy) + at(ix + 1, iy) + at(ix, iy - 1) + at(ix, iy + 1))
59 - 5.0 * at(ix, iy);`;
60 return `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
61@group(0) @binding(1) var<storage, read> src: array<f32>;
62${at}
63@compute @workgroup_size(${WORKGROUP_SIZE})
64fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
65 let i = gid.x;
66 if (i >= ${npts}u) { return; }
67 let ix = i32(i % ${nx}u);
68 let iy = i32(i / ${nx}u);
69${body}
70 dst[i] = ${inv}f * s;
71}
72`;
73}
75/** Compiled stencil pipelines for one grid. Shared by every plan on it. */
76export class StencilPlan {
77 readonly grid: StencilGrid;
78 readonly npts: number;
80 #device: GPUDevice;
81 #layout: GPUBindGroupLayout;
82 #pipelines = new Map<StencilKind, GPUComputePipeline>();
84 private constructor(device: GPUDevice, grid: StencilGrid, layout: GPUBindGroupLayout) {
85 this.#device = device;
86 this.grid = grid;
87 this.npts = grid.nx * grid.ny;
88 this.#layout = layout;
89 }
91 static create(device: GPUDevice, grid: StencilGrid): StencilPlan {
92 const layout = device.createBindGroupLayout({
93 label: 'stencil',
94 entries: [
95 {
96 binding: 0,
97 visibility: GPUShaderStage.COMPUTE,
98 buffer: { type: 'storage' },
99 },
100 {
101 binding: 1,
102 visibility: GPUShaderStage.COMPUTE,
103 buffer: { type: 'read-only-storage' },
104 },
105 ],
106 });
107 return new StencilPlan(device, grid, layout);
108 }
110 /**
111 * The pipeline for one stencil, compiled on first use. Both are cheap, but
112 * a model uses one of them, and compiling only what is asked for keeps the
113 * op sequence honest about what the .m actually costs.
114 */
115 async pipeline(kind: StencilKind): Promise<GPUComputePipeline> {
116 const existing = this.#pipelines.get(kind);
117 if (existing) return existing;
118 const code = stencilWGSL(kind, this.grid, this.npts);
119 const module = this.#device.createShaderModule({ code, label: kind });
120 let pipeline: GPUComputePipeline;
121 try {
122 pipeline = await this.#device.createComputePipelineAsync({
123 layout: this.#device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
124 compute: { module, entryPoint: 'main' },
125 label: kind,
126 });
127 } catch (e) {
128 // No error scope here either; see makePipeline in plan.ts.
129 throw new UnsupportedOnGpu(
130 `stencil '${kind}': ${e instanceof Error ? e.message : String(e)}`,
131 );
132 }
133 this.#pipelines.set(kind, pipeline);
134 return pipeline;
135 }
137 createBinding(src: GPUBuffer, dst: GPUBuffer): GPUBindGroup {
138 return this.#device.createBindGroup({
139 layout: this.#layout,
140 entries: [
141 { binding: 0, resource: { buffer: dst } },
142 { binding: 1, resource: { buffer: src } },
143 ],
144 });
145 }
147 get workgroups(): number {
148 return Math.ceil(this.npts / WORKGROUP_SIZE);
149 }
150}